m-trieu commented on code in PR #32774:
URL: https://github.com/apache/beam/pull/32774#discussion_r1813325167


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/AbstractWindmillStream.java:
##########
@@ -278,21 +354,67 @@ public String backendWorkerToken() {
   }
 
   @Override
-  public void shutdown() {
-    if (isShutdown.compareAndSet(false, true)) {
-      requestObserver()
-          .onError(new WindmillStreamShutdownException("Explicit call to 
shutdown stream."));
+  public final void shutdown() {
+    // Don't lock on "this" as isShutdown checks are used in the stream to 
free blocked
+    // threads or as exit conditions to loops.
+    synchronized (shutdownLock) {
+      if (!isShutdown) {
+        isShutdown = true;
+        shutdownTime.set(DateTime.now());
+        requestObserver.onError(
+            new WindmillStreamShutdownException("Explicit call to shutdown 
stream."));
+        shutdownInternal();
+      }
     }
   }
 
-  private void setLastError(String error) {
-    lastError.set(error);
-    lastErrorTime.set(DateTime.now());
+  private void recordRestartReason(String error) {
+    lastRestartReason.set(error);
+    lastRestartTime.set(DateTime.now());
   }
 
-  public static class WindmillStreamShutdownException extends RuntimeException 
{
-    public WindmillStreamShutdownException(String message) {
-      super(message);
+  protected abstract void shutdownInternal();
+
+  /**
+   * Request observer that allows resetting its internal delegate using the 
given {@link
+   * #requestObserverSupplier}.
+   */
+  @ThreadSafe
+  private static class ResettableRequestObserver<RequestT> implements 
StreamObserver<RequestT> {
+
+    private final Supplier<StreamObserver<RequestT>> requestObserverSupplier;
+
+    @GuardedBy("this")
+    private @Nullable StreamObserver<RequestT> delegateRequestObserver;
+
+    private ResettableRequestObserver(Supplier<StreamObserver<RequestT>> 
requestObserverSupplier) {
+      this.requestObserverSupplier = requestObserverSupplier;
+      this.delegateRequestObserver = null;
+    }
+
+    private synchronized StreamObserver<RequestT> delegate() {
+      return Preconditions.checkNotNull(
+          delegateRequestObserver,
+          "requestObserver cannot be null. Missing a call to startStream() to 
initialize.");
+    }
+
+    private synchronized void reset() {
+      delegateRequestObserver = requestObserverSupplier.get();
+    }
+
+    @Override
+    public void onNext(RequestT requestT) {
+      delegate().onNext(requestT);

Review Comment:
   the delegate() request observer has its own internal locking
   added comment



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to