[Dataflow Streaming] Refactor ResettableThrowingStreamObserver to inject instead of use factory (#34940)

Author: scwhittle

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/AbstractWindmillStream.java

@@ -28,9 +28,11 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Function;
+import java.util.function.Supplier;
 import javax.annotation.concurrent.GuardedBy;
 import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.observers.StreamObserverCancelledException;
 import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.observers.StreamObserverFactory;
+import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.observers.TerminatingStreamObserver;
 import org.apache.beam.sdk.util.BackOff;
 import org.apache.beam.vendor.grpc.v1p69p0.com.google.api.client.util.Sleeper;
 import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.Status;
@@ -83,6 +85,8 @@ public abstract class AbstractWindmillStream<RequestT, ResponseT> implements Win
   private final int logEveryNStreamFailures;
   private final String backendWorkerToken;
   private final ResettableThrowingStreamObserver<RequestT> requestObserver;
+
+  private final Supplier<TerminatingStreamObserver<RequestT>> requestObserverFactory;
   private final StreamDebugMetrics debugMetrics;
   private final AtomicBoolean isHealthCheckScheduled;
 
@@ -120,13 +124,11 @@ protected AbstractWindmillStream(
     this.isHealthCheckScheduled = new AtomicBoolean(false);
     this.finishLatch = new CountDownLatch(1);
     this.logger = logger;
-    this.requestObserver =
-        new ResettableThrowingStreamObserver<>(
-            () ->
-                streamObserverFactory.from(
-                    clientFactory,
-                    new AbstractWindmillStream<RequestT, ResponseT>.ResponseObserver()),
-            logger);
+    this.requestObserver = new ResettableThrowingStreamObserver<>(logger);
+    this.requestObserverFactory =
+        () ->
+            streamObserverFactory.from(
+                clientFactory, new AbstractWindmillStream<RequestT, ResponseT>.ResponseObserver());
     this.sleeper = Sleeper.DEFAULT;
     this.debugMetrics = StreamDebugMetrics.create();
   }
@@ -184,7 +186,7 @@ private void startStream() {
       try {
         synchronized (this) {
           debugMetrics.recordStart();
-          requestObserver.reset();
+          requestObserver.reset(requestObserverFactory.get());
           onNewStream();
           if (clientClosed) {
             halfClose();

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/ResettableThrowingStreamObserver.java

@@ -17,7 +17,6 @@
  */
 package org.apache.beam.runners.dataflow.worker.windmill.client;
 
-import java.util.function.Supplier;
 import javax.annotation.Nullable;
 import javax.annotation.concurrent.GuardedBy;
 import javax.annotation.concurrent.ThreadSafe;
@@ -29,7 +28,7 @@
 import org.slf4j.Logger;
 
 /**
- * Request observer that allows resetting its internal delegate using the given {@link
+ * Request observer that allows resetting its internal delegate using a {@link
  * #streamObserverFactory}.
  *
  * @implNote {@link StreamObserver}s generated by {@link #streamObserverFactory} are expected to be
@@ -40,7 +39,7 @@
 @ThreadSafe
 @Internal
 final class ResettableThrowingStreamObserver<T> {
-  private final Supplier<TerminatingStreamObserver<T>> streamObserverFactory;
+
   private final Logger logger;
 
   @GuardedBy("this")
@@ -57,9 +56,7 @@ final class ResettableThrowingStreamObserver<T> {
   @GuardedBy("this")
   private boolean isCurrentStreamClosed = true;
 
-  ResettableThrowingStreamObserver(
-      Supplier<TerminatingStreamObserver<T>> streamObserverFactory, Logger logger) {
-    this.streamObserverFactory = streamObserverFactory;
+  ResettableThrowingStreamObserver(Logger logger) {
     this.logger = logger;
     this.delegateStreamObserver = null;
   }
@@ -79,12 +76,13 @@ private synchronized StreamObserver<T> delegate()
   }
 
   /** Creates a new delegate to use for future {@link StreamObserver} methods. */
-  synchronized void reset() throws WindmillStreamShutdownException {
+  synchronized void reset(TerminatingStreamObserver<T> observer)
+      throws WindmillStreamShutdownException {
     if (isPoisoned) {
       throw new WindmillStreamShutdownException("Stream is already shutdown.");
     }
 
-    delegateStreamObserver = streamObserverFactory.get();
+    delegateStreamObserver = observer;
     isCurrentStreamClosed = false;
   }
 

runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/ResettableThrowingStreamObserverTest.java

@@ -24,11 +24,7 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoInteractions;
 
-import java.util.ArrayList;
-import java.util.List;
-import java.util.function.Supplier;
 import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.observers.TerminatingStreamObserver;
-import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.StreamObserver;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
@@ -57,36 +53,36 @@ public void terminate(Throwable terminationException) {}
 
   @Test
   public void testPoison_beforeDelegateSet() {
-    ResettableThrowingStreamObserver<Integer> observer = newStreamObserver(() -> delegate);
+    ResettableThrowingStreamObserver<Integer> observer = newStreamObserver();
     observer.poison();
     verifyNoInteractions(delegate);
   }
 
   @Test
   public void testPoison_afterDelegateSet() throws WindmillStreamShutdownException {
-    ResettableThrowingStreamObserver<Integer> observer = newStreamObserver(() -> delegate);
-    observer.reset();
+    ResettableThrowingStreamObserver<Integer> observer = newStreamObserver();
+    observer.reset(delegate);
     observer.poison();
     verify(delegate).terminate(isA(WindmillStreamShutdownException.class));
   }
 
   @Test
   public void testReset_afterPoisonedThrows() {
-    ResettableThrowingStreamObserver<Integer> observer = newStreamObserver(() -> delegate);
+    ResettableThrowingStreamObserver<Integer> observer = newStreamObserver();
     observer.poison();
-    assertThrows(WindmillStreamShutdownException.class, observer::reset);
+    assertThrows(WindmillStreamShutdownException.class, () -> observer.reset(newDelegate()));
   }
 
   @Test
   public void testOnNext_afterPoisonedThrows() {
-    ResettableThrowingStreamObserver<Integer> observer = newStreamObserver(() -> delegate);
+    ResettableThrowingStreamObserver<Integer> observer = newStreamObserver();
     observer.poison();
     assertThrows(WindmillStreamShutdownException.class, () -> observer.onNext(1));
   }
 
   @Test
   public void testOnError_afterPoisonedThrows() {
-    ResettableThrowingStreamObserver<Integer> observer = newStreamObserver(() -> delegate);
+    ResettableThrowingStreamObserver<Integer> observer = newStreamObserver();
     observer.poison();
     assertThrows(
         WindmillStreamShutdownException.class,
@@ -95,7 +91,7 @@ public void testOnError_afterPoisonedThrows() {
 
   @Test
   public void testOnCompleted_afterPoisonedThrows() {
-    ResettableThrowingStreamObserver<Integer> observer = newStreamObserver(() -> delegate);
+    ResettableThrowingStreamObserver<Integer> observer = newStreamObserver();
     observer.poison();
     assertThrows(WindmillStreamShutdownException.class, observer::onCompleted);
   }
@@ -104,28 +100,20 @@ public void testOnCompleted_afterPoisonedThrows() {
   public void testReset_usesNewDelegate()
       throws WindmillStreamShutdownException,
           ResettableThrowingStreamObserver.StreamClosedException {
-    List<StreamObserver<Integer>> delegates = new ArrayList<>();
-    ResettableThrowingStreamObserver<Integer> observer =
-        newStreamObserver(
-            () -> {
-              TerminatingStreamObserver<Integer> delegate = newDelegate();
-              delegates.add(delegate);
-              return delegate;
-            });
-    observer.reset();
+    ResettableThrowingStreamObserver<Integer> observer = newStreamObserver();
+    TerminatingStreamObserver<Integer> firstObserver = newDelegate();
+    observer.reset(firstObserver);
     observer.onNext(1);
-    observer.reset();
-    observer.onNext(2);
 
-    StreamObserver<Integer> firstObserver = delegates.get(0);
-    StreamObserver<Integer> secondObserver = delegates.get(1);
+    TerminatingStreamObserver<Integer> secondObserver = newDelegate();
+    observer.reset(secondObserver);
+    observer.onNext(2);
 
     verify(firstObserver).onNext(eq(1));
     verify(secondObserver).onNext(eq(2));
   }
 
-  private <T> ResettableThrowingStreamObserver<T> newStreamObserver(
-      Supplier<TerminatingStreamObserver<T>> delegate) {
-    return new ResettableThrowingStreamObserver<>(delegate, LoggerFactory.getLogger(getClass()));
+  private <T> ResettableThrowingStreamObserver<T> newStreamObserver() {
+    return new ResettableThrowingStreamObserver<>(LoggerFactory.getLogger(getClass()));
   }
 }