address PR comments

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java

@@ -75,9 +75,8 @@
 import org.apache.beam.runners.dataflow.worker.windmill.client.commits.WorkCommitter;
 import org.apache.beam.runners.dataflow.worker.windmill.client.getdata.ApplianceGetDataClient;
 import org.apache.beam.runners.dataflow.worker.windmill.client.getdata.GetDataClient;
-import org.apache.beam.runners.dataflow.worker.windmill.client.getdata.StreamingEngineGetDataClient;
+import org.apache.beam.runners.dataflow.worker.windmill.client.getdata.StreamPoolGetDataClient;
 import org.apache.beam.runners.dataflow.worker.windmill.client.getdata.ThrottlingGetDataMetricTracker;
-import org.apache.beam.runners.dataflow.worker.windmill.client.getdata.WorkRefreshClient;
 import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.ChannelzServlet;
 import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.GrpcDispatcherClient;
 import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.GrpcWindmillServer;
@@ -217,12 +216,16 @@ private StreamingDataflowWorker(
 
     this.workCommitter =
         windmillServiceEnabled
-            ? StreamingEngineWorkCommitter.create(
-                WindmillStreamPool.create(
-                        numCommitThreads, COMMIT_STREAM_TIMEOUT, windmillServer::commitWorkStream)
-                    ::getCloseableStream,
-                numCommitThreads,
-                this::onCompleteCommit)
+            ? StreamingEngineWorkCommitter.builder()
+                .setCommitWorkStreamFactory(
+                    WindmillStreamPool.create(
+                            numCommitThreads,
+                            COMMIT_STREAM_TIMEOUT,
+                            windmillServer::commitWorkStream)
+                        ::getCloseableStream)
+                .setNumCommitSenders(numCommitThreads)
+                .setOnCommitComplete(this::onCompleteCommit)
+                .build()
             : StreamingApplianceWorkCommitter.create(
                 windmillServer::commitWork, this::onCompleteCommit);
 
@@ -252,31 +255,26 @@ private StreamingDataflowWorker(
     ThrottlingGetDataMetricTracker getDataMetricTracker =
         new ThrottlingGetDataMetricTracker(memoryMonitor);
 
-    WindmillStreamPool<GetDataStream> getDataStreamPool =
-        WindmillStreamPool.create(
-            Math.max(1, options.getWindmillGetDataStreamCount()),
-            GET_DATA_STREAM_TIMEOUT,
-            windmillServer::getDataStream);
-
-    // Register standard file systems.
-    FileSystems.setDefaultPipelineOptions(options);
-
-    int stuckCommitDurationMillis =
-        windmillServiceEnabled && options.getStuckCommitDurationMillis() > 0
-            ? options.getStuckCommitDurationMillis()
-            : 0;
-
-    WorkRefreshClient workRefreshClient;
+    int stuckCommitDurationMillis;
     if (windmillServiceEnabled) {
-      StreamingEngineGetDataClient streamingEngineGetDataClient =
-          new StreamingEngineGetDataClient(getDataMetricTracker, getDataStreamPool);
-      this.getDataClient = streamingEngineGetDataClient;
-      workRefreshClient = streamingEngineGetDataClient;
+      WindmillStreamPool<GetDataStream> getDataStreamPool =
+          WindmillStreamPool.create(
+              Math.max(1, options.getWindmillGetDataStreamCount()),
+              GET_DATA_STREAM_TIMEOUT,
+              windmillServer::getDataStream);
+      this.getDataClient = new StreamPoolGetDataClient(getDataMetricTracker, getDataStreamPool);
+      this.heartbeatSender =
+          new StreamPoolHeartbeatSender(
+              options.getUseSeparateWindmillHeartbeatStreams()
+                  ? WindmillStreamPool.create(
+                      1, GET_DATA_STREAM_TIMEOUT, windmillServer::getDataStream)
+                  : getDataStreamPool);
+      stuckCommitDurationMillis =
+          options.getStuckCommitDurationMillis() > 0 ? options.getStuckCommitDurationMillis() : 0;
     } else {
-      ApplianceGetDataClient applianceGetDataClient =
-          new ApplianceGetDataClient(windmillServer, getDataMetricTracker);
-      this.getDataClient = applianceGetDataClient;
-      workRefreshClient = applianceGetDataClient;
+      this.getDataClient = new ApplianceGetDataClient(windmillServer, getDataMetricTracker);
+      this.heartbeatSender = new ApplianceHeartbeatSender(windmillServer::getData);
+      stuckCommitDurationMillis = 0;
     }
 
     this.activeWorkRefresher =
@@ -287,7 +285,7 @@ private StreamingDataflowWorker(
             computationStateCache::getAllPresentComputations,
             sampler,
             executorSupplier.apply("RefreshWork"),
-            workRefreshClient::refreshActiveWork);
+            getDataMetricTracker::trackHeartbeats);
 
     WorkerStatusPages workerStatusPages =
         WorkerStatusPages.create(DEFAULT_STATUS_PORT, memoryMonitor);
@@ -333,14 +331,8 @@ private StreamingDataflowWorker(
             ID_GENERATOR,
             stageInfoMap);
 
-    this.heartbeatSender =
-        options.isEnableStreamingEngine()
-            ? new StreamPoolHeartbeatSender(
-                options.getUseSeparateWindmillHeartbeatStreams()
-                    ? WindmillStreamPool.create(
-                        1, GET_DATA_STREAM_TIMEOUT, windmillServer::getDataStream)
-                    : getDataStreamPool)
-            : new ApplianceHeartbeatSender(windmillServer::getData);
+    // Register standard file systems.
+    FileSystems.setDefaultPipelineOptions(options);
 
     LOG.debug("windmillServiceEnabled: {}", windmillServiceEnabled);
     LOG.debug("WindmillServiceEndpoint: {}", options.getWindmillServiceEndpoint());

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WorkItemCancelledException.java

@@ -26,6 +26,10 @@ public WorkItemCancelledException(long sharding_key) {
     super("Work item cancelled for key " + sharding_key);
   }
 
+  public WorkItemCancelledException(Throwable e) {
+    super(e);
+  }
+
   /** Returns whether an exception was caused by a {@link WorkItemCancelledException}. */
   public static boolean isWorkItemCancelledException(Throwable t) {
     while (t != null) {

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ActiveWorkState.java

@@ -206,14 +206,17 @@ synchronized ImmutableListMultimap<ShardedKey, RefreshableWork> getReadOnlyActiv
         .collect(
             flatteningToImmutableListMultimap(
                 Entry::getKey,
-                e -> e.getValue().stream().map(ExecutableWork::work).map(Work::refreshableView)));
+                e ->
+                    e.getValue().stream()
+                        .map(ExecutableWork::work)
+                        .map(work -> (RefreshableWork) work)));
   }
 
   synchronized ImmutableList<RefreshableWork> getRefreshableWork(Instant refreshDeadline) {
     return activeWork.values().stream()
         .flatMap(Deque::stream)
         .map(ExecutableWork::work)
-        .filter(work -> work.isRefreshable(refreshDeadline))
+        .filter(work -> !work.isFailed() && work.getStartTime().isBefore(refreshDeadline))
         .collect(toImmutableList());
   }
 

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/RefreshableWork.java

@@ -22,7 +22,6 @@
 import org.apache.beam.runners.dataflow.worker.windmill.work.refresh.HeartbeatSender;
 import org.apache.beam.sdk.annotations.Internal;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
-import org.joda.time.Instant;
 
 /** View of {@link Work} that exposes an interface for work refreshing. */
 @Internal
@@ -32,8 +31,6 @@ public interface RefreshableWork {
 
   ShardedKey getShardedKey();
 
-  boolean isRefreshable(Instant refreshDeadline);
-
   HeartbeatSender heartbeatSender();
 
   ImmutableList<Windmill.LatencyAttribution> getHeartbeatLatencyAttributions(

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/Work.java

@@ -150,10 +150,6 @@ private static LatencyAttribution.Builder createLatencyAttributionWithActiveLate
     return latencyAttribution;
   }
 
-  public RefreshableWork refreshableView() {
-    return this;
-  }
-
   public WorkItem getWorkItem() {
     return workItem;
   }
@@ -209,11 +205,6 @@ public String getLatencyTrackingId() {
     return latencyTrackingId;
   }
 
-  @Override
-  public boolean isRefreshable(Instant refreshDeadline) {
-    return !isFailed && getStartTime().isBefore(refreshDeadline);
-  }
-
   @Override
   public HeartbeatSender heartbeatSender() {
     return processingContext.heartbeatSender();

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/sideinput/SideInputStateFetcher.java

@@ -34,6 +34,7 @@
 import org.apache.beam.runners.dataflow.worker.windmill.Windmill;
 import org.apache.beam.runners.dataflow.worker.windmill.Windmill.GlobalData;
 import org.apache.beam.runners.dataflow.worker.windmill.Windmill.GlobalDataRequest;
+import org.apache.beam.sdk.annotations.Internal;
 import org.apache.beam.sdk.coders.Coder;
 import org.apache.beam.sdk.coders.IterableCoder;
 import org.apache.beam.sdk.coders.KvCoder;
@@ -52,6 +53,7 @@
 
 /** Class responsible for fetching side input state from the streaming backend. */
 @NotThreadSafe
+@Internal
 public class SideInputStateFetcher {
   private static final Logger LOG = LoggerFactory.getLogger(SideInputStateFetcher.class);
 

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

@@ -85,10 +85,11 @@ public abstract class AbstractWindmillStream<RequestT, ResponseT> implements Win
   private final Supplier<StreamObserver<RequestT>> requestObserverSupplier;
   // Indicates if the current stream in requestObserver is closed by calling close() method
   private final AtomicBoolean streamClosed;
-  private @Nullable StreamObserver<RequestT> requestObserver;
   private final String backendWorkerToken;
+  private @Nullable StreamObserver<RequestT> requestObserver;
 
   protected AbstractWindmillStream(
+      String debugStreamType,
       Function<StreamObserver<ResponseT>, StreamObserver<RequestT>> clientFactory,
       BackOff backoff,
       StreamObserverFactory streamObserverFactory,
@@ -100,7 +101,7 @@ protected AbstractWindmillStream(
         Executors.newSingleThreadExecutor(
             new ThreadFactoryBuilder()
                 .setDaemon(true)
-                .setNameFormat(createThreadName(streamType(), backendWorkerToken))
+                .setNameFormat(createThreadName(debugStreamType, backendWorkerToken))
                 .build());
     this.backoff = backoff;
     this.streamRegistry = streamRegistry;
@@ -122,10 +123,10 @@ protected AbstractWindmillStream(
                 clientFactory, new AbstractWindmillStream<RequestT, ResponseT>.ResponseObserver());
   }
 
-  private static String createThreadName(Type streamType, String backendWorkerToken) {
+  private static String createThreadName(String streamType, String backendWorkerToken) {
     return !backendWorkerToken.isEmpty()
-        ? String.format("%s-%s-WindmillStream-thread", streamType.name(), backendWorkerToken)
-        : String.format("%s-WindmillStream-thread", streamType.name());
+        ? String.format("%s-%s-WindmillStream-thread", streamType, backendWorkerToken)
+        : String.format("%s-WindmillStream-thread", streamType);
   }
 
   private static long debugDuration(long nowMs, long startMs) {
@@ -151,6 +152,11 @@ private static long debugDuration(long nowMs, long startMs) {
    */
   protected abstract void startThrottleTimer();
 
+  /** Reflects that {@link #shutdown()} was explicitly called. */
+  protected boolean isShutdown() {
+    return isShutdown.get();
+  }
+
   private StreamObserver<RequestT> requestObserver() {
     if (requestObserver == null) {
       throw new NullPointerException(
@@ -274,15 +280,11 @@ public String backendWorkerToken() {
   @Override
   public void shutdown() {
     if (isShutdown.compareAndSet(false, true)) {
-      halfClose();
+      requestObserver()
+          .onError(new WindmillStreamShutdownException("Explicit call to shutdown stream."));
     }
   }
 
-  @Override
-  public boolean isShutdown() {
-    return isShutdown.get();
-  }
-
   private void setLastError(String error) {
     lastError.set(error);
     lastErrorTime.set(DateTime.now());
@@ -313,7 +315,7 @@ public void onCompleted() {
 
     private void onStreamFinished(@Nullable Throwable t) {
       synchronized (this) {
-        if (clientClosed.get() && !hasPendingRequests()) {
+        if (isShutdown.get() || (clientClosed.get() && !hasPendingRequests())) {
           streamRegistry.remove(AbstractWindmillStream.this);
           finishLatch.countDown();
           return;

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

@@ -52,18 +52,6 @@ public interface WindmillStream {
    */
   void shutdown();
 
-  /** Reflects that {@link #shutdown()} was explicitly called. */
-  boolean isShutdown();
-
-  Type streamType();
-
-  enum Type {
-    GET_WORKER_METADATA,
-    GET_WORK,
-    GET_DATA,
-    COMMIT_WORK,
-  }
-
   /** Handle representing a stream of GetWork responses. */
   @ThreadSafe
   interface GetWorkStream extends WindmillStream {
@@ -72,11 +60,6 @@ interface GetWorkStream extends WindmillStream {
 
     /** Returns the remaining in-flight {@link GetWorkBudget}. */
     GetWorkBudget remainingBudget();
-
-    @Override
-    default Type streamType() {
-      return Type.GET_WORK;
-    }
   }
 
   /** Interface for streaming GetDataRequests to Windmill. */
@@ -93,11 +76,6 @@ Windmill.KeyedGetDataResponse requestKeyedData(
     void refreshActiveWork(Map<String, Collection<HeartbeatRequest>> heartbeats);
 
     void onHeartbeatResponse(List<Windmill.ComputationHeartbeatResponse> responses);
-
-    @Override
-    default Type streamType() {
-      return Type.GET_DATA;
-    }
   }
 
   /** Interface for streaming CommitWorkRequests to Windmill. */
@@ -109,11 +87,6 @@ interface CommitWorkStream extends WindmillStream {
      */
     CommitWorkStream.RequestBatcher batcher();
 
-    @Override
-    default Type streamType() {
-      return Type.COMMIT_WORK;
-    }
-
     @NotThreadSafe
     interface RequestBatcher extends Closeable {
       /**
@@ -140,10 +113,11 @@ default void close() {
 
   /** Interface for streaming GetWorkerMetadata requests to Windmill. */
   @ThreadSafe
-  interface GetWorkerMetadataStream extends WindmillStream {
-    @Override
-    default Type streamType() {
-      return Type.GET_WORKER_METADATA;
+  interface GetWorkerMetadataStream extends WindmillStream {}
+
+  class WindmillStreamShutdownException extends RuntimeException {
+    public WindmillStreamShutdownException(String message) {
+      super(message);
     }
   }
 }