From 51c40abd408ab0637f3b65cf5697a4ee85a544a4 Mon Sep 17 00:00:00 2001 From: Noah Dietz Date: Mon, 5 Apr 2021 14:00:06 -0700 Subject: [PATCH] feat: wrap non-retryable RPCs in retry machinery (#1328) * feat: wrap non-retryable RPCs in retry machinery * address feedback * add contextual retry settings to timeout tests * add tracing test for non-retryable callable * add server streaming callable tests --- .../com/google/api/gax/grpc/TimeoutTest.java | 130 ++++++++++++--- .../com/google/api/gax/rpc/Callables.java | 33 ++-- .../com/google/api/gax/rpc/CallableTest.java | 148 ++++++++++++++++++ .../api/gax/tracing/TracedCallableTest.java | 119 ++++++++++++++ 4 files changed, 395 insertions(+), 35 deletions(-) create mode 100644 gax/src/test/java/com/google/api/gax/rpc/CallableTest.java create mode 100644 gax/src/test/java/com/google/api/gax/tracing/TracedCallableTest.java diff --git a/gax-grpc/src/test/java/com/google/api/gax/grpc/TimeoutTest.java b/gax-grpc/src/test/java/com/google/api/gax/grpc/TimeoutTest.java index bb35d21ef..3ff64b7b2 100644 --- a/gax-grpc/src/test/java/com/google/api/gax/grpc/TimeoutTest.java +++ b/gax-grpc/src/test/java/com/google/api/gax/grpc/TimeoutTest.java @@ -74,9 +74,12 @@ public class TimeoutTest { private static final int DEADLINE_IN_MINUTES = 10; private static final int DEADLINE_IN_SECONDS = 20; private static final ImmutableSet emptyRetryCodes = ImmutableSet.of(); + private static final ImmutableSet retryUnknownCode = + ImmutableSet.of(StatusCode.Code.UNKNOWN); private static final Duration totalTimeout = Duration.ofDays(DEADLINE_IN_DAYS); private static final Duration maxRpcTimeout = Duration.ofMinutes(DEADLINE_IN_MINUTES); private static final Duration initialRpcTimeout = Duration.ofSeconds(DEADLINE_IN_SECONDS); + private static final GrpcCallContext defaultCallContext = GrpcCallContext.createDefault(); @Rule public MockitoRule mockitoRule = MockitoJUnit.rule().strictness(Strictness.STRICT_STUBS); @Mock private Marshaller stringMarshaller; @@ -97,7 +100,8 @@ public void testNonRetryUnarySettings() { .setRpcTimeoutMultiplier(1.0) .setMaxRpcTimeout(maxRpcTimeout) .build(); - CallOptions callOptionsUsed = setupUnaryCallable(retrySettings); + CallOptions callOptionsUsed = + setupUnaryCallable(retrySettings, emptyRetryCodes, defaultCallContext); // Verify that the gRPC channel used the CallOptions with our custom timeout of ~2 Days. assertThat(callOptionsUsed.getDeadline()).isNotNull(); @@ -108,6 +112,46 @@ public void testNonRetryUnarySettings() { assertThat(callOptionsUsed.getAuthority()).isEqualTo(CALL_OPTIONS_AUTHORITY); } + @Test + public void testNonRetryUnarySettingsContextWithRetry() { + RetrySettings retrySettings = + RetrySettings.newBuilder() + .setTotalTimeout(totalTimeout) + .setInitialRetryDelay(Duration.ZERO) + .setRetryDelayMultiplier(1.0) + .setMaxRetryDelay(Duration.ZERO) + .setMaxAttempts(1) + .setJittered(true) + .setInitialRpcTimeout(initialRpcTimeout) + .setRpcTimeoutMultiplier(1.0) + .setMaxRpcTimeout(maxRpcTimeout) + .build(); + Duration newTimeout = Duration.ofSeconds(5); + RetrySettings contextRetrySettings = + retrySettings + .toBuilder() + .setInitialRpcTimeout(newTimeout) + .setMaxRpcTimeout(newTimeout) + .setMaxAttempts(3) + .build(); + GrpcCallContext retryingContext = + defaultCallContext + .withRetrySettings(contextRetrySettings) + .withRetryableCodes(retryUnknownCode); + CallOptions callOptionsUsed = + setupUnaryCallable(retrySettings, emptyRetryCodes, retryingContext); + + // Verify that the gRPC channel used the CallOptions the initial timeout of ~5 seconds. + // This indicates that the context retry settings were used on a callable that was instantiated + // with non-retryable settings. + assertThat(callOptionsUsed.getDeadline()).isNotNull(); + assertThat(callOptionsUsed.getDeadline()) + .isGreaterThan(Deadline.after(newTimeout.toSecondsPart() - 1, TimeUnit.SECONDS)); + assertThat(callOptionsUsed.getDeadline()) + .isLessThan(Deadline.after(newTimeout.toSecondsPart(), TimeUnit.SECONDS)); + assertThat(callOptionsUsed.getAuthority()).isEqualTo(CALL_OPTIONS_AUTHORITY); + } + @Test public void testNonRetryUnarySettingsWithoutInitialRpcTimeout() { RetrySettings retrySettings = @@ -121,7 +165,8 @@ public void testNonRetryUnarySettingsWithoutInitialRpcTimeout() { .setRpcTimeoutMultiplier(1.0) .setMaxRpcTimeout(maxRpcTimeout) .build(); - CallOptions callOptionsUsed = setupUnaryCallable(retrySettings); + CallOptions callOptionsUsed = + setupUnaryCallable(retrySettings, emptyRetryCodes, defaultCallContext); // Verify that the gRPC channel used the CallOptions with our custom timeout of ~2 Days. assertThat(callOptionsUsed.getDeadline()).isNotNull(); @@ -145,7 +190,8 @@ public void testNonRetryUnarySettingsWithoutIndividualRpcTimeout() { .setRpcTimeoutMultiplier(1.0) .setRpcTimeoutMultiplier(1.0) .build(); - CallOptions callOptionsUsed = setupUnaryCallable(retrySettings); + CallOptions callOptionsUsed = + setupUnaryCallable(retrySettings, emptyRetryCodes, defaultCallContext); // Verify that the gRPC channel used the CallOptions with our custom timeout of ~2 Days. assertThat(callOptionsUsed.getDeadline()).isNotNull(); @@ -170,7 +216,8 @@ public void testNonRetryServerStreamingSettings() { .setRpcTimeoutMultiplier(1.0) .setMaxRpcTimeout(maxRpcTimeout) .build(); - CallOptions callOptionsUsed = setupServerStreamingCallable(retrySettings); + CallOptions callOptionsUsed = + setupServerStreamingCallable(retrySettings, emptyRetryCodes, defaultCallContext); // Verify that the gRPC channel used the CallOptions with our custom timeout of ~2 Days. assertThat(callOptionsUsed.getDeadline()).isNotNull(); @@ -181,6 +228,41 @@ public void testNonRetryServerStreamingSettings() { assertThat(callOptionsUsed.getAuthority()).isEqualTo(CALL_OPTIONS_AUTHORITY); } + @Test + public void testNonRetryServerStreamingSettingsContextWithRetry() { + RetrySettings retrySettings = + RetrySettings.newBuilder() + .setTotalTimeout(totalTimeout) + .setInitialRetryDelay(Duration.ZERO) + .setRetryDelayMultiplier(1.0) + .setMaxRetryDelay(Duration.ZERO) + .setMaxAttempts(1) + .setJittered(true) + .setInitialRpcTimeout(initialRpcTimeout) + .setRpcTimeoutMultiplier(1.0) + .setMaxRpcTimeout(maxRpcTimeout) + .build(); + Duration newTimeout = Duration.ofSeconds(5); + RetrySettings contextRetrySettings = + retrySettings.toBuilder().setTotalTimeout(newTimeout).setMaxAttempts(3).build(); + GrpcCallContext retryingContext = + defaultCallContext + .withRetrySettings(contextRetrySettings) + .withRetryableCodes(retryUnknownCode); + CallOptions callOptionsUsed = + setupServerStreamingCallable(retrySettings, emptyRetryCodes, retryingContext); + + // Verify that the gRPC channel used the CallOptions the total timeout of ~5 seconds. + // This indicates that the context retry settings were used on a callable that was instantiated + // with non-retryable settings. + assertThat(callOptionsUsed.getDeadline()).isNotNull(); + assertThat(callOptionsUsed.getDeadline()) + .isGreaterThan(Deadline.after(newTimeout.toSecondsPart() - 1, TimeUnit.SECONDS)); + assertThat(callOptionsUsed.getDeadline()) + .isLessThan(Deadline.after(newTimeout.toSecondsPart(), TimeUnit.SECONDS)); + assertThat(callOptionsUsed.getAuthority()).isEqualTo(CALL_OPTIONS_AUTHORITY); + } + @Test public void testNonRetryServerStreamingSettingsWithoutInitialRpcTimeout() { RetrySettings retrySettings = @@ -194,7 +276,8 @@ public void testNonRetryServerStreamingSettingsWithoutInitialRpcTimeout() { .setRpcTimeoutMultiplier(1.0) .setMaxRpcTimeout(maxRpcTimeout) .build(); - CallOptions callOptionsUsed = setupServerStreamingCallable(retrySettings); + CallOptions callOptionsUsed = + setupServerStreamingCallable(retrySettings, emptyRetryCodes, defaultCallContext); // Verify that the gRPC channel used the CallOptions with our custom timeout of ~2 Days. assertThat(callOptionsUsed.getDeadline()).isNotNull(); @@ -218,7 +301,8 @@ public void testNonRetryServerStreamingSettingsWithoutIndividualRpcTimeout() { .setRpcTimeoutMultiplier(1.0) .setRpcTimeoutMultiplier(1.0) .build(); - CallOptions callOptionsUsed = setupServerStreamingCallable(retrySettings); + CallOptions callOptionsUsed = + setupServerStreamingCallable(retrySettings, emptyRetryCodes, defaultCallContext); // Verify that the gRPC channel used the CallOptions with our custom timeout of ~2 Days. assertThat(callOptionsUsed.getDeadline()).isNotNull(); @@ -229,7 +313,10 @@ public void testNonRetryServerStreamingSettingsWithoutIndividualRpcTimeout() { assertThat(callOptionsUsed.getAuthority()).isEqualTo(CALL_OPTIONS_AUTHORITY); } - private CallOptions setupUnaryCallable(RetrySettings retrySettings) { + private CallOptions setupUnaryCallable( + RetrySettings retrySettings, + ImmutableSet retryableCodes, + GrpcCallContext callContext) { MethodDescriptor methodDescriptor = MethodDescriptor.newBuilder() .setSchemaDescriptor("yaml") @@ -248,8 +335,8 @@ private CallOptions setupUnaryCallable(RetrySettings retrySettings) { // Clobber the "authority" property with an identifier that allows us to trace // the use of this CallOptions variable. CallOptions spyCallOptions = CallOptions.DEFAULT.withAuthority("RETRYING_TEST"); - GrpcCallContext grpcCallContext = - GrpcCallContext.createDefault().withChannel(managedChannel).withCallOptions(spyCallOptions); + GrpcCallContext context = + callContext.withChannel(managedChannel).withCallOptions(spyCallOptions); ArgumentCaptor callOptionsArgumentCaptor = ArgumentCaptor.forClass(CallOptions.class); @@ -266,16 +353,16 @@ private CallOptions setupUnaryCallable(RetrySettings retrySettings) { .setMethodDescriptor(methodDescriptor) .setParamsExtractor(paramsExtractor) .build(); - UnaryCallSettings nonRetriedCallSettings = + UnaryCallSettings unaryCallSettings = UnaryCallSettings.newUnaryCallSettingsBuilder() .setRetrySettings(retrySettings) - .setRetryableCodes(emptyRetryCodes) + .setRetryableCodes(retryableCodes) .build(); UnaryCallable callable = GrpcCallableFactory.createUnaryCallable( grpcCallSettings, - nonRetriedCallSettings, - ClientContext.newBuilder().setDefaultCallContext(grpcCallContext).build()); + unaryCallSettings, + ClientContext.newBuilder().setDefaultCallContext(context).build()); try { ApiFuture future = callable.futureCall("Is your refrigerator running?"); @@ -287,7 +374,10 @@ private CallOptions setupUnaryCallable(RetrySettings retrySettings) { return callOptionsArgumentCaptor.getValue(); } - private CallOptions setupServerStreamingCallable(RetrySettings retrySettings) { + private CallOptions setupServerStreamingCallable( + RetrySettings retrySettings, + ImmutableSet retryableCodes, + GrpcCallContext callContext) { MethodDescriptor methodDescriptor = MethodDescriptor.newBuilder() .setSchemaDescriptor("yaml") @@ -306,8 +396,8 @@ private CallOptions setupServerStreamingCallable(RetrySettings retrySettings) { // Clobber the "authority" property with an identifier that allows us to trace // the use of this CallOptions variable. CallOptions spyCallOptions = CallOptions.DEFAULT.withAuthority("RETRYING_TEST"); - GrpcCallContext grpcCallContext = - GrpcCallContext.createDefault().withChannel(managedChannel).withCallOptions(spyCallOptions); + GrpcCallContext context = + callContext.withChannel(managedChannel).withCallOptions(spyCallOptions); ArgumentCaptor callOptionsArgumentCaptor = ArgumentCaptor.forClass(CallOptions.class); @@ -324,16 +414,16 @@ private CallOptions setupServerStreamingCallable(RetrySettings retrySettings) { .setMethodDescriptor(methodDescriptor) .setParamsExtractor(paramsExtractor) .build(); - ServerStreamingCallSettings nonRetriedCallSettings = + ServerStreamingCallSettings serverStreamingCallSettings = ServerStreamingCallSettings.newBuilder() .setRetrySettings(retrySettings) - .setRetryableCodes(emptyRetryCodes) + .setRetryableCodes(retryableCodes) .build(); ServerStreamingCallable callable = GrpcCallableFactory.createServerStreamingCallable( grpcCallSettings, - nonRetriedCallSettings, - ClientContext.newBuilder().setDefaultCallContext(grpcCallContext).build()); + serverStreamingCallSettings, + ClientContext.newBuilder().setDefaultCallContext(context).build()); try { ServerStream stream = callable.call("Is your refrigerator running?"); diff --git a/gax/src/main/java/com/google/api/gax/rpc/Callables.java b/gax/src/main/java/com/google/api/gax/rpc/Callables.java index a1c34a02c..939e66b7f 100644 --- a/gax/src/main/java/com/google/api/gax/rpc/Callables.java +++ b/gax/src/main/java/com/google/api/gax/rpc/Callables.java @@ -56,19 +56,21 @@ public static UnaryCallable retrying( UnaryCallSettings callSettings, ClientContext clientContext) { - if (areRetriesDisabled(callSettings.getRetryableCodes(), callSettings.getRetrySettings())) { + UnaryCallSettings settings = callSettings; + + if (areRetriesDisabled(settings.getRetryableCodes(), settings.getRetrySettings())) { // When retries are disabled, the total timeout can be treated as the rpc timeout. - return innerCallable.withDefaultCallContext( - clientContext - .getDefaultCallContext() - .withTimeout(callSettings.getRetrySettings().getTotalTimeout())); + settings = + settings + .toBuilder() + .setSimpleTimeoutNoRetries(settings.getRetrySettings().getTotalTimeout()) + .build(); } RetryAlgorithm retryAlgorithm = new RetryAlgorithm<>( new ApiResultRetryAlgorithm(), - new ExponentialRetryAlgorithm( - callSettings.getRetrySettings(), clientContext.getClock())); + new ExponentialRetryAlgorithm(settings.getRetrySettings(), clientContext.getClock())); ScheduledRetryingExecutor retryingExecutor = new ScheduledRetryingExecutor<>(retryAlgorithm, clientContext.getExecutor()); return new RetryingCallable<>( @@ -81,25 +83,26 @@ public static ServerStreamingCallable ServerStreamingCallSettings callSettings, ClientContext clientContext) { - if (areRetriesDisabled(callSettings.getRetryableCodes(), callSettings.getRetrySettings())) { + ServerStreamingCallSettings settings = callSettings; + if (areRetriesDisabled(settings.getRetryableCodes(), settings.getRetrySettings())) { // When retries are disabled, the total timeout can be treated as the rpc timeout. - return innerCallable.withDefaultCallContext( - clientContext - .getDefaultCallContext() - .withTimeout(callSettings.getRetrySettings().getTotalTimeout())); + settings = + settings + .toBuilder() + .setSimpleTimeoutNoRetries(settings.getRetrySettings().getTotalTimeout()) + .build(); } StreamingRetryAlgorithm retryAlgorithm = new StreamingRetryAlgorithm<>( new ApiResultRetryAlgorithm(), - new ExponentialRetryAlgorithm( - callSettings.getRetrySettings(), clientContext.getClock())); + new ExponentialRetryAlgorithm(settings.getRetrySettings(), clientContext.getClock())); ScheduledRetryingExecutor retryingExecutor = new ScheduledRetryingExecutor<>(retryAlgorithm, clientContext.getExecutor()); return new RetryingServerStreamingCallable<>( - innerCallable, retryingExecutor, callSettings.getResumptionStrategy()); + innerCallable, retryingExecutor, settings.getResumptionStrategy()); } @BetaApi("The surface for streaming is not stable yet and may change in the future.") diff --git a/gax/src/test/java/com/google/api/gax/rpc/CallableTest.java b/gax/src/test/java/com/google/api/gax/rpc/CallableTest.java new file mode 100644 index 000000000..eaa0be04c --- /dev/null +++ b/gax/src/test/java/com/google/api/gax/rpc/CallableTest.java @@ -0,0 +1,148 @@ +/* + * Copyright 2021 Google LLC + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are + * met: + * + * * Redistributions of source code must retain the above copyright + * notice, this list of conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above + * copyright notice, this list of conditions and the following disclaimer + * in the documentation and/or other materials provided with the + * distribution. + * * Neither the name of Google LLC nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS + * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT + * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR + * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT + * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, + * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT + * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, + * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY + * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT + * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + */ +package com.google.api.gax.rpc; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.atLeastOnce; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.google.api.core.SettableApiFuture; +import com.google.api.gax.retrying.RetrySettings; +import com.google.api.gax.rpc.testing.FakeCallContext; +import org.junit.Rule; +import org.junit.Test; +import org.mockito.Mock; +import org.mockito.Spy; +import org.mockito.junit.MockitoJUnit; +import org.mockito.junit.MockitoRule; +import org.mockito.quality.Strictness; +import org.threeten.bp.Duration; + +public class CallableTest { + + @Rule + public final MockitoRule mockitoRule = MockitoJUnit.rule().strictness(Strictness.STRICT_STUBS); + + @Mock private UnaryCallable innerCallable; + private SettableApiFuture innerResult; + + @Mock private ServerStreamingCallable innerServerStreamingCallable; + + private RetrySettings retrySettings = + RetrySettings.newBuilder() + .setInitialRpcTimeout(Duration.ofMillis(5L)) + .setMaxRpcTimeout(Duration.ofMillis(5L)) + .setTotalTimeout(Duration.ofMillis(10L)) + .build(); + + @Spy private ApiCallContext callContext = FakeCallContext.createDefault(); + + @Spy + private ApiCallContext callContextWithRetrySettings = + FakeCallContext.createDefault().withRetrySettings(retrySettings); + + private ClientContext clientContext = + ClientContext.newBuilder().setDefaultCallContext(callContext).build(); + + @Test + public void testNonRetriedCallable() throws Exception { + innerResult = SettableApiFuture.create(); + when(innerCallable.futureCall(anyString(), any(ApiCallContext.class))).thenReturn(innerResult); + Duration timeout = Duration.ofMillis(5L); + + UnaryCallSettings callSettings = + UnaryCallSettings.newUnaryCallSettingsBuilder().setSimpleTimeoutNoRetries(timeout).build(); + UnaryCallable callable = + Callables.retrying(innerCallable, callSettings, clientContext); + innerResult.set("No, my refrigerator is not running!"); + + callable.futureCall("Is your refrigerator running?", callContext); + verify(callContext, atLeastOnce()).getRetrySettings(); + verify(callContext).getTimeout(); + verify(callContext).withTimeout(timeout); + } + + @Test + public void testNonRetriedCallableWithRetrySettings() throws Exception { + innerResult = SettableApiFuture.create(); + when(innerCallable.futureCall(anyString(), any(ApiCallContext.class))).thenReturn(innerResult); + + UnaryCallSettings callSettings = + UnaryCallSettings.newUnaryCallSettingsBuilder() + .setSimpleTimeoutNoRetries(Duration.ofMillis(10L)) + .build(); + UnaryCallable callable = + Callables.retrying(innerCallable, callSettings, clientContext); + innerResult.set("No, my refrigerator is not running!"); + + Duration timeout = retrySettings.getInitialRpcTimeout(); + + callable.futureCall("Is your refrigerator running?", callContextWithRetrySettings); + + verify(callContextWithRetrySettings, atLeastOnce()).getRetrySettings(); + verify(callContextWithRetrySettings).getTimeout(); + verify(callContextWithRetrySettings).withTimeout(timeout); + } + + @Test + public void testNonRetriedServerStreamingCallable() throws Exception { + Duration timeout = Duration.ofMillis(5L); + ServerStreamingCallSettings callSettings = + ServerStreamingCallSettings.newBuilder().setSimpleTimeoutNoRetries(timeout).build(); + ServerStreamingCallable callable = + Callables.retrying(innerServerStreamingCallable, callSettings, clientContext); + + callable.call("Is your refrigerator running?", callContext); + + verify(callContext, atLeastOnce()).getRetrySettings(); + verify(callContext).getTimeout(); + verify(callContext).withTimeout(timeout); + } + + @Test + public void testNonRetriedServerStreamingCallableWithRetrySettings() throws Exception { + ServerStreamingCallSettings callSettings = + ServerStreamingCallSettings.newBuilder() + .setSimpleTimeoutNoRetries(Duration.ofMillis(10L)) + .build(); + ServerStreamingCallable callable = + Callables.retrying(innerServerStreamingCallable, callSettings, clientContext); + + Duration timeout = retrySettings.getTotalTimeout(); + + callable.call("Is your refrigerator running?", callContextWithRetrySettings); + + verify(callContextWithRetrySettings, atLeastOnce()).getRetrySettings(); + verify(callContextWithRetrySettings).getTimeout(); + verify(callContextWithRetrySettings).withTimeout(timeout); + } +} diff --git a/gax/src/test/java/com/google/api/gax/tracing/TracedCallableTest.java b/gax/src/test/java/com/google/api/gax/tracing/TracedCallableTest.java new file mode 100644 index 000000000..bb96d4543 --- /dev/null +++ b/gax/src/test/java/com/google/api/gax/tracing/TracedCallableTest.java @@ -0,0 +1,119 @@ +/* + * Copyright 2021 Google LLC + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are + * met: + * + * * Redistributions of source code must retain the above copyright + * notice, this list of conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above + * copyright notice, this list of conditions and the following disclaimer + * in the documentation and/or other materials provided with the + * distribution. + * * Neither the name of Google LLC nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS + * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT + * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR + * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT + * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, + * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT + * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, + * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY + * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT + * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + */ +package com.google.api.gax.tracing; + +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.Mockito.any; +import static org.mockito.Mockito.anyString; +import static org.mockito.Mockito.eq; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoMoreInteractions; +import static org.mockito.Mockito.when; + +import com.google.api.core.ApiFuture; +import com.google.api.core.SettableApiFuture; +import com.google.api.gax.rpc.ApiCallContext; +import com.google.api.gax.rpc.Callables; +import com.google.api.gax.rpc.ClientContext; +import com.google.api.gax.rpc.UnaryCallSettings; +import com.google.api.gax.rpc.UnaryCallable; +import com.google.api.gax.rpc.testing.FakeCallContext; +import com.google.api.gax.tracing.ApiTracerFactory.OperationType; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.mockito.Mock; +import org.mockito.junit.MockitoJUnit; +import org.mockito.junit.MockitoRule; +import org.mockito.quality.Strictness; +import org.threeten.bp.Duration; + +@RunWith(JUnit4.class) +public class TracedCallableTest { + private static final SpanName SPAN_NAME = SpanName.of("FakeClient", "FakeRpc"); + + @Rule + public final MockitoRule mockitoRule = MockitoJUnit.rule().strictness(Strictness.STRICT_STUBS); + + @Mock private ApiTracerFactory tracerFactory; + private ApiTracer parentTracer; + @Mock private ApiTracer tracer; + @Mock private UnaryCallable innerCallable; + private SettableApiFuture innerResult; + + private ApiCallContext callContext; + private ClientContext clientContext; + + @Before + public void setUp() { + parentTracer = NoopApiTracer.getInstance(); + + // Wire the mock tracer factory + when(tracerFactory.newTracer( + any(ApiTracer.class), any(SpanName.class), eq(OperationType.Unary))) + .thenReturn(tracer); + + // Wire the mock inner callable + innerResult = SettableApiFuture.create(); + when(innerCallable.futureCall(anyString(), any(ApiCallContext.class))).thenReturn(innerResult); + + callContext = FakeCallContext.createDefault(); + clientContext = ClientContext.newBuilder().setDefaultCallContext(callContext).build(); + } + + public UnaryCallable setupTracedUnaryCallable( + UnaryCallSettings callSettings) { + UnaryCallable callable = + Callables.retrying(innerCallable, callSettings, clientContext); + return new TracedUnaryCallable<>(callable, tracerFactory, SPAN_NAME); + } + + @Test + public void testNonRetriedCallable() throws Exception { + // Verify that callables configured to not retry have the appropriate tracer interactions. + UnaryCallSettings callSettings = + UnaryCallSettings.newUnaryCallSettingsBuilder() + .setSimpleTimeoutNoRetries(Duration.ofMillis(5L)) + .build(); + UnaryCallable callable = setupTracedUnaryCallable(callSettings); + innerResult.set("No, my refrigerator is not running!"); + + ApiFuture future = callable.futureCall("Is your refrigerator running?", callContext); + + verify(tracerFactory, times(1)).newTracer(parentTracer, SPAN_NAME, OperationType.Unary); + verify(tracer, times(1)).attemptStarted(anyInt()); + verify(tracer, times(1)).attemptSucceeded(); + verify(tracer, times(1)).operationSucceeded(); + verifyNoMoreInteractions(tracer); + } +}