Add WorkProvider interfaces and implementations (#31883)

apache · Aug 5, 2024 · ca744ae · ca744ae
1 parent 7e75087
commit ca744ae
Show file tree

Hide file tree

Showing 21 changed files with 839 additions and 608 deletions.
diff --git a/...worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java b/...worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java
diff --git a/...ll/client/grpc/StreamingEngineClient.java → ...s/FanOutStreamingEngineWorkerHarness.java b/...ll/client/grpc/StreamingEngineClient.java → ...s/FanOutStreamingEngineWorkerHarness.java
@@ -15,7 +15,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.beam.runners.dataflow.worker.windmill.client.grpc;
+package org.apache.beam.runners.dataflow.worker.streaming.harness;
 
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap.toImmutableMap;
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet.toImmutableSet;
@@ -47,6 +47,8 @@
 import org.apache.beam.runners.dataflow.worker.windmill.client.commits.WorkCommitter;
 import org.apache.beam.runners.dataflow.worker.windmill.client.getdata.StreamGetDataClient;
 import org.apache.beam.runners.dataflow.worker.windmill.client.getdata.ThrottlingGetDataMetricTracker;
+import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.GrpcDispatcherClient;
+import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.GrpcWindmillStreamFactory;
 import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs.ChannelCachingStubFactory;
 import org.apache.beam.runners.dataflow.worker.windmill.client.throttling.ThrottleTimer;
 import org.apache.beam.runners.dataflow.worker.windmill.work.WorkItemScheduler;
@@ -68,16 +70,19 @@
 import org.slf4j.LoggerFactory;
 
 /**
- * Client for StreamingEngine. Given a {@link GetWorkBudget}, divides the budget and starts the
- * {@link WindmillStream.GetWorkStream}(s).
+ * {@link StreamingWorkerHarness} implementation that manages fan out to multiple backend
+ * destinations. Given a {@link GetWorkBudget}, divides the budget and starts the {@link
+ * WindmillStream.GetWorkStream}(s).
  */
 @Internal
 @CheckReturnValue
 @ThreadSafe
-public final class StreamingEngineClient {
-  private static final Logger LOG = LoggerFactory.getLogger(StreamingEngineClient.class);
+public final class FanOutStreamingEngineWorkerHarness implements StreamingWorkerHarness {
+  private static final Logger LOG =
+      LoggerFactory.getLogger(FanOutStreamingEngineWorkerHarness.class);
   private static final String PUBLISH_NEW_WORKER_METADATA_THREAD = "PublishNewWorkerMetadataThread";
   private static final String CONSUME_NEW_WORKER_METADATA_THREAD = "ConsumeNewWorkerMetadataThread";
+
   private final JobHeader jobHeader;
   private final GrpcWindmillStreamFactory streamFactory;
   private final WorkItemScheduler workItemScheduler;
@@ -101,7 +106,7 @@ public final class StreamingEngineClient {
   private volatile boolean started;
 
   @SuppressWarnings("FutureReturnValueIgnored")
-  private StreamingEngineClient(
+  private FanOutStreamingEngineWorkerHarness(
       JobHeader jobHeader,
       GetWorkBudget totalGetWorkBudget,
       GrpcWindmillStreamFactory streamFactory,
@@ -152,23 +157,15 @@ private StreamingEngineClient(
 
   private static ExecutorService singleThreadedExecutorServiceOf(String threadName) {
     return Executors.newSingleThreadScheduledExecutor(
-        new ThreadFactoryBuilder()
-            .setNameFormat(threadName)
-            .setUncaughtExceptionHandler(
-                (t, e) -> {
-                  LOG.error(
-                      "{} failed due to uncaught exception during execution. ", t.getName(), e);
-                  throw new StreamingEngineClientException(e);
-                })
-            .build());
+        new ThreadFactoryBuilder().setNameFormat(threadName).build());
   }
 
   /**
-   * Creates an instance of {@link StreamingEngineClient} in a non-started state.
+   * Creates an instance of {@link FanOutStreamingEngineWorkerHarness} in a non-started state.
    *
    * @implNote Does not block the calling thread. Callers must explicitly call {@link #start()}.
    */
-  public static StreamingEngineClient create(
+  public static FanOutStreamingEngineWorkerHarness create(
       JobHeader jobHeader,
       GetWorkBudget totalGetWorkBudget,
       GrpcWindmillStreamFactory streamingEngineStreamFactory,
@@ -178,7 +175,7 @@ public static StreamingEngineClient create(
       GrpcDispatcherClient dispatcherClient,
       Function<WindmillStream.CommitWorkStream, WorkCommitter> workCommitterFactory,
       ThrottlingGetDataMetricTracker getDataMetricTracker) {
-    return new StreamingEngineClient(
+    return new FanOutStreamingEngineWorkerHarness(
         jobHeader,
         totalGetWorkBudget,
         streamingEngineStreamFactory,
@@ -192,7 +189,7 @@ public static StreamingEngineClient create(
   }
 
   @VisibleForTesting
-  static StreamingEngineClient forTesting(
+  static FanOutStreamingEngineWorkerHarness forTesting(
       JobHeader jobHeader,
       GetWorkBudget totalGetWorkBudget,
       GrpcWindmillStreamFactory streamFactory,
@@ -203,8 +200,8 @@ static StreamingEngineClient forTesting(
       long clientId,
       Function<WindmillStream.CommitWorkStream, WorkCommitter> workCommitterFactory,
       ThrottlingGetDataMetricTracker getDataMetricTracker) {
-    StreamingEngineClient streamingEngineClient =
-        new StreamingEngineClient(
+    FanOutStreamingEngineWorkerHarness fanOutStreamingEngineWorkProvider =
+        new FanOutStreamingEngineWorkerHarness(
             jobHeader,
             totalGetWorkBudget,
             streamFactory,
@@ -215,11 +212,12 @@ static StreamingEngineClient forTesting(
             clientId,
             workCommitterFactory,
             getDataMetricTracker);
-    streamingEngineClient.start();
-    return streamingEngineClient;
+    fanOutStreamingEngineWorkProvider.start();
+    return fanOutStreamingEngineWorkProvider;
   }
 
   @SuppressWarnings("ReturnValueIgnored")
+  @Override
   public synchronized void start() {
     Preconditions.checkState(!started, "StreamingEngineClient cannot start twice.");
     // Starts the stream, this value is memoized.
@@ -270,7 +268,8 @@ private void startWorkerMetadataConsumer() {
   }
 
   @VisibleForTesting
-  public synchronized void finish() {
+  @Override
+  public synchronized void shutdown() {
     Preconditions.checkState(started, "StreamingEngineClient never started.");
     getWorkerMetadataStream.get().halfClose();
     getWorkBudgetRefresher.stop();
@@ -334,10 +333,13 @@ private synchronized ImmutableMap<Endpoint, WindmillConnection> createNewWindmil
         .collect(
             toImmutableMap(
                 Function.identity(),
-                // Reuse existing stubs if they exist.
                 endpoint ->
-                    currentConnections.getOrDefault(
-                        endpoint, WindmillConnection.from(endpoint, this::createWindmillStub))));
+                    // Reuse existing stubs if they exist. Optional.orElseGet only calls the
+                    // supplier if the value is not present, preventing constructing expensive
+                    // objects.
+                    Optional.ofNullable(currentConnections.get(endpoint))
+                        .orElseGet(
+                            () -> WindmillConnection.from(endpoint, this::createWindmillStub))));
   }
 
   private synchronized ImmutableMap<WindmillConnection, WindmillStreamSender>
@@ -423,11 +425,4 @@ private CloudWindmillServiceV1Alpha1Stub createWindmillStub(Endpoint endpoint) {
         .map(channelCachingStubFactory::createWindmillServiceStub)
         .orElseGet(dispatcherClient::getWindmillServiceStub);
   }
-
-  private static class StreamingEngineClientException extends IllegalStateException {
-
-    private StreamingEngineClientException(Throwable exception) {
-      super(exception);
-    }
-  }
 }