You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by da...@apache.org on 2025/04/16 14:00:18 UTC
(beam) branch master updated: consume flow control settings updates from job settings (#34539)
This is an automated email from the ASF dual-hosted git repository.
damccorm pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 9331662c934 consume flow control settings updates from job settings (#34539)
9331662c934 is described below
commit 9331662c93409246d88be1ef5b79b696db49e0e0
Author: martin trieu <ma...@google.com>
AuthorDate: Wed Apr 16 07:00:11 2025 -0700
consume flow control settings updates from job settings (#34539)
* consume flow control settings updates from job settings, lazily create streams/channels to use new flow control settings
* address PR comments
* address PR comments
* address PR comments
* address PR comments
---
.../dataflow/worker/StreamingDataflowWorker.java | 42 +++++---
.../config/StreamingGlobalConfigHandle.java | 3 +-
.../harness/StreamingWorkerStatusPages.java | 11 +-
.../worker/windmill/WindmillConnection.java | 16 ++-
.../windmill/client/grpc/GrpcDispatcherClient.java | 21 +---
.../windmill/client/grpc/GrpcWindmillServer.java | 6 +-
.../client/grpc/GrpcWindmillStreamFactory.java | 6 +-
.../windmill/client/grpc/stubs/ChannelCache.java | 66 ++++++++++--
.../client/grpc/stubs/IsolationChannel.java | 34 ++++--
.../client/grpc/stubs/WindmillChannelFactory.java | 119 ++-------------------
...llChannelFactory.java => WindmillChannels.java} | 100 +++++++++++++----
.../grpc/stubs/WindmillStubFactoryFactoryImpl.java | 22 ++--
.../worker/StreamingDataflowWorkerTest.java | 5 +-
.../FanOutStreamingEngineWorkerHarnessTest.java | 6 +-
.../harness/WindmillStreamSenderTest.java | 2 +-
.../client/grpc/GrpcDirectGetWorkStreamTest.java | 3 +-
.../client/grpc/GrpcWindmillServerTest.java | 4 +-
.../client/grpc/stubs/ChannelCacheTest.java | 103 +++++++++++++++++-
.../windmill/testing/FakeWindmillStubFactory.java | 4 +-
19 files changed, 344 insertions(+), 229 deletions(-)
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java
index 30782735d77..daca54fe158 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java
+++ b/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.GrpcWindmill
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 @@ public final class 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 @@ public final class StreamingDataflowWorker {
processingContext,
getWorkStreamLatencies);
}),
- createFanOutStubFactory(options),
+ ChannelCachingRemoteStubFactory.create(options.getGcpCredential(), channelCache),
GetWorkBudgetDistributors.distributeEvenly(),
Preconditions.checkNotNull(dispatcherClient),
commitWorkStream ->
@@ -384,6 +384,7 @@ public final class StreamingDataflowWorker {
.setChannelzServlet(channelzServlet)
.setGetDataStatusProvider(getDataStatusProvider)
.setCurrentActiveCommitBytes(currentActiveCommitBytesProvider)
+ .setChannelCache(channelCache)
.build();
LOG.debug("isDirectPathEnabled: {}", options.getIsWindmillServiceDirectPathEnabled());
@@ -620,19 +621,28 @@ public final class StreamingDataflowWorker {
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
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/config/StreamingGlobalConfigHandle.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/config/StreamingGlobalConfigHandle.java
index 6f75ba88747..78810b35c88 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/config/StreamingGlobalConfigHandle.java
+++ b/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);
}
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/harness/StreamingWorkerStatusPages.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/harness/StreamingWorkerStatusPages.java
index ddfc6809231..4be367e6cb7 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/harness/StreamingWorkerStatusPages.java
+++ b/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.streaming.config.StreamingGlobalC
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 @@ public final class StreamingWorkerStatusPages {
}
writer.println(config.userWorkerJobSettings().toString());
});
+ if (channelCache != null) {
+ statusPages.addStatusDataProvider("ChannelCache", "ChannelCache", channelCache);
+ }
}
private boolean isStreamingEngine() {
@@ -276,6 +283,8 @@ public final class StreamingWorkerStatusPages {
Builder setGlobalConfigHandle(StreamingGlobalConfigHandle globalConfigHandle);
+ Builder setChannelCache(@Nullable ChannelCache channelCache);
+
StreamingWorkerStatusPages build();
}
}
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/WindmillConnection.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/WindmillConnection.java
index 7d199afc086..4ffd2861b48 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/WindmillConnection.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/WindmillConnection.java
@@ -20,6 +20,7 @@ package org.apache.beam.runners.dataflow.worker.windmill;
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 abstract class WindmillConnection {
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 abstract class WindmillConnection {
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();
}
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcDispatcherClient.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcDispatcherClient.java
index 23488883177..3603bacf461 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcDispatcherClient.java
+++ b/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 @@ public class GrpcDispatcherClient {
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 @@ public class GrpcDispatcherClient {
WindmillServiceAddress.create(endpoint));
}
- private int size() {
- return dispatcherEndpoints().size();
- }
-
- private boolean hasInitializedEndpoints() {
- return size() > 0;
- }
-
abstract ImmutableSet<HostAndPort> dispatcherEndpoints();
abstract ImmutableList<CloudWindmillServiceV1Alpha1Stub> windmillServiceStubs();
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcWindmillServer.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcWindmillServer.java
index 8fe1f14caa7..e4a7e1a2f92 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcWindmillServer.java
+++ b/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;
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcWindmillStreamFactory.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcWindmillStreamFactory.java
index 5d35fc5aedf..fd97c6b48e1 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcWindmillStreamFactory.java
+++ b/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 class GrpcWindmillStreamFactory implements StatusDataProvider {
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 class GrpcWindmillStreamFactory implements StatusDataProvider {
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 class GrpcWindmillStreamFactory implements StatusDataProvider {
WindmillConnection connection, ThrottleTimer commitWorkThrottleTimer) {
return GrpcCommitWorkStream.create(
connection.backendWorkerToken(),
- responseObserver -> connection.stub().commitWorkStream(responseObserver),
+ responseObserver -> connection.currentStub().commitWorkStream(responseObserver),
grpcBackOff.get(),
newStreamObserverFactory(),
streamRegistry,
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/stubs/ChannelCache.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/stubs/ChannelCache.java
index 9a401690ca2..d34612725a1 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/stubs/ChannelCache.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/stubs/ChannelCache.java
@@ -17,13 +17,16 @@
*/
package org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs;
+import static org.apache.beam.runners.dataflow.worker.windmill.WindmillServiceAddress.Kind.AUTHENTICATED_GCP_SERVICE_ADDRESS;
+
import java.io.PrintWriter;
import java.util.concurrent.Executor;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
-import java.util.function.Function;
+import javax.annotation.concurrent.GuardedBy;
import javax.annotation.concurrent.ThreadSafe;
import org.apache.beam.runners.dataflow.worker.status.StatusDataProvider;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.UserWorkerGrpcFlowControlSettings;
import org.apache.beam.runners.dataflow.worker.windmill.WindmillServiceAddress;
import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.ManagedChannel;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
@@ -34,6 +37,7 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.RemovalL
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.RemovalListeners;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.MoreExecutors;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.ThreadFactoryBuilder;
+import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -48,10 +52,16 @@ import org.slf4j.LoggerFactory;
@ThreadSafe
public final class ChannelCache implements StatusDataProvider {
private static final Logger LOG = LoggerFactory.getLogger(ChannelCache.class);
+ private static final String DIRECTPATH = "Directpath";
+ private static final String CLOUDPATH = "Cloudpath";
private final LoadingCache<WindmillServiceAddress, ManagedChannel> channelCache;
+ @GuardedBy("this")
+ @MonotonicNonNull
+ private UserWorkerGrpcFlowControlSettings currentFlowControlSettings = null;
+
private ChannelCache(
- Function<WindmillServiceAddress, ManagedChannel> channelFactory,
+ WindmillChannelFactory channelFactory,
RemovalListener<WindmillServiceAddress, ManagedChannel> onChannelRemoved,
Executor channelCloser) {
this.channelCache =
@@ -61,13 +71,25 @@ public final class ChannelCache implements StatusDataProvider {
new CacheLoader<WindmillServiceAddress, ManagedChannel>() {
@Override
public ManagedChannel load(WindmillServiceAddress key) {
- return channelFactory.apply(key);
+ return channelFactory.create(resolveFlowControlSettings(key.getKind()), key);
+ }
+
+ private UserWorkerGrpcFlowControlSettings resolveFlowControlSettings(
+ WindmillServiceAddress.Kind addressType) {
+ synchronized (ChannelCache.this) {
+ if (currentFlowControlSettings == null) {
+ return addressType == AUTHENTICATED_GCP_SERVICE_ADDRESS
+ ? WindmillChannels.DEFAULT_DIRECTPATH_FLOW_CONTROL_SETTINGS
+ : WindmillChannels.DEFAULT_CLOUDPATH_FLOW_CONTROL_SETTINGS;
+ }
+
+ return currentFlowControlSettings;
+ }
}
});
}
- public static ChannelCache create(
- Function<WindmillServiceAddress, ManagedChannel> channelFactory) {
+ public static ChannelCache create(WindmillChannelFactory channelFactory) {
return new ChannelCache(
channelFactory,
// Shutdown the channels as they get removed from the cache, so they do not leak.
@@ -78,7 +100,7 @@ public final class ChannelCache implements StatusDataProvider {
@VisibleForTesting
public static ChannelCache forTesting(
- Function<WindmillServiceAddress, ManagedChannel> channelFactory, Runnable onChannelShutdown) {
+ WindmillChannelFactory channelFactory, Runnable onChannelShutdown) {
return new ChannelCache(
channelFactory,
// Shutdown the channels as they get removed from the cache, so they do not leak.
@@ -98,7 +120,7 @@ public final class ChannelCache implements StatusDataProvider {
try {
channel.awaitTermination(10, TimeUnit.SECONDS);
} catch (InterruptedException e) {
- LOG.error("Couldn't close gRPC channel={}", channel, e);
+ LOG.error("Couldn't gracefully close gRPC channel={}", channel, e);
}
channel.shutdownNow();
}
@@ -107,6 +129,18 @@ public final class ChannelCache implements StatusDataProvider {
return channelCache.getUnchecked(windmillServiceAddress);
}
+ public synchronized void consumeFlowControlSettings(
+ UserWorkerGrpcFlowControlSettings flowControlSettings) {
+ if (!flowControlSettings.equals(currentFlowControlSettings)) {
+ // Refreshing the cache will asynchronously terminate the old channels via the removalListener
+ // and return a newly created one on the next Cache.load(address). This could be expensive so
+ // only do it when we have received new flow control settings.
+ LOG.debug("Updating flow control settings {}.", flowControlSettings);
+ currentFlowControlSettings = flowControlSettings;
+ channelCache.asMap().keySet().forEach(channelCache::refresh);
+ }
+ }
+
public void remove(WindmillServiceAddress windmillServiceAddress) {
channelCache.invalidate(windmillServiceAddress);
}
@@ -128,12 +162,28 @@ public final class ChannelCache implements StatusDataProvider {
@Override
public void appendSummaryHtml(PrintWriter writer) {
+ synchronized (this) {
+ if (currentFlowControlSettings == null) {
+ writer.format(
+ "Directpath gRPC flow control settings: [%s]",
+ WindmillChannels.DEFAULT_DIRECTPATH_FLOW_CONTROL_SETTINGS);
+ writer.format(
+ "Cloudpath gRPC flow control settings: [%s]",
+ WindmillChannels.DEFAULT_CLOUDPATH_FLOW_CONTROL_SETTINGS);
+ } else {
+ writer.format("gRPC flow control settings: [%s]", currentFlowControlSettings);
+ }
+ }
writer.write("Active gRPC Channels:<br>");
channelCache
.asMap()
.forEach(
(address, channel) -> {
- writer.format("Address: [%s]; Channel: [%s].", address, channel);
+ writer.format(
+ "Address: [%s]; Channel: [%s]; AddressType:[%s].",
+ address,
+ channel,
+ address.getKind() == AUTHENTICATED_GCP_SERVICE_ADDRESS ? DIRECTPATH : CLOUDPATH);
writer.write("<br>");
});
}
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/stubs/IsolationChannel.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/stubs/IsolationChannel.java
index d78dd421dcc..e0cb032a5f8 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/stubs/IsolationChannel.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/stubs/IsolationChannel.java
@@ -34,7 +34,6 @@ import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.ManagedChannel;
import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.Metadata;
import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.MethodDescriptor;
import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.Status;
-import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -48,6 +47,7 @@ public class IsolationChannel extends ManagedChannel {
private static final Logger LOG = LoggerFactory.getLogger(IsolationChannel.class);
private final Supplier<ManagedChannel> channelFactory;
+ private final int onReadyThresholdBytes;
@GuardedBy("this")
private final List<ManagedChannel> channelCache;
@@ -55,24 +55,29 @@ public class IsolationChannel extends ManagedChannel {
@GuardedBy("this")
private final Set<ManagedChannel> usedChannels;
- @GuardedBy("this")
- private boolean shutdownStarted = false;
-
private final String authority;
- // Expected that supplier returns channels to the same endpoint with the same authority.
- public static IsolationChannel create(Supplier<ManagedChannel> channelFactory) {
- return new IsolationChannel(channelFactory);
- }
+ @GuardedBy("this")
+ private boolean shutdownStarted = false;
- @VisibleForTesting
- IsolationChannel(Supplier<ManagedChannel> channelFactory) {
+ private IsolationChannel(Supplier<ManagedChannel> channelFactory, int onReadyThresholdBytes) {
this.channelFactory = channelFactory;
this.channelCache = new ArrayList<>();
ManagedChannel channel = channelFactory.get();
this.authority = channel.authority();
this.channelCache.add(channel);
this.usedChannels = new HashSet<>();
+ this.onReadyThresholdBytes = onReadyThresholdBytes;
+ }
+
+ // Expected that supplier returns channels to the same endpoint with the same authority.
+ public static IsolationChannel create(Supplier<ManagedChannel> channelFactory) {
+ return new IsolationChannel(channelFactory, -1);
+ }
+
+ public static IsolationChannel create(
+ Supplier<ManagedChannel> channelFactory, int onReadyThresholdBytes) {
+ return new IsolationChannel(channelFactory, onReadyThresholdBytes);
}
@Override
@@ -85,7 +90,14 @@ public class IsolationChannel extends ManagedChannel {
public <ReqT, RespT> ClientCall<ReqT, RespT> newCall(
MethodDescriptor<ReqT, RespT> methodDescriptor, CallOptions callOptions) {
ManagedChannel channel = getChannel();
- return new ReleasingClientCall<>(this, channel, channel.newCall(methodDescriptor, callOptions));
+ return new ReleasingClientCall<>(
+ this,
+ channel,
+ channel.newCall(
+ methodDescriptor,
+ onReadyThresholdBytes > 0
+ ? callOptions.withOnReadyThreshold(onReadyThresholdBytes)
+ : callOptions));
}
private synchronized ManagedChannel getChannel() {
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/stubs/WindmillChannelFactory.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/stubs/WindmillChannelFactory.java
index 8ca962e9560..3a64528ca98 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/stubs/WindmillChannelFactory.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/stubs/WindmillChannelFactory.java
@@ -17,120 +17,13 @@
*/
package org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs;
-import java.util.concurrent.TimeUnit;
-import javax.net.ssl.SSLException;
-import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.UserWorkerGrpcFlowControlSettings;
import org.apache.beam.runners.dataflow.worker.windmill.WindmillServiceAddress;
-import org.apache.beam.runners.dataflow.worker.windmill.WindmillServiceAddress.AuthenticatedGcpServiceAddress;
-import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.Channel;
import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.ManagedChannel;
-import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.alts.AltsChannelCredentials;
-import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.inprocess.InProcessChannelBuilder;
-import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.netty.GrpcSslContexts;
-import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.netty.NegotiationType;
-import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.netty.NettyChannelBuilder;
-import org.apache.beam.vendor.grpc.v1p69p0.io.netty.handler.ssl.SslContext;
-import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.net.HostAndPort;
-/** Utility class used to create different RPC Channels. */
-public final class WindmillChannelFactory {
- public static final String LOCALHOST = "localhost";
- private static final int MAX_REMOTE_TRACE_EVENTS = 100;
- // 10MiB.
- private static final int WINDMILL_MAX_FLOW_CONTROL_WINDOW =
- NettyChannelBuilder.DEFAULT_FLOW_CONTROL_WINDOW * 10;
-
- private WindmillChannelFactory() {}
-
- public static ManagedChannel inProcessChannel(String channelName) {
- return InProcessChannelBuilder.forName(channelName).directExecutor().build();
- }
-
- public static Channel localhostChannel(int port) {
- return NettyChannelBuilder.forAddress(LOCALHOST, port)
- .maxInboundMessageSize(Integer.MAX_VALUE)
- .negotiationType(NegotiationType.PLAINTEXT)
- .build();
- }
-
- public static ManagedChannel remoteChannel(
- WindmillServiceAddress windmillServiceAddress, int windmillServiceRpcChannelTimeoutSec) {
- switch (windmillServiceAddress.getKind()) {
- case GCP_SERVICE_ADDRESS:
- return remoteChannel(
- windmillServiceAddress.gcpServiceAddress(), windmillServiceRpcChannelTimeoutSec);
- case AUTHENTICATED_GCP_SERVICE_ADDRESS:
- return remoteDirectChannel(
- windmillServiceAddress.authenticatedGcpServiceAddress(),
- windmillServiceRpcChannelTimeoutSec);
- default:
- throw new UnsupportedOperationException(
- "Only GCP_SERVICE_ADDRESS and AUTHENTICATED_GCP_SERVICE_ADDRESS are supported"
- + " WindmillServiceAddresses.");
- }
- }
-
- private static ManagedChannel remoteDirectChannel(
- AuthenticatedGcpServiceAddress authenticatedGcpServiceAddress,
- int windmillServiceRpcChannelTimeoutSec) {
- return withDefaultChannelOptions(
- NettyChannelBuilder.forAddress(
- authenticatedGcpServiceAddress.gcpServiceAddress().getHost(),
- // Ports are required for direct channels.
- authenticatedGcpServiceAddress.gcpServiceAddress().getPort(),
- new AltsChannelCredentials.Builder().build())
- .overrideAuthority(authenticatedGcpServiceAddress.authenticatingService()),
- windmillServiceRpcChannelTimeoutSec)
- .build();
- }
-
- public static ManagedChannel remoteChannel(
- HostAndPort endpoint, int windmillServiceRpcChannelTimeoutSec) {
- return withDefaultChannelOptions(
- NettyChannelBuilder.forAddress(
- endpoint.getHost(),
- endpoint.hasPort()
- ? endpoint.getPort()
- : WindmillEndpoints.DEFAULT_WINDMILL_SERVICE_PORT),
- windmillServiceRpcChannelTimeoutSec)
- .negotiationType(NegotiationType.TLS)
- .sslContext(dataflowGrpcSslContext(endpoint))
- .build();
- }
-
- @SuppressWarnings("nullness")
- private static SslContext dataflowGrpcSslContext(HostAndPort endpoint) {
- try {
- // Set ciphers(null) to not use GCM, which is disabled for Dataflow
- // due to it being horribly slow.
- return GrpcSslContexts.forClient().ciphers(null).build();
- } catch (SSLException sslException) {
- throw new WindmillChannelCreationException(endpoint, sslException);
- }
- }
-
- private static NettyChannelBuilder withDefaultChannelOptions(
- NettyChannelBuilder channelBuilder, int windmillServiceRpcChannelTimeoutSec) {
- if (windmillServiceRpcChannelTimeoutSec > 0) {
- channelBuilder
- .keepAliveTime(windmillServiceRpcChannelTimeoutSec, TimeUnit.SECONDS)
- .keepAliveTimeout(windmillServiceRpcChannelTimeoutSec, TimeUnit.SECONDS)
- .keepAliveWithoutCalls(true);
- }
-
- return channelBuilder
- .maxInboundMessageSize(Integer.MAX_VALUE)
- .maxTraceEvents(MAX_REMOTE_TRACE_EVENTS)
- // 1MiB
- .maxInboundMetadataSize(1024 * 1024)
- .flowControlWindow(WINDMILL_MAX_FLOW_CONTROL_WINDOW);
- }
-
- private static class WindmillChannelCreationException extends IllegalStateException {
- private WindmillChannelCreationException(HostAndPort endpoint, SSLException sourceException) {
- super(
- String.format("Exception thrown when trying to create channel to %s", endpoint),
- sourceException);
- }
- }
+@FunctionalInterface
+public interface WindmillChannelFactory {
+ ManagedChannel create(
+ UserWorkerGrpcFlowControlSettings flowControlSettings,
+ WindmillServiceAddress windmillServiceAddress);
}
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/stubs/WindmillChannelFactory.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/stubs/WindmillChannels.java
similarity index 55%
copy from runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/stubs/WindmillChannelFactory.java
copy to runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/stubs/WindmillChannels.java
index 8ca962e9560..375cd354ac9 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/stubs/WindmillChannelFactory.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/stubs/WindmillChannels.java
@@ -19,6 +19,7 @@ package org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs;
import java.util.concurrent.TimeUnit;
import javax.net.ssl.SSLException;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.UserWorkerGrpcFlowControlSettings;
import org.apache.beam.runners.dataflow.worker.windmill.WindmillEndpoints;
import org.apache.beam.runners.dataflow.worker.windmill.WindmillServiceAddress;
import org.apache.beam.runners.dataflow.worker.windmill.WindmillServiceAddress.AuthenticatedGcpServiceAddress;
@@ -33,14 +34,40 @@ import org.apache.beam.vendor.grpc.v1p69p0.io.netty.handler.ssl.SslContext;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.net.HostAndPort;
/** Utility class used to create different RPC Channels. */
-public final class WindmillChannelFactory {
+public final class WindmillChannels {
public static final String LOCALHOST = "localhost";
private static final int MAX_REMOTE_TRACE_EVENTS = 100;
- // 10MiB.
- private static final int WINDMILL_MAX_FLOW_CONTROL_WINDOW =
+ // 1MiB.
+ private static final int MAX_INBOUND_METADATA_SIZE_BYTES = 1024 * 1024;
+ // 10MiB. Roughly 2x max message size.
+ private static final int WINDMILL_MIN_FLOW_CONTROL_WINDOW =
NettyChannelBuilder.DEFAULT_FLOW_CONTROL_WINDOW * 10;
+ public static final UserWorkerGrpcFlowControlSettings DEFAULT_CLOUDPATH_FLOW_CONTROL_SETTINGS =
+ UserWorkerGrpcFlowControlSettings.newBuilder()
+ .setFlowControlWindowBytes(WINDMILL_MIN_FLOW_CONTROL_WINDOW)
+ .setEnableAutoFlowControl(false)
+ .build();
+ // 100MiB.
+ private static final int DIRECTPATH_MIN_FLOW_CONTROL_WINDOW =
+ WINDMILL_MIN_FLOW_CONTROL_WINDOW * 10;
+ // 100MiB.
+ private static final int DIRECTPATH_ON_READY_THRESHOLD = WINDMILL_MIN_FLOW_CONTROL_WINDOW * 10;
+ // User a bigger flow control window and onready threshold for directpath to prevent churn when
+ // gRPC is trying to flush gRPCs over the wire. If onReadyThreshold and flowControlWindowBytes are
+ // too low, it was observed in testing that the gRPC channel can get stuck in a "not-ready" state
+ // until stream deadline.
+ public static final UserWorkerGrpcFlowControlSettings DEFAULT_DIRECTPATH_FLOW_CONTROL_SETTINGS =
+ UserWorkerGrpcFlowControlSettings.newBuilder()
+ .setFlowControlWindowBytes(DIRECTPATH_MIN_FLOW_CONTROL_WINDOW)
+ .setOnReadyThresholdBytes(DIRECTPATH_ON_READY_THRESHOLD)
+ // Prevent gRPC from automatically resizing the window. If we have things to send/receive
+ // from Windmill we want to do it right away. There are internal pushback mechanisms in
+ // the user worker and Windmill that attempt to guard the process from OOMing (i.e
+ // MemoryMonitor.java).
+ .setEnableAutoFlowControl(false)
+ .build();
- private WindmillChannelFactory() {}
+ private WindmillChannels() {}
public static ManagedChannel inProcessChannel(String channelName) {
return InProcessChannelBuilder.forName(channelName).directExecutor().build();
@@ -54,15 +81,20 @@ public final class WindmillChannelFactory {
}
public static ManagedChannel remoteChannel(
- WindmillServiceAddress windmillServiceAddress, int windmillServiceRpcChannelTimeoutSec) {
+ WindmillServiceAddress windmillServiceAddress,
+ int windmillServiceRpcChannelTimeoutSec,
+ UserWorkerGrpcFlowControlSettings flowControlSettings) {
switch (windmillServiceAddress.getKind()) {
case GCP_SERVICE_ADDRESS:
return remoteChannel(
- windmillServiceAddress.gcpServiceAddress(), windmillServiceRpcChannelTimeoutSec);
+ windmillServiceAddress.gcpServiceAddress(),
+ windmillServiceRpcChannelTimeoutSec,
+ flowControlSettings);
case AUTHENTICATED_GCP_SERVICE_ADDRESS:
return remoteDirectChannel(
windmillServiceAddress.authenticatedGcpServiceAddress(),
- windmillServiceRpcChannelTimeoutSec);
+ windmillServiceRpcChannelTimeoutSec,
+ flowControlSettings);
default:
throw new UnsupportedOperationException(
"Only GCP_SERVICE_ADDRESS and AUTHENTICATED_GCP_SERVICE_ADDRESS are supported"
@@ -72,32 +104,56 @@ public final class WindmillChannelFactory {
private static ManagedChannel remoteDirectChannel(
AuthenticatedGcpServiceAddress authenticatedGcpServiceAddress,
- int windmillServiceRpcChannelTimeoutSec) {
- return withDefaultChannelOptions(
+ int windmillServiceRpcChannelTimeoutSec,
+ UserWorkerGrpcFlowControlSettings flowControlSettings) {
+ NettyChannelBuilder channelBuilder =
+ withDefaultChannelOptions(
NettyChannelBuilder.forAddress(
authenticatedGcpServiceAddress.gcpServiceAddress().getHost(),
// Ports are required for direct channels.
authenticatedGcpServiceAddress.gcpServiceAddress().getPort(),
new AltsChannelCredentials.Builder().build())
.overrideAuthority(authenticatedGcpServiceAddress.authenticatingService()),
- windmillServiceRpcChannelTimeoutSec)
+ windmillServiceRpcChannelTimeoutSec);
+
+ return withFlowControlSettings(
+ channelBuilder, flowControlSettings, DIRECTPATH_MIN_FLOW_CONTROL_WINDOW)
.build();
}
public static ManagedChannel remoteChannel(
- HostAndPort endpoint, int windmillServiceRpcChannelTimeoutSec) {
- return withDefaultChannelOptions(
- NettyChannelBuilder.forAddress(
- endpoint.getHost(),
- endpoint.hasPort()
- ? endpoint.getPort()
- : WindmillEndpoints.DEFAULT_WINDMILL_SERVICE_PORT),
- windmillServiceRpcChannelTimeoutSec)
- .negotiationType(NegotiationType.TLS)
- .sslContext(dataflowGrpcSslContext(endpoint))
+ HostAndPort endpoint,
+ int windmillServiceRpcChannelTimeoutSec,
+ UserWorkerGrpcFlowControlSettings flowControlSettings) {
+ NettyChannelBuilder channelBuilder =
+ withDefaultChannelOptions(
+ NettyChannelBuilder.forAddress(
+ endpoint.getHost(),
+ endpoint.hasPort()
+ ? endpoint.getPort()
+ : WindmillEndpoints.DEFAULT_WINDMILL_SERVICE_PORT),
+ windmillServiceRpcChannelTimeoutSec)
+ .negotiationType(NegotiationType.TLS)
+ .sslContext(dataflowGrpcSslContext(endpoint));
+
+ return withFlowControlSettings(
+ channelBuilder, flowControlSettings, WINDMILL_MIN_FLOW_CONTROL_WINDOW)
.build();
}
+ private static NettyChannelBuilder withFlowControlSettings(
+ NettyChannelBuilder channelBuilder,
+ UserWorkerGrpcFlowControlSettings flowControlSettings,
+ int defaultFlowControlBytes) {
+ int flowControlWindowSizeBytes =
+ Math.max(defaultFlowControlBytes, flowControlSettings.getFlowControlWindowBytes());
+ return flowControlSettings.getEnableAutoFlowControl()
+ // Enable auto flow control with an initial window of flowControlWindowSizeBytes. gRPC may
+ // resize this value based on throughput.
+ ? channelBuilder.initialFlowControlWindow(flowControlWindowSizeBytes)
+ : channelBuilder.flowControlWindow(flowControlWindowSizeBytes);
+ }
+
@SuppressWarnings("nullness")
private static SslContext dataflowGrpcSslContext(HostAndPort endpoint) {
try {
@@ -121,9 +177,7 @@ public final class WindmillChannelFactory {
return channelBuilder
.maxInboundMessageSize(Integer.MAX_VALUE)
.maxTraceEvents(MAX_REMOTE_TRACE_EVENTS)
- // 1MiB
- .maxInboundMetadataSize(1024 * 1024)
- .flowControlWindow(WINDMILL_MAX_FLOW_CONTROL_WINDOW);
+ .maxInboundMetadataSize(MAX_INBOUND_METADATA_SIZE_BYTES);
}
private static class WindmillChannelCreationException extends IllegalStateException {
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/stubs/WindmillStubFactoryFactoryImpl.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/stubs/WindmillStubFactoryFactoryImpl.java
index 98ff93733d6..6de6b633753 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/stubs/WindmillStubFactoryFactoryImpl.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/stubs/WindmillStubFactoryFactoryImpl.java
@@ -17,16 +17,12 @@
*/
package org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs;
-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.auth.Credentials;
-import java.util.function.Function;
import org.apache.beam.runners.dataflow.options.DataflowWorkerHarnessOptions;
-import org.apache.beam.runners.dataflow.worker.windmill.WindmillServiceAddress;
-import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.ManagedChannel;
public class WindmillStubFactoryFactoryImpl implements WindmillStubFactoryFactory {
-
private final int windmillServiceRpcChannelAliveTimeoutSec;
private final Credentials gcpCredential;
@@ -38,17 +34,23 @@ public class WindmillStubFactoryFactoryImpl implements WindmillStubFactoryFactor
@Override
public WindmillStubFactory makeWindmillStubFactory(boolean useIsolatedChannels) {
- Function<WindmillServiceAddress, ManagedChannel> channelFactory =
- serviceAddress -> remoteChannel(serviceAddress, windmillServiceRpcChannelAliveTimeoutSec);
ChannelCache channelCache =
ChannelCache.create(
- serviceAddress ->
+ (flowControlSettings, serviceAddress) ->
// IsolationChannel will create and manage separate RPC channels to the same
// serviceAddress via calling the channelFactory, else just directly return the
// RPC channel.
useIsolatedChannels
- ? IsolationChannel.create(() -> channelFactory.apply(serviceAddress))
- : channelFactory.apply(serviceAddress));
+ ? IsolationChannel.create(
+ () ->
+ remoteChannel(
+ serviceAddress.getServiceAddress(),
+ windmillServiceRpcChannelAliveTimeoutSec,
+ flowControlSettings))
+ : remoteChannel(
+ serviceAddress.getServiceAddress(),
+ windmillServiceRpcChannelAliveTimeoutSec,
+ flowControlSettings));
return ChannelCachingRemoteStubFactory.create(gcpCredential, channelCache);
}
}
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java
index 539520a2edc..c60228f336e 100644
--- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java
@@ -132,7 +132,7 @@ import org.apache.beam.runners.dataflow.worker.windmill.Windmill.Timer.Type;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WatermarkHold;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkItemCommitRequest;
import org.apache.beam.runners.dataflow.worker.windmill.client.getdata.FakeGetDataClient;
-import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs.WindmillChannelFactory;
+import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs.WindmillChannels;
import org.apache.beam.runners.dataflow.worker.windmill.testing.FakeWindmillStubFactory;
import org.apache.beam.runners.dataflow.worker.windmill.testing.FakeWindmillStubFactoryFactory;
import org.apache.beam.runners.dataflow.worker.windmill.work.refresh.HeartbeatSender;
@@ -883,8 +883,7 @@ public class StreamingDataflowWorkerTest {
new FakeWindmillStubFactoryFactory(
new FakeWindmillStubFactory(
() ->
- WindmillChannelFactory.inProcessChannel(
- "StreamingDataflowWorkerTestChannel"))));
+ WindmillChannels.inProcessChannel("StreamingDataflowWorkerTestChannel"))));
this.computationStateCache = worker.getComputationStateCache();
return worker;
}
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/harness/FanOutStreamingEngineWorkerHarnessTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/harness/FanOutStreamingEngineWorkerHarnessTest.java
index ff98b7ed322..b3aa321b980 100644
--- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/harness/FanOutStreamingEngineWorkerHarnessTest.java
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/harness/FanOutStreamingEngineWorkerHarnessTest.java
@@ -51,7 +51,7 @@ import org.apache.beam.runners.dataflow.worker.windmill.client.getdata.Throttlin
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.grpc.stubs.WindmillChannelFactory;
+import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs.WindmillChannels;
import org.apache.beam.runners.dataflow.worker.windmill.testing.FakeWindmillStubFactory;
import org.apache.beam.runners.dataflow.worker.windmill.testing.FakeWindmillStubFactoryFactory;
import org.apache.beam.runners.dataflow.worker.windmill.work.WorkItemScheduler;
@@ -81,7 +81,7 @@ import org.junit.runners.JUnit4;
public class FanOutStreamingEngineWorkerHarnessTest {
private static final String CHANNEL_NAME = "FanOutStreamingEngineWorkerHarnessTest";
private static final WindmillServiceAddress DEFAULT_WINDMILL_SERVICE_ADDRESS =
- WindmillServiceAddress.create(HostAndPort.fromParts(WindmillChannelFactory.LOCALHOST, 443));
+ WindmillServiceAddress.create(HostAndPort.fromParts(WindmillChannels.LOCALHOST, 443));
private static final ImmutableMap<String, WorkerMetadataResponse.Endpoint> DEFAULT =
ImmutableMap.of(
"global_data",
@@ -107,7 +107,7 @@ public class FanOutStreamingEngineWorkerHarnessTest {
spy(GrpcWindmillStreamFactory.of(JOB_HEADER).build());
private final ChannelCachingStubFactory stubFactory =
new FakeWindmillStubFactory(
- () -> grpcCleanup.register(WindmillChannelFactory.inProcessChannel(CHANNEL_NAME)));
+ () -> grpcCleanup.register(WindmillChannels.inProcessChannel(CHANNEL_NAME)));
private final GrpcDispatcherClient dispatcherClient =
GrpcDispatcherClient.forTesting(
PipelineOptionsFactory.as(DataflowWorkerHarnessOptions.class),
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/harness/WindmillStreamSenderTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/harness/WindmillStreamSenderTest.java
index 21d47b2cd83..e570a83e345 100644
--- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/harness/WindmillStreamSenderTest.java
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/harness/WindmillStreamSenderTest.java
@@ -82,7 +82,7 @@ public class WindmillStreamSenderTest {
grpcCleanup.register(inProcessChannel);
connection =
WindmillConnection.builder()
- .setStub(CloudWindmillServiceV1Alpha1Grpc.newStub(inProcessChannel))
+ .setStubSupplier(() -> CloudWindmillServiceV1Alpha1Grpc.newStub(inProcessChannel))
.build();
}
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcDirectGetWorkStreamTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcDirectGetWorkStreamTest.java
index 6cb68191f2d..b9592e9aacc 100644
--- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcDirectGetWorkStreamTest.java
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcDirectGetWorkStreamTest.java
@@ -143,7 +143,8 @@ public class GrpcDirectGetWorkStreamTest {
.build()
.createDirectGetWorkStream(
WindmillConnection.builder()
- .setStub(CloudWindmillServiceV1Alpha1Grpc.newStub(inProcessChannel))
+ .setStubSupplier(
+ () -> CloudWindmillServiceV1Alpha1Grpc.newStub(inProcessChannel))
.build(),
Windmill.GetWorkRequest.newBuilder()
.setClientId(TEST_JOB_HEADER.getClientId())
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcWindmillServerTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcWindmillServerTest.java
index bdb482ba768..211e926a073 100644
--- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcWindmillServerTest.java
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcWindmillServerTest.java
@@ -72,7 +72,7 @@ import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStream.Co
import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStream.GetDataStream;
import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStream.GetWorkStream;
import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStreamShutdownException;
-import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs.WindmillChannelFactory;
+import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs.WindmillChannels;
import org.apache.beam.runners.dataflow.worker.windmill.testing.FakeWindmillStubFactory;
import org.apache.beam.runners.dataflow.worker.windmill.testing.FakeWindmillStubFactoryFactory;
import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString;
@@ -156,7 +156,7 @@ public class GrpcWindmillServerTest {
new FakeWindmillStubFactory(
() -> {
ManagedChannel channel =
- grpcCleanup.register(WindmillChannelFactory.inProcessChannel(name));
+ grpcCleanup.register(WindmillChannels.inProcessChannel(name));
openedChannels.add(channel);
return channel;
})));
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/stubs/ChannelCacheTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/stubs/ChannelCacheTest.java
index 704a9c9fdb2..311bed75ccc 100644
--- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/stubs/ChannelCacheTest.java
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/stubs/ChannelCacheTest.java
@@ -18,6 +18,7 @@
package org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs;
import static com.google.common.truth.Truth.assertThat;
+import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertSame;
import static org.junit.Assert.assertTrue;
import static org.mockito.ArgumentMatchers.eq;
@@ -25,11 +26,18 @@ import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+import java.util.Queue;
import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Function;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.UserWorkerGrpcFlowControlSettings;
import org.apache.beam.runners.dataflow.worker.windmill.WindmillServiceAddress;
import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.ManagedChannel;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
import org.junit.After;
import org.junit.Rule;
import org.junit.Test;
@@ -45,7 +53,7 @@ public class ChannelCacheTest {
private static ChannelCache newCache(
Function<WindmillServiceAddress, ManagedChannel> channelFactory) {
- return ChannelCache.forTesting(channelFactory, () -> {});
+ return ChannelCache.forTesting((ignored, address) -> channelFactory.apply(address), () -> {});
}
@After
@@ -56,7 +64,7 @@ public class ChannelCacheTest {
}
private ManagedChannel newChannel(String channelName) {
- return WindmillChannelFactory.inProcessChannel(channelName);
+ return WindmillChannels.inProcessChannel(channelName);
}
@Test
@@ -108,7 +116,8 @@ public class ChannelCacheTest {
CountDownLatch notifyWhenChannelClosed = new CountDownLatch(1);
cache =
ChannelCache.forTesting(
- ignored -> newChannel(channelName), notifyWhenChannelClosed::countDown);
+ (ignoredFlowControlSettings, ignoredServiceAddress) -> newChannel(channelName),
+ notifyWhenChannelClosed::countDown);
WindmillServiceAddress someAddress = mock(WindmillServiceAddress.class);
ManagedChannel cachedChannel = cache.get(someAddress);
@@ -132,7 +141,8 @@ public class ChannelCacheTest {
CountDownLatch notifyWhenChannelClosed = new CountDownLatch(1);
cache =
ChannelCache.forTesting(
- ignored -> newChannel(channelName), notifyWhenChannelClosed::countDown);
+ (ignoredFlowControlSettings, ignoredServiceAddress) -> newChannel(channelName),
+ notifyWhenChannelClosed::countDown);
WindmillServiceAddress someAddress = mock(WindmillServiceAddress.class);
ManagedChannel cachedChannel = cache.get(someAddress);
cache.clear();
@@ -140,4 +150,89 @@ public class ChannelCacheTest {
assertTrue(cache.isEmpty());
assertTrue(cachedChannel.isShutdown());
}
+
+ @Test
+ public void testConsumeFlowControlSettings() throws InterruptedException {
+ String channelName = "channel";
+ CountDownLatch notifyWhenChannelClosed = new CountDownLatch(1);
+ AtomicInteger newChannelsCreated = new AtomicInteger();
+ AtomicReference<UserWorkerGrpcFlowControlSettings> consumedFlowControlSettings =
+ new AtomicReference<>();
+ cache =
+ ChannelCache.forTesting(
+ (newFlowControlSettings, ignoredServiceAddress) -> {
+ ManagedChannel channel = newChannel(channelName);
+ newChannelsCreated.incrementAndGet();
+ consumedFlowControlSettings.set(newFlowControlSettings);
+ return channel;
+ },
+ notifyWhenChannelClosed::countDown);
+ WindmillServiceAddress someAddress = mock(WindmillServiceAddress.class);
+ when(someAddress.getKind())
+ .thenReturn(WindmillServiceAddress.Kind.AUTHENTICATED_GCP_SERVICE_ADDRESS);
+
+ // Load the cache w/ this first get.
+ ManagedChannel cachedChannel = cache.get(someAddress);
+
+ // Load flow control settings.
+ UserWorkerGrpcFlowControlSettings flowControlSettings =
+ UserWorkerGrpcFlowControlSettings.newBuilder()
+ .setEnableAutoFlowControl(true)
+ .setOnReadyThresholdBytes(1)
+ .setFlowControlWindowBytes(1)
+ .build();
+ cache.consumeFlowControlSettings(flowControlSettings);
+
+ // This get should reload the cache, since flow control settings have changed.
+ ManagedChannel reloadedChannel = cache.get(someAddress);
+ notifyWhenChannelClosed.await();
+ assertThat(cachedChannel).isNotSameInstanceAs(reloadedChannel);
+ assertTrue(cachedChannel.isShutdown());
+ assertFalse(reloadedChannel.isShutdown());
+ assertThat(newChannelsCreated.get()).isEqualTo(2);
+ assertThat(cache.get(someAddress)).isSameInstanceAs(reloadedChannel);
+ assertThat(consumedFlowControlSettings.get()).isEqualTo(flowControlSettings);
+ }
+
+ @Test
+ public void testConsumeFlowControlSettings_sameFlowControlSettings() {
+ String channelName = "channel";
+ AtomicInteger newChannelsCreated = new AtomicInteger();
+ UserWorkerGrpcFlowControlSettings flowControlSettings =
+ UserWorkerGrpcFlowControlSettings.newBuilder()
+ .setEnableAutoFlowControl(true)
+ .setOnReadyThresholdBytes(1)
+ .setFlowControlWindowBytes(1)
+ .build();
+ Queue<UserWorkerGrpcFlowControlSettings> consumedFlowControlSettings =
+ new LinkedBlockingQueue<>();
+ cache =
+ ChannelCache.forTesting(
+ (newFlowControlSettings, ignoredServiceAddress) -> {
+ ManagedChannel channel = newChannel(channelName);
+ newChannelsCreated.incrementAndGet();
+ consumedFlowControlSettings.add(newFlowControlSettings);
+ return channel;
+ },
+ () -> {});
+ // Load flow control settings.
+ cache.consumeFlowControlSettings(flowControlSettings);
+ WindmillServiceAddress someAddress = mock(WindmillServiceAddress.class);
+ when(someAddress.getKind())
+ .thenReturn(WindmillServiceAddress.Kind.AUTHENTICATED_GCP_SERVICE_ADDRESS);
+ // Load the cache w/ this first get.
+ ManagedChannel cachedChannel = cache.get(someAddress);
+
+ // Load the same flow control settings.
+ cache.consumeFlowControlSettings(flowControlSettings);
+ // Because the flow control settings are the same, cache should not reload the value for the
+ // same address.
+ ManagedChannel reloadedChannel = cache.get(someAddress);
+ assertThat(cachedChannel).isSameInstanceAs(reloadedChannel);
+ assertFalse(cachedChannel.isShutdown());
+ assertThat(newChannelsCreated.get()).isEqualTo(1);
+ assertThat(consumedFlowControlSettings.size()).isEqualTo(1);
+ assertThat(Iterables.getOnlyElement(consumedFlowControlSettings))
+ .isEqualTo(flowControlSettings);
+ }
}
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/testing/FakeWindmillStubFactory.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/testing/FakeWindmillStubFactory.java
index b0897574d9a..79297e51a55 100644
--- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/testing/FakeWindmillStubFactory.java
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/testing/FakeWindmillStubFactory.java
@@ -31,7 +31,9 @@ public final class FakeWindmillStubFactory implements ChannelCachingStubFactory
private final ChannelCache channelCache;
public FakeWindmillStubFactory(Supplier<ManagedChannel> channelFactory) {
- this.channelCache = ChannelCache.forTesting(ignored -> channelFactory.get(), () -> {});
+ this.channelCache =
+ ChannelCache.forTesting(
+ (ignoredFlowControlSettings, ignoredServiceAddress) -> channelFactory.get(), () -> {});
}
@Override