fix: retain context timeouts in ServerStreamingAttemptCallable (#1155)

Author: noahdietz

gax/src/main/java/com/google/api/gax/rpc/ServerStreamingAttemptCallable.java

@@ -181,11 +181,12 @@ public void cancel() {
  }
  isStarted = true;
 
- // Propagate the totalTimeout as the overall stream deadline.
+ // Propagate the totalTimeout as the overall stream deadline, so long as the user
+ // has not provided a timeout via the ApiCallContext. If they have, retain it.
  Duration totalTimeout =
  outerRetryingFuture.getAttemptSettings().getGlobalSettings().getTotalTimeout();
 
- if (totalTimeout != null && context != null) {
+ if (totalTimeout != null && context != null && context.getTimeout() == null) {
  context = context.withTimeout(totalTimeout);
  }
 
@@ -217,7 +218,10 @@ public Void call() {
 
  ApiCallContext attemptContext = context;
 
- if (!outerRetryingFuture.getAttemptSettings().getRpcTimeout().isZero()) {
+ // Set the streamWaitTimeout to the attempt RPC Timeout, only if the context
+ // does not already have a timeout set by a user via withStreamWaitTimeout.
+ if (!outerRetryingFuture.getAttemptSettings().getRpcTimeout().isZero()
+ && attemptContext.getStreamWaitTimeout() == null) {
  attemptContext =
  attemptContext.withStreamWaitTimeout(
  outerRetryingFuture.getAttemptSettings().getRpcTimeout());

gax/src/test/java/com/google/api/gax/rpc/ServerStreamingAttemptCallableTest.java

@@ -41,6 +41,7 @@
 import com.google.api.gax.rpc.testing.FakeCallContext;
 import com.google.api.gax.rpc.testing.MockStreamingApi.MockServerStreamingCall;
 import com.google.api.gax.rpc.testing.MockStreamingApi.MockServerStreamingCallable;
+import com.google.api.gax.tracing.NoopApiTracer;
 import com.google.common.collect.Queues;
 import com.google.common.truth.Truth;
 import java.util.concurrent.BlockingDeque;
@@ -51,6 +52,7 @@
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
+import org.mockito.Mockito;
 import org.threeten.bp.Duration;
 
 @RunWith(JUnit4.class)
@@ -59,29 +61,107 @@ public class ServerStreamingAttemptCallableTest {
  private AccumulatingObserver observer;
  private FakeRetryingFuture fakeRetryingFuture;
  private StreamResumptionStrategy<String, String> resumptionStrategy;
+ private static Duration totalTimeout = Duration.ofHours(1);
+ private FakeCallContext mockedCallContext;
 
  @Before
  public void setUp() {
  innerCallable = new MockServerStreamingCallable<>();
  observer = new AccumulatingObserver(true);
  resumptionStrategy = new MyStreamResumptionStrategy();
+ mockedCallContext = Mockito.mock(FakeCallContext.class);
  }
 
  private ServerStreamingAttemptCallable<String, String> createCallable() {
+ return createCallable(FakeCallContext.createDefault());
+ }
+
+ private ServerStreamingAttemptCallable<String, String> createCallable(ApiCallContext context) {
  ServerStreamingAttemptCallable<String, String> callable =
  new ServerStreamingAttemptCallable<>(
- innerCallable,
- resumptionStrategy,
- "request",
- FakeCallContext.createDefault(),
- observer);
+ innerCallable, resumptionStrategy, "request", context, observer);
 
  fakeRetryingFuture = new FakeRetryingFuture(callable);
  callable.setExternalFuture(fakeRetryingFuture);
 
  return callable;
  }
 
+ @Test
+ public void testUserProvidedContextTimeout() {
+ // Mock up the ApiCallContext as if the user provided a timeout and streamWaitTimeout.
+ Mockito.doReturn(NoopApiTracer.getInstance()).when(mockedCallContext).getTracer();
+ Mockito.doReturn(Duration.ofHours(5)).when(mockedCallContext).getTimeout();
+ Mockito.doReturn(Duration.ofHours(5)).when(mockedCallContext).getStreamWaitTimeout();
+
+ ServerStreamingAttemptCallable<String, String> callable = createCallable(mockedCallContext);
+ callable.start();
+
+ // Ensure that the callable did not overwrite the user provided timeouts
+ Mockito.verify(mockedCallContext, Mockito.times(1)).getTimeout();
+ Mockito.verify(mockedCallContext, Mockito.never()).withTimeout(totalTimeout);
+ Mockito.verify(mockedCallContext, Mockito.times(1)).getStreamWaitTimeout();
+ Mockito.verify(mockedCallContext, Mockito.never())
+ .withStreamWaitTimeout(Mockito.any(Duration.class));
+
+ // Should notify outer observer
+ Truth.assertThat(observer.controller).isNotNull();
+
+ // Should configure the inner controller correctly.
+ MockServerStreamingCall<String, String> call = innerCallable.popLastCall();
+ Truth.assertThat(call.getController().isAutoFlowControlEnabled()).isTrue();
+ Truth.assertThat(call.getRequest()).isEqualTo("request");
+
+ // Send a response in auto flow mode.
+ call.getController().getObserver().onResponse("response1");
+ call.getController().getObserver().onResponse("response2");
+ call.getController().getObserver().onComplete();
+
+ // Make sure the responses are received
+ Truth.assertThat(observer.responses).containsExactly("response1", "response2").inOrder();
+ fakeRetryingFuture.assertSuccess();
+ }
+
+ @Test
+ public void testNoUserProvidedContextTimeout() {
+ // Mock up the ApiCallContext as if the user did not provide custom timeouts.
+ Mockito.doReturn(NoopApiTracer.getInstance()).when(mockedCallContext).getTracer();
+ Mockito.doReturn(null).when(mockedCallContext).getTimeout();
+ Mockito.doReturn(null).when(mockedCallContext).getStreamWaitTimeout();
+ Mockito.doReturn(mockedCallContext).when(mockedCallContext).withTimeout(totalTimeout);
+ Mockito.doReturn(mockedCallContext)
+ .when(mockedCallContext)
+ .withStreamWaitTimeout(Mockito.any(Duration.class));
+
+ ServerStreamingAttemptCallable<String, String> callable = createCallable(mockedCallContext);
+ callable.start();
+
+ // Ensure that the callable configured the timeouts via the Settings in the
+ // absence of user-defined timeouts.
+ Mockito.verify(mockedCallContext, Mockito.times(1)).getTimeout();
+ Mockito.verify(mockedCallContext, Mockito.times(1)).withTimeout(totalTimeout);
+ Mockito.verify(mockedCallContext, Mockito.times(1)).getStreamWaitTimeout();
+ Mockito.verify(mockedCallContext, Mockito.times(1))
+ .withStreamWaitTimeout(Mockito.any(Duration.class));
+
+ // Should notify outer observer
+ Truth.assertThat(observer.controller).isNotNull();
+
+ // Should configure the inner controller correctly.
+ MockServerStreamingCall<String, String> call = innerCallable.popLastCall();
+ Truth.assertThat(call.getController().isAutoFlowControlEnabled()).isTrue();
+ Truth.assertThat(call.getRequest()).isEqualTo("request");
+
+ // Send a response in auto flow mode.
+ call.getController().getObserver().onResponse("response1");
+ call.getController().getObserver().onResponse("response2");
+ call.getController().getObserver().onComplete();
+
+ // Make sure the responses are received
+ Truth.assertThat(observer.responses).containsExactly("response1", "response2").inOrder();
+ fakeRetryingFuture.assertSuccess();
+ }
+
  @Test
  public void testNoErrorsAutoFlow() {
  ServerStreamingAttemptCallable<String, String> callable = createCallable();
@@ -396,8 +476,7 @@ private static class FakeRetryingFuture extends AbstractApiFuture<Void>
  this.attemptCallable = attemptCallable;
  attemptSettings =
  TimedAttemptSettings.newBuilder()
- .setGlobalSettings(
- RetrySettings.newBuilder().setTotalTimeout(Duration.ofHours(1)).build())
+ .setGlobalSettings(RetrySettings.newBuilder().setTotalTimeout(totalTimeout).build())
  .setFirstAttemptStartTimeNanos(0)
  .setAttemptCount(0)
  .setOverallAttemptCount(0)