consume flow control settings updates from job settings

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

@@ -17,7 +17,7 @@
  */
 package org.apache.beam.runners.dataflow.worker;
 
-import static org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs.WindmillChannelFactory.remoteChannel;
+import static org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs.WindmillChannels.remoteChannel;
 
 import com.google.api.services.dataflow.model.MapTask;
 import com.google.auto.value.AutoValue;
@@ -87,7 +87,6 @@
 import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.GrpcWindmillStreamFactory;
 import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs.ChannelCache;
 import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs.ChannelCachingRemoteStubFactory;
-import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs.ChannelCachingStubFactory;
 import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs.IsolationChannel;
 import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs.WindmillStubFactoryFactory;
 import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs.WindmillStubFactoryFactoryImpl;
@@ -243,10 +242,11 @@ private StreamingDataflowWorker(
     @Nullable ChannelzServlet channelzServlet = null;
     Consumer<PrintWriter> getDataStatusProvider;
     Supplier<Long> currentActiveCommitBytesProvider;
-
+    ChannelCache channelCache = null;
     if (options.isEnableStreamingEngine() && options.getIsWindmillServiceDirectPathEnabled()) {
       // Direct path pipelines.
       WeightedSemaphore<Commit> maxCommitByteSemaphore = Commits.maxCommitByteSemaphore();
+      channelCache = createChannelCache(options, configFetcher);
       FanOutStreamingEngineWorkerHarness fanOutStreamingEngineWorkerHarness =
           FanOutStreamingEngineWorkerHarness.create(
               createJobHeader(options, clientId),
@@ -273,7 +273,7 @@ private StreamingDataflowWorker(
                                 processingContext,
                                 getWorkStreamLatencies);
                           }),
-              createFanOutStubFactory(options),
+              ChannelCachingRemoteStubFactory.create(options.getGcpCredential(), channelCache),
               GetWorkBudgetDistributors.distributeEvenly(),
               Preconditions.checkNotNull(dispatcherClient),
               commitWorkStream ->
@@ -384,6 +384,7 @@ private StreamingDataflowWorker(
             .setChannelzServlet(channelzServlet)
             .setGetDataStatusProvider(getDataStatusProvider)
             .setCurrentActiveCommitBytes(currentActiveCommitBytesProvider)
+            .setChannelCache(channelCache)
             .build();
 
     LOG.debug("isDirectPathEnabled: {}", options.getIsWindmillServiceDirectPathEnabled());
@@ -620,19 +621,28 @@ private static void validateWorkerOptions(DataflowWorkerHarnessOptions options)
         StreamingDataflowWorker.class.getSimpleName());
   }
 
-  private static ChannelCachingStubFactory createFanOutStubFactory(
-      DataflowWorkerHarnessOptions workerOptions) {
-    return ChannelCachingRemoteStubFactory.create(
-        workerOptions.getGcpCredential(),
+  private static ChannelCache createChannelCache(
+      DataflowWorkerHarnessOptions workerOptions, ComputationConfig.Fetcher configFetcher) {
+    ChannelCache channelCache =
         ChannelCache.create(
-            serviceAddress ->
-                // IsolationChannel will create and manage separate RPC channels to the same
-                // serviceAddress.
-                IsolationChannel.create(
-                    () ->
-                        remoteChannel(
-                            serviceAddress,
-                            workerOptions.getWindmillServiceRpcChannelAliveTimeoutSec()))));
+            (currentFlowControlSettings, serviceAddress) -> {
+              // IsolationChannel will create and manage separate RPC channels to the same
+              // serviceAddress.
+              return IsolationChannel.create(
+                  () ->
+                      remoteChannel(
+                          serviceAddress,
+                          workerOptions.getWindmillServiceRpcChannelAliveTimeoutSec(),
+                          currentFlowControlSettings),
+                  currentFlowControlSettings.getOnReadyThresholdBytes());
+            });
+    configFetcher
+        .getGlobalConfigHandle()
+        .registerConfigObserver(
+            config ->
+                channelCache.consumeFlowControlSettings(
+                    config.userWorkerJobSettings().getFlowControlSettings()));
+    return channelCache;
   }
 
   @VisibleForTesting

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

@@ -18,7 +18,6 @@
 package org.apache.beam.runners.dataflow.worker.streaming.config;
 
 import java.util.function.Consumer;
-import javax.annotation.Nonnull;
 import javax.annotation.concurrent.ThreadSafe;
 import org.apache.beam.sdk.annotations.Internal;
 
@@ -33,5 +32,5 @@ public interface StreamingGlobalConfigHandle {
    * Subscribe to config updates by registering a callback. Callback should be called the first time
    * with settings, if any. The callback could execute inline before the method returns.
    */
-  void registerConfigObserver(@Nonnull Consumer<StreamingGlobalConfig> callback);
+  void registerConfigObserver(Consumer<StreamingGlobalConfig> callback);
 }

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

@@ -44,6 +44,7 @@
 import org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor;
 import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.ChannelzServlet;
 import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.GrpcWindmillStreamFactory;
+import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs.ChannelCache;
 import org.apache.beam.runners.dataflow.worker.windmill.state.WindmillStateCache;
 import org.apache.beam.sdk.annotations.Internal;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
@@ -79,6 +80,7 @@ public final class StreamingWorkerStatusPages {
   private final @Nullable GrpcWindmillStreamFactory windmillStreamFactory;
   private final DebugCapture.@Nullable Manager debugCapture;
   private final @Nullable ChannelzServlet channelzServlet;
+  private final @Nullable ChannelCache channelCache;
 
   private final AtomicReference<StreamingGlobalConfig> globalConfig = new AtomicReference<>();
 
@@ -96,7 +98,8 @@ public final class StreamingWorkerStatusPages {
       Consumer<PrintWriter> getDataStatusProvider,
       BoundedQueueExecutor workUnitExecutor,
       ScheduledExecutorService statusPageDumper,
-      StreamingGlobalConfigHandle globalConfigHandle) {
+      StreamingGlobalConfigHandle globalConfigHandle,
+      @Nullable ChannelCache channelCache) {
     this.clock = clock;
     this.clientId = clientId;
     this.isRunning = isRunning;
@@ -111,6 +114,7 @@ public final class StreamingWorkerStatusPages {
     this.workUnitExecutor = workUnitExecutor;
     this.statusPageDumper = statusPageDumper;
     globalConfigHandle.registerConfigObserver(globalConfig::set);
+    this.channelCache = channelCache;
   }
 
   public static StreamingWorkerStatusPages.Builder builder() {
@@ -168,6 +172,9 @@ private void addStreamingEngineStatusPages() {
           }
           writer.println(config.userWorkerJobSettings().toString());
         });
+    if (channelCache != null) {
+      statusPages.addStatusDataProvider("ChannelCache", "ChannelCache", channelCache);
+    }
   }
 
   private boolean isStreamingEngine() {
@@ -276,6 +283,8 @@ public interface Builder {
 
     Builder setGlobalConfigHandle(StreamingGlobalConfigHandle globalConfigHandle);
 
+    Builder setChannelCache(@Nullable ChannelCache channelCache);
+
     StreamingWorkerStatusPages build();
   }
 }

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

@@ -20,6 +20,7 @@
 import com.google.auto.value.AutoValue;
 import java.util.Optional;
 import java.util.function.Function;
+import java.util.function.Supplier;
 import org.apache.beam.runners.dataflow.worker.windmill.CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alpha1Stub;
 import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints.Endpoint;
 import org.apache.beam.sdk.annotations.Internal;
@@ -36,7 +37,7 @@ public static WindmillConnection from(
 
     windmillEndpoint.workerToken().ifPresent(windmillWorkerConnection::setBackendWorkerToken);
     windmillEndpoint.directEndpoint().ifPresent(windmillWorkerConnection::setDirectEndpoint);
-    windmillWorkerConnection.setStub(endpointToStubFn.apply(windmillEndpoint));
+    windmillWorkerConnection.setStubSupplier(() -> endpointToStubFn.apply(windmillEndpoint));
 
     return windmillWorkerConnection.build();
   }
@@ -48,17 +49,22 @@ public static Builder builder() {
 
   public abstract String backendWorkerToken();
 
-  public abstract Optional<WindmillServiceAddress> directEndpoint();
+  abstract Optional<WindmillServiceAddress> directEndpoint();
 
-  public abstract CloudWindmillServiceV1Alpha1Stub stub();
+  abstract Supplier<CloudWindmillServiceV1Alpha1Stub> stubSupplier();
+
+  public final CloudWindmillServiceV1Alpha1Stub currentStub() {
+    return stubSupplier().get();
+  }
 
   @AutoValue.Builder
   public abstract static class Builder {
     abstract Builder setBackendWorkerToken(String backendWorkerToken);
 
-    public abstract Builder setDirectEndpoint(WindmillServiceAddress value);
+    abstract Builder setDirectEndpoint(WindmillServiceAddress value);
 
-    public abstract Builder setStub(CloudWindmillServiceV1Alpha1Stub stub);
+    public abstract Builder setStubSupplier(
+        Supplier<CloudWindmillServiceV1Alpha1Stub> stubSupplier);
 
     public abstract WindmillConnection build();
   }

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

@@ -17,8 +17,8 @@
  */
 package org.apache.beam.runners.dataflow.worker.windmill.client.grpc;
 
-import static org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs.WindmillChannelFactory.LOCALHOST;
-import static org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs.WindmillChannelFactory.localhostChannel;
+import static org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs.WindmillChannels.LOCALHOST;
+import static org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs.WindmillChannels.localhostChannel;
 
 import com.google.auto.value.AutoValue;
 import java.util.List;
@@ -164,15 +164,6 @@ private synchronized <T> T randomlySelectNextStub(List<T> stubs) {
     return stubs.get(rand.nextInt(stubs.size()));
   }
 
-  /**
-   * Returns whether the {@link DispatcherStubs} have been set. Once initially set, {@link
-   * #dispatcherStubs} will always have a value as empty updates will trigger an {@link
-   * IllegalStateException}.
-   */
-  public boolean hasInitializedEndpoints() {
-    return dispatcherStubs.get().hasInitializedEndpoints();
-  }
-
   public void onJobConfig(StreamingGlobalConfig config) {
     if (config.windmillServiceEndpoints().isEmpty()) {
       LOG.warn("Dispatcher client received empty windmill service endpoints from global config");
@@ -282,14 +273,6 @@ private static CloudWindmillMetadataServiceV1Alpha1Stub createWindmillMetadataSe
           WindmillServiceAddress.create(endpoint));
     }
 
-    private int size() {
-      return dispatcherEndpoints().size();
-    }
-
-    private boolean hasInitializedEndpoints() {
-      return size() > 0;
-    }
-
     abstract ImmutableSet<HostAndPort> dispatcherEndpoints();
 
     abstract ImmutableList<CloudWindmillServiceV1Alpha1Stub> windmillServiceStubs();

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

@@ -17,9 +17,9 @@
  */
 package org.apache.beam.runners.dataflow.worker.windmill.client.grpc;
 
-import static org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs.WindmillChannelFactory.LOCALHOST;
-import static org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs.WindmillChannelFactory.inProcessChannel;
-import static org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs.WindmillChannelFactory.localhostChannel;
+import static org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs.WindmillChannels.LOCALHOST;
+import static org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs.WindmillChannels.inProcessChannel;
+import static org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs.WindmillChannels.localhostChannel;
 
 import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
 import java.io.IOException;

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

@@ -230,7 +230,7 @@ public GetWorkStream createDirectGetWorkStream(
       WorkItemScheduler workItemScheduler) {
     return GrpcDirectGetWorkStream.create(
         connection.backendWorkerToken(),
-        responseObserver -> connection.stub().getWorkStream(responseObserver),
+        responseObserver -> connection.currentStub().getWorkStream(responseObserver),
         request,
         grpcBackOff.get(),
         newStreamObserverFactory(),
@@ -265,7 +265,7 @@ public GetDataStream createDirectGetDataStream(
       WindmillConnection connection, ThrottleTimer getDataThrottleTimer) {
     return GrpcGetDataStream.create(
         connection.backendWorkerToken(),
-        responseObserver -> connection.stub().getDataStream(responseObserver),
+        responseObserver -> connection.currentStub().getDataStream(responseObserver),
         grpcBackOff.get(),
         newStreamObserverFactory(),
         streamRegistry,
@@ -297,7 +297,7 @@ public CommitWorkStream createDirectCommitWorkStream(
       WindmillConnection connection, ThrottleTimer commitWorkThrottleTimer) {
     return GrpcCommitWorkStream.create(
         connection.backendWorkerToken(),
-        responseObserver -> connection.stub().commitWorkStream(responseObserver),
+        responseObserver -> connection.currentStub().commitWorkStream(responseObserver),
         grpcBackOff.get(),
         newStreamObserverFactory(),
         streamRegistry,