You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by sc...@apache.org on 2024/05/14 07:54:51 UTC
(beam) branch master updated: add windmill computation state cache, config loaders, and status pages (#31133)
This is an automated email from the ASF dual-hosted git repository.
scwhittle 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 31e81ffe12a add windmill computation state cache, config loaders, and status pages (#31133)
31e81ffe12a is described below
commit 31e81ffe12ac5b8d7b7d56aab995d536a7036c85
Author: martin trieu <ma...@google.com>
AuthorDate: Tue May 14 00:54:44 2024 -0700
add windmill computation state cache, config loaders, and status pages (#31133)
* add computation state cache, config loaders, and status pages
* clean up StreamingDataflowWorkerTest to use builder pattern when constructor StreamingDataflowWorker instance
---
.../dataflow/worker/StreamingDataflowWorker.java | 792 +++++++--------------
.../dataflow/worker/status/WorkerStatusPages.java | 4 +
.../worker/streaming/ComputationState.java | 27 +-
.../worker/streaming/ComputationStateCache.java | 218 ++++++
.../worker/streaming/config/ComputationConfig.java | 59 ++
...StreamingApplianceComputationConfigFetcher.java | 136 ++++
.../StreamingEngineComputationConfigFetcher.java | 283 ++++++++
.../config/StreamingEnginePipelineConfig.java | 57 ++
.../streaming/harness/MetricsDataProvider.java | 83 +++
.../harness/StreamingWorkerStatusPages.java | 261 +++++++
.../worker/util/common/worker/WorkExecutor.java | 2 +-
.../worker/windmill/WindmillServerBase.java | 2 +-
.../worker/windmill/WindmillServerStub.java | 3 +-
.../windmill/client/grpc/GrpcDispatcherClient.java | 2 +-
.../windmill/client/grpc/GrpcGetDataStream.java | 2 +-
.../windmill/client/grpc/GrpcWindmillServer.java | 30 +-
.../worker/StreamingDataflowWorkerTest.java | 407 ++++++-----
.../streaming/ComputationStateCacheTest.java | 282 ++++++++
...amingApplianceComputationConfigFetcherTest.java | 142 ++++
...treamingEngineComputationConfigFetcherTest.java | 219 ++++++
20 files changed, 2263 insertions(+), 748 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 660680cda3e..f18d5fac721 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,30 +17,19 @@
*/
package org.apache.beam.runners.dataflow.worker;
-import static java.nio.charset.StandardCharsets.UTF_8;
-import static java.util.stream.Collectors.toConcurrentMap;
import static org.apache.beam.runners.dataflow.DataflowRunner.hasExperiment;
+import static org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs.WindmillChannelFactory.remoteChannel;
import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
import com.google.api.services.dataflow.model.CounterUpdate;
import com.google.api.services.dataflow.model.MapTask;
-import com.google.api.services.dataflow.model.StreamingComputationConfig;
-import com.google.api.services.dataflow.model.StreamingConfigTask;
-import com.google.api.services.dataflow.model.WorkItem;
-import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
-import java.io.File;
-import java.io.IOException;
-import java.io.PrintWriter;
-import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
-import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Random;
-import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.ConcurrentMap;
@@ -48,18 +37,17 @@ import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.Supplier;
-import javax.servlet.http.HttpServletRequest;
-import javax.servlet.http.HttpServletResponse;
+import org.apache.beam.repackaged.core.org.apache.commons.lang3.tuple.Pair;
import org.apache.beam.runners.core.metrics.MetricsLogger;
import org.apache.beam.runners.dataflow.DataflowRunner;
import org.apache.beam.runners.dataflow.internal.CustomSources;
import org.apache.beam.runners.dataflow.options.DataflowWorkerHarnessOptions;
import org.apache.beam.runners.dataflow.util.CloudObject;
import org.apache.beam.runners.dataflow.util.CloudObjects;
-import org.apache.beam.runners.dataflow.worker.apiary.FixMultiOutputInfosOnParDoInstructions;
import org.apache.beam.runners.dataflow.worker.counters.DataflowCounterUpdateExtractor;
import org.apache.beam.runners.dataflow.worker.counters.NameContext;
import org.apache.beam.runners.dataflow.worker.graph.Edges.Edge;
@@ -70,13 +58,10 @@ import org.apache.beam.runners.dataflow.worker.graph.Nodes.Node;
import org.apache.beam.runners.dataflow.worker.graph.Nodes.ParallelInstructionNode;
import org.apache.beam.runners.dataflow.worker.logging.DataflowWorkerLoggingMDC;
import org.apache.beam.runners.dataflow.worker.profiler.ScopedProfiler;
-import org.apache.beam.runners.dataflow.worker.status.BaseStatusServlet;
import org.apache.beam.runners.dataflow.worker.status.DebugCapture;
-import org.apache.beam.runners.dataflow.worker.status.DebugCapture.Capturable;
-import org.apache.beam.runners.dataflow.worker.status.LastExceptionDataProvider;
-import org.apache.beam.runners.dataflow.worker.status.StatusDataProvider;
import org.apache.beam.runners.dataflow.worker.status.WorkerStatusPages;
import org.apache.beam.runners.dataflow.worker.streaming.ComputationState;
+import org.apache.beam.runners.dataflow.worker.streaming.ComputationStateCache;
import org.apache.beam.runners.dataflow.worker.streaming.ExecutionState;
import org.apache.beam.runners.dataflow.worker.streaming.KeyCommitTooLargeException;
import org.apache.beam.runners.dataflow.worker.streaming.ShardedKey;
@@ -84,7 +69,12 @@ import org.apache.beam.runners.dataflow.worker.streaming.StageInfo;
import org.apache.beam.runners.dataflow.worker.streaming.Work;
import org.apache.beam.runners.dataflow.worker.streaming.Work.State;
import org.apache.beam.runners.dataflow.worker.streaming.WorkHeartbeatResponseProcessor;
+import org.apache.beam.runners.dataflow.worker.streaming.config.ComputationConfig;
+import org.apache.beam.runners.dataflow.worker.streaming.config.StreamingApplianceComputationConfigFetcher;
+import org.apache.beam.runners.dataflow.worker.streaming.config.StreamingEngineComputationConfigFetcher;
+import org.apache.beam.runners.dataflow.worker.streaming.config.StreamingEnginePipelineConfig;
import org.apache.beam.runners.dataflow.worker.streaming.harness.StreamingCounters;
+import org.apache.beam.runners.dataflow.worker.streaming.harness.StreamingWorkerStatusPages;
import org.apache.beam.runners.dataflow.worker.streaming.harness.StreamingWorkerStatusReporter;
import org.apache.beam.runners.dataflow.worker.streaming.sideinput.SideInputStateFetcher;
import org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor;
@@ -97,6 +87,7 @@ import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.LatencyAttribution;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkItemCommitRequest;
import org.apache.beam.runners.dataflow.worker.windmill.WindmillServerStub;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillServiceAddress;
import org.apache.beam.runners.dataflow.worker.windmill.appliance.JniWindmillApplianceServer;
import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStream.GetWorkStream;
import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStreamPool;
@@ -106,8 +97,13 @@ import org.apache.beam.runners.dataflow.worker.windmill.client.commits.Streaming
import org.apache.beam.runners.dataflow.worker.windmill.client.commits.StreamingEngineWorkCommitter;
import org.apache.beam.runners.dataflow.worker.windmill.client.commits.WorkCommitter;
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;
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.state.WindmillStateCache;
import org.apache.beam.runners.dataflow.worker.windmill.state.WindmillStateReader;
import org.apache.beam.runners.dataflow.worker.windmill.work.processing.failures.FailureTracker;
@@ -118,7 +114,6 @@ import org.apache.beam.runners.dataflow.worker.windmill.work.refresh.ActiveWorkR
import org.apache.beam.runners.dataflow.worker.windmill.work.refresh.ActiveWorkRefreshers;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.KvCoder;
-import org.apache.beam.sdk.extensions.gcp.util.Transport;
import org.apache.beam.sdk.fn.IdGenerator;
import org.apache.beam.sdk.fn.IdGenerators;
import org.apache.beam.sdk.fn.JvmInitializers;
@@ -126,16 +121,12 @@ import org.apache.beam.sdk.io.FileSystems;
import org.apache.beam.sdk.io.gcp.bigquery.BigQuerySinkMetrics;
import org.apache.beam.sdk.metrics.MetricName;
import org.apache.beam.sdk.metrics.MetricsEnvironment;
-import org.apache.beam.sdk.util.BackOff;
-import org.apache.beam.sdk.util.BackOffUtils;
-import org.apache.beam.sdk.util.FluentBackoff;
-import org.apache.beam.sdk.util.Sleeper;
import org.apache.beam.sdk.util.WindowedValue.WindowedValueCoder;
import org.apache.beam.sdk.util.construction.CoderTranslation;
import org.apache.beam.vendor.grpc.v1p60p1.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p60p1.io.grpc.ManagedChannel;
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.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Splitter;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.Cache;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.CacheBuilder;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.*;
@@ -151,7 +142,6 @@ import org.slf4j.LoggerFactory;
/** Implements a Streaming Dataflow worker. */
@SuppressWarnings({
- "rawtypes", // TODO(https://github.com/apache/beam/issues/20447)
"nullness" // TODO(https://github.com/apache/beam/issues/20497)
})
public class StreamingDataflowWorker {
@@ -179,13 +169,8 @@ public class StreamingDataflowWorker {
private static final Logger LOG = LoggerFactory.getLogger(StreamingDataflowWorker.class);
/** The idGenerator to generate unique id globally. */
- private static final IdGenerator idGenerator = IdGenerators.decrementingLongs();
- /**
- * Fix up MapTask representation because MultiOutputInfos are missing from system generated
- * ParDoInstructions.
- */
- private static final Function<MapTask, MapTask> fixMultiOutputInfos =
- new FixMultiOutputInfosOnParDoInstructions(idGenerator);
+ private static final IdGenerator ID_GENERATOR = IdGenerators.decrementingLongs();
+
/**
* Function which converts map tasks to their network representation for execution.
*
@@ -194,8 +179,8 @@ public class StreamingDataflowWorker {
* <li>Remove flatten instructions by rewiring edges.
* </ul>
*/
- private static final Function<MapTask, MutableNetwork<Node, Edge>> mapTaskToBaseNetwork =
- new MapTaskToNetworkFunction(idGenerator);
+ private static final Function<MapTask, MutableNetwork<Node, Edge>> MAP_TASK_TO_BASE_NETWORK_FN =
+ new MapTaskToNetworkFunction(ID_GENERATOR);
private static final int DEFAULT_STATUS_PORT = 8081;
// Maximum size of the result of a GetWork request.
@@ -206,19 +191,16 @@ public class StreamingDataflowWorker {
private static final Random clientIdGenerator = new Random();
private static final String CHANNELZ_PATH = "/channelz";
final WindmillStateCache stateCache;
+ private final StreamingWorkerStatusPages statusPages;
+ private final ComputationConfig.Fetcher configFetcher;
+ private final ComputationStateCache computationStateCache;
// Maps from computation ids to per-computation state.
- private final ConcurrentMap<String, ComputationState> computationMap;
// Cache of tokens to commit callbacks.
// Using Cache with time eviction policy helps us to prevent memory leak when callback ids are
// discarded by Dataflow service and calling commitCallback is best-effort.
private final Cache<Long, Runnable> commitCallbacks =
CacheBuilder.newBuilder().expireAfterWrite(5L, TimeUnit.MINUTES).build();
- // Map of user state names to system state names.
- // TODO(drieber): obsolete stateNameMap. Use transformUserNameToStateFamily in
- // ComputationState instead.
- private final ConcurrentMap<String, String> stateNameMap = new ConcurrentHashMap<>();
- private final ConcurrentMap<String, String> systemNameToComputationIdMap =
- new ConcurrentHashMap<>();
+
private final BoundedQueueExecutor workUnitExecutor;
private final WindmillServerStub windmillServer;
private final Thread dispatchThread;
@@ -236,28 +218,17 @@ public class StreamingDataflowWorker {
private final MemoryMonitor memoryMonitor;
private final Thread memoryMonitorThread;
- private final WorkerStatusPages statusPages;
// Limit on bytes sinked (committed) in a work item.
private final long maxSinkBytes; // = MAX_SINK_BYTES unless disabled in options.
private final ReaderCache readerCache;
- private final WorkUnitClient workUnitClient;
- private final CompletableFuture<Void> isDoneFuture;
private final Function<MapTask, MutableNetwork<Node, Edge>> mapTaskToNetwork;
private final ReaderRegistry readerRegistry = ReaderRegistry.defaultRegistry();
private final SinkRegistry sinkRegistry = SinkRegistry.defaultRegistry();
private final Supplier<Instant> clock;
- private final Function<String, ScheduledExecutorService> executorSupplier;
private final DataflowMapTaskExecutorFactory mapTaskExecutorFactory;
private final HotKeyLogger hotKeyLogger;
- // Periodic sender of debug information to the debug capture service.
- private final DebugCapture.@Nullable Manager debugCaptureManager;
- // Collection of ScheduledExecutorServices that are running periodic functions.
- private final ArrayList<ScheduledExecutorService> scheduledExecutors = new ArrayList<>();
- // Periodically fires a global config request to dataflow service. Only used when windmill service
- // is enabled.
// Possibly overridden by streaming engine config.
- private int maxWorkItemCommitBytes = Integer.MAX_VALUE;
-
+ private final AtomicInteger maxWorkItemCommitBytes;
private final DataflowExecutionStateSampler sampler = DataflowExecutionStateSampler.instance();
private final ActiveWorkRefresher activeWorkRefresher;
private final WorkCommitter workCommitter;
@@ -266,15 +237,15 @@ public class StreamingDataflowWorker {
private final WorkFailureProcessor workFailureProcessor;
private final StreamingCounters streamingCounters;
- private StreamingDataflowWorker(
+ StreamingDataflowWorker(
WindmillServerStub windmillServer,
long clientId,
- ConcurrentMap<String, ComputationState> computationMap,
+ ComputationConfig.Fetcher configFetcher,
+ ComputationStateCache computationStateCache,
ConcurrentMap<String, StageInfo> stageInfoMap,
WindmillStateCache windmillStateCache,
BoundedQueueExecutor workUnitExecutor,
DataflowMapTaskExecutorFactory mapTaskExecutorFactory,
- WorkUnitClient workUnitClient,
DataflowWorkerHarnessOptions options,
HotKeyLogger hotKeyLogger,
Supplier<Instant> clock,
@@ -283,8 +254,11 @@ public class StreamingDataflowWorker {
WorkFailureProcessor workFailureProcessor,
StreamingCounters streamingCounters,
MemoryMonitor memoryMonitor,
+ AtomicInteger maxWorkItemCommitBytes,
+ GrpcWindmillStreamFactory windmillStreamFactory,
Function<String, ScheduledExecutorService> executorSupplier) {
- this.computationMap = computationMap;
+ this.configFetcher = configFetcher;
+ this.computationStateCache = computationStateCache;
this.stageInfoMap = stageInfoMap;
this.stateCache = windmillStateCache;
this.readerCache =
@@ -292,21 +266,27 @@ public class StreamingDataflowWorker {
Duration.standardSeconds(options.getReaderCacheTimeoutSec()),
Executors.newCachedThreadPool());
this.mapTaskExecutorFactory = mapTaskExecutorFactory;
- this.workUnitClient = workUnitClient;
this.options = options;
this.hotKeyLogger = hotKeyLogger;
this.clock = clock;
- this.executorSupplier = executorSupplier;
+ this.maxWorkItemCommitBytes = maxWorkItemCommitBytes;
this.windmillServiceEnabled = options.isEnableStreamingEngine();
- this.statusPages = WorkerStatusPages.create(DEFAULT_STATUS_PORT, memoryMonitor, () -> true);
- if (windmillServiceEnabled) {
- this.debugCaptureManager =
- new DebugCapture.Manager(options, statusPages.getDebugCapturePages());
- } else {
- this.debugCaptureManager = null;
+ int numCommitThreads = 1;
+ if (windmillServiceEnabled && options.getWindmillServiceCommitThreads() > 0) {
+ numCommitThreads = options.getWindmillServiceCommitThreads();
}
- this.isDoneFuture = new CompletableFuture<>();
+ this.workCommitter =
+ windmillServiceEnabled
+ ? StreamingEngineWorkCommitter.create(
+ WindmillStreamPool.create(
+ NUM_COMMIT_STREAMS, COMMIT_STREAM_TIMEOUT, windmillServer::commitWorkStream)
+ ::getCloseableStream,
+ numCommitThreads,
+ this::onCompleteCommit)
+ : StreamingApplianceWorkCommitter.create(
+ windmillServer::commitWork, this::onCompleteCommit);
+
this.workUnitExecutor = workUnitExecutor;
maxSinkBytes =
@@ -343,26 +323,11 @@ public class StreamingDataflowWorker {
this.sideInputStateFetcher =
new SideInputStateFetcher(metricTrackingWindmillServer::getSideInputData, options);
- int numCommitThreads = 1;
- if (windmillServiceEnabled && options.getWindmillServiceCommitThreads() > 0) {
- numCommitThreads = options.getWindmillServiceCommitThreads();
- }
-
- this.workCommitter =
- windmillServiceEnabled
- ? StreamingEngineWorkCommitter.create(
- WindmillStreamPool.create(
- NUM_COMMIT_STREAMS, COMMIT_STREAM_TIMEOUT, windmillServer::commitWorkStream)
- ::getCloseableStream,
- numCommitThreads,
- this::onCompleteCommit)
- : StreamingApplianceWorkCommitter.create(
- windmillServer::commitWork, this::onCompleteCommit);
// Register standard file systems.
FileSystems.setDefaultPipelineOptions(options);
- this.mapTaskToNetwork = mapTaskToBaseNetwork;
+ this.mapTaskToNetwork = MAP_TASK_TO_BASE_NETWORK_FN;
int stuckCommitDurationMillis =
windmillServiceEnabled && options.getStuckCommitDurationMillis() > 0
@@ -373,10 +338,35 @@ public class StreamingDataflowWorker {
clock,
options.getActiveWorkRefreshPeriodMillis(),
stuckCommitDurationMillis,
- () -> Collections.unmodifiableCollection(computationMap.values()),
+ computationStateCache::getAllPresentComputations,
sampler,
metricTrackingWindmillServer::refreshActiveWork,
executorSupplier.apply("RefreshWork"));
+
+ WorkerStatusPages workerStatusPages =
+ WorkerStatusPages.create(DEFAULT_STATUS_PORT, memoryMonitor);
+ StreamingWorkerStatusPages.Builder statusPagesBuilder =
+ StreamingWorkerStatusPages.builder()
+ .setClock(clock)
+ .setClientId(clientId)
+ .setIsRunning(running)
+ .setStatusPages(workerStatusPages)
+ .setStateCache(stateCache)
+ .setComputationStateCache(computationStateCache)
+ .setCurrentActiveCommitBytes(workCommitter::currentActiveCommitBytes)
+ .setGetDataStatusProvider(metricTrackingWindmillServer::printHtml)
+ .setWorkUnitExecutor(workUnitExecutor);
+
+ this.statusPages =
+ windmillServiceEnabled
+ ? statusPagesBuilder
+ .setDebugCapture(
+ new DebugCapture.Manager(options, workerStatusPages.getDebugCapturePages()))
+ .setChannelzServlet(new ChannelzServlet(CHANNELZ_PATH, options, windmillServer))
+ .setWindmillStreamFactory(windmillStreamFactory)
+ .build()
+ : statusPagesBuilder.build();
+
this.workerStatusReporter = workerStatusReporter;
this.failureTracker = failureTracker;
this.workFailureProcessor = workFailureProcessor;
@@ -387,24 +377,56 @@ public class StreamingDataflowWorker {
LOG.debug("WindmillServiceEndpoint: {}", options.getWindmillServiceEndpoint());
LOG.debug("WindmillServicePort: {}", options.getWindmillServicePort());
LOG.debug("LocalWindmillHostport: {}", options.getLocalWindmillHostport());
- LOG.debug("maxWorkItemCommitBytes: {}", maxWorkItemCommitBytes);
+ LOG.debug("maxWorkItemCommitBytes: {}", maxWorkItemCommitBytes.get());
}
public static StreamingDataflowWorker fromOptions(DataflowWorkerHarnessOptions options) {
long clientId = clientIdGenerator.nextLong();
MemoryMonitor memoryMonitor = MemoryMonitor.fromOptions(options);
- ConcurrentMap<String, ComputationState> computationMap = new ConcurrentHashMap<>();
ConcurrentMap<String, StageInfo> stageInfo = new ConcurrentHashMap<>();
StreamingCounters streamingCounters = StreamingCounters.create();
-
+ WorkUnitClient dataflowServiceClient = new DataflowWorkUnitClient(options, LOG);
+ BoundedQueueExecutor workExecutor = createWorkUnitExecutor(options);
+ AtomicInteger maxWorkItemCommitBytes = new AtomicInteger(Integer.MAX_VALUE);
+ WindmillStateCache windmillStateCache =
+ WindmillStateCache.ofSizeMbs(options.getWorkerCacheMb());
+ Function<String, ScheduledExecutorService> executorSupplier =
+ threadName ->
+ Executors.newSingleThreadScheduledExecutor(
+ new ThreadFactoryBuilder().setNameFormat(threadName).build());
GrpcWindmillStreamFactory windmillStreamFactory =
createWindmillStreamFactory(options, clientId);
- WindmillServerStub windmillServer =
- createWindmillServerStub(
+ GrpcDispatcherClient dispatcherClient = GrpcDispatcherClient.create(createStubFactory(options));
+
+ // If ComputationConfig.Fetcher is the Streaming Appliance implementation, WindmillServerStub
+ // can be created without a heartbeat response processor, as appliance does not send heartbeats.
+ Pair<ComputationConfig.Fetcher, Optional<WindmillServerStub>> configFetcherAndWindmillClient =
+ createConfigFetcherAndWindmillClient(
options,
- windmillStreamFactory,
- new WorkHeartbeatResponseProcessor(
- computationId -> Optional.ofNullable(computationMap.get(computationId))));
+ dataflowServiceClient,
+ dispatcherClient,
+ maxWorkItemCommitBytes,
+ windmillStreamFactory);
+
+ ComputationStateCache computationStateCache =
+ ComputationStateCache.create(
+ configFetcherAndWindmillClient.getLeft(),
+ workExecutor,
+ windmillStateCache::forComputation,
+ ID_GENERATOR);
+
+ // If WindmillServerStub is not present, it is a Streaming Engine job. We now have all the
+ // components created to initialize the GrpcWindmillServer.
+ WindmillServerStub windmillServer =
+ configFetcherAndWindmillClient
+ .getRight()
+ .orElseGet(
+ () ->
+ GrpcWindmillServer.create(
+ options,
+ windmillStreamFactory,
+ dispatcherClient,
+ new WorkHeartbeatResponseProcessor(computationStateCache::get)));
FailureTracker failureTracker =
options.isEnableStreamingEngine()
@@ -415,8 +437,6 @@ public class StreamingDataflowWorker {
options.getMaxStackTraceDepthToReport(),
windmillServer::reportStats);
- WorkUnitClient dataflowServiceClient = new DataflowWorkUnitClient(options, LOG);
- BoundedQueueExecutor workExecutor = createWorkUnitExecutor(options);
Supplier<Instant> clock = Instant::now;
WorkFailureProcessor workFailureProcessor =
WorkFailureProcessor.create(
@@ -438,12 +458,12 @@ public class StreamingDataflowWorker {
return new StreamingDataflowWorker(
windmillServer,
clientId,
- computationMap,
+ configFetcherAndWindmillClient.getLeft(),
+ computationStateCache,
stageInfo,
WindmillStateCache.ofSizeMbs(options.getWorkerCacheMb()),
workExecutor,
IntrinsicMapTaskExecutorFactory.defaultFactory(),
- dataflowServiceClient,
options,
new HotKeyLogger(),
clock,
@@ -452,14 +472,42 @@ public class StreamingDataflowWorker {
workFailureProcessor,
streamingCounters,
memoryMonitor,
- threadName ->
- Executors.newSingleThreadScheduledExecutor(
- new ThreadFactoryBuilder().setNameFormat(threadName).build()));
+ maxWorkItemCommitBytes,
+ windmillStreamFactory,
+ executorSupplier);
+ }
+
+ private static Pair<ComputationConfig.Fetcher, Optional<WindmillServerStub>>
+ createConfigFetcherAndWindmillClient(
+ DataflowWorkerHarnessOptions options,
+ WorkUnitClient dataflowServiceClient,
+ GrpcDispatcherClient dispatcherClient,
+ AtomicInteger maxWorkItemCommitBytes,
+ GrpcWindmillStreamFactory windmillStreamFactory) {
+ ComputationConfig.Fetcher configFetcher;
+ @Nullable WindmillServerStub windmillServer = null;
+ if (options.isEnableStreamingEngine()) {
+ configFetcher =
+ StreamingEngineComputationConfigFetcher.create(
+ options.getGlobalConfigRefreshPeriod().getMillis(),
+ dataflowServiceClient,
+ config ->
+ onPipelineConfig(
+ config,
+ dispatcherClient::consumeWindmillDispatcherEndpoints,
+ maxWorkItemCommitBytes));
+ } else {
+ windmillServer =
+ createWindmillServerStub(options, windmillStreamFactory, dispatcherClient, ignored -> {});
+ configFetcher = new StreamingApplianceComputationConfigFetcher(windmillServer::getConfig);
+ }
+
+ return Pair.of(configFetcher, Optional.ofNullable(windmillServer));
}
@VisibleForTesting
static StreamingDataflowWorker forTesting(
- ConcurrentMap<String, ComputationState> computationMap,
+ Map<String, String> prePopulatedStateNameMappings,
WindmillServerStub windmillServer,
List<MapTask> mapTasks,
DataflowMapTaskExecutorFactory mapTaskExecutorFactory,
@@ -469,12 +517,39 @@ public class StreamingDataflowWorker {
HotKeyLogger hotKeyLogger,
Supplier<Instant> clock,
Function<String, ScheduledExecutorService> executorSupplier,
- int localRetryTimeoutMs) {
+ int localRetryTimeoutMs,
+ int maxWorkItemCommitBytesOverrides) {
ConcurrentMap<String, StageInfo> stageInfo = new ConcurrentHashMap<>();
+ AtomicInteger maxWorkItemCommitBytes = new AtomicInteger(maxWorkItemCommitBytesOverrides);
BoundedQueueExecutor workExecutor = createWorkUnitExecutor(options);
WindmillStateCache stateCache = WindmillStateCache.ofSizeMbs(options.getWorkerCacheMb());
- computationMap.putAll(
- createComputationMapForTesting(mapTasks, workExecutor, stateCache::forComputation));
+ ComputationConfig.Fetcher configFetcher =
+ options.isEnableStreamingEngine()
+ ? StreamingEngineComputationConfigFetcher.forTesting(
+ /* hasReceivedGlobalConfig= */ true,
+ options.getGlobalConfigRefreshPeriod().getMillis(),
+ workUnitClient,
+ executorSupplier,
+ config ->
+ onPipelineConfig(
+ config,
+ windmillServer::setWindmillServiceEndpoints,
+ maxWorkItemCommitBytes))
+ : new StreamingApplianceComputationConfigFetcher(windmillServer::getConfig);
+ ConcurrentMap<String, String> stateNameMap =
+ new ConcurrentHashMap<>(prePopulatedStateNameMappings);
+ ComputationStateCache computationStateCache =
+ ComputationStateCache.forTesting(
+ configFetcher, workExecutor, stateCache::forComputation, ID_GENERATOR, stateNameMap);
+ computationStateCache.loadCacheForTesting(
+ mapTasks,
+ mapTask ->
+ new ComputationState(
+ mapTask.getStageName(),
+ mapTask,
+ workExecutor,
+ stateNameMap,
+ stateCache.forComputation(mapTask.getStageName())));
MemoryMonitor memoryMonitor = MemoryMonitor.fromOptions(options);
StreamingCounters streamingCounters = StreamingCounters.create();
FailureTracker failureTracker =
@@ -505,15 +580,16 @@ public class StreamingDataflowWorker {
executorSupplier,
options.getWindmillHarnessUpdateReportingPeriod().getMillis(),
options.getPerWorkerMetricsUpdateReportingPeriodMillis());
+
return new StreamingDataflowWorker(
windmillServer,
1L,
- computationMap,
+ configFetcher,
+ computationStateCache,
stageInfo,
stateCache,
workExecutor,
mapTaskExecutorFactory,
- workUnitClient,
options,
hotKeyLogger,
clock,
@@ -522,9 +598,25 @@ public class StreamingDataflowWorker {
workFailureProcessor,
streamingCounters,
memoryMonitor,
+ maxWorkItemCommitBytes,
+ createWindmillStreamFactory(options, 1),
executorSupplier);
}
+ private static void onPipelineConfig(
+ StreamingEnginePipelineConfig config,
+ Consumer<ImmutableSet<HostAndPort>> consumeWindmillServiceEndpoints,
+ AtomicInteger maxWorkItemCommitBytes) {
+ if (config.maxWorkItemCommitBytes() != maxWorkItemCommitBytes.get()) {
+ LOG.info("Setting maxWorkItemCommitBytes to {}", maxWorkItemCommitBytes);
+ maxWorkItemCommitBytes.set((int) config.maxWorkItemCommitBytes());
+ }
+
+ if (!config.windmillServiceEndpoints().isEmpty()) {
+ consumeWindmillServiceEndpoints.accept(config.windmillServiceEndpoints());
+ }
+ }
+
private static GrpcWindmillStreamFactory createWindmillStreamFactory(
DataflowWorkerHarnessOptions options, long clientId) {
Duration maxBackoff =
@@ -550,26 +642,6 @@ public class StreamingDataflowWorker {
workerStatusReporter.reportPeriodicWorkerUpdates();
}
- private static ConcurrentMap<String, ComputationState> createComputationMapForTesting(
- List<MapTask> mapTasks,
- BoundedQueueExecutor workUnitExecutor,
- Function<String, WindmillStateCache.ForComputation> forComputationStateCacheFactory) {
- return mapTasks.stream()
- .map(fixMultiOutputInfos)
- .map(
- mapTask -> {
- LOG.info("Adding config for {}: {}", mapTask.getSystemName(), mapTask);
- String computationId = mapTask.getStageName();
- return new ComputationState(
- computationId,
- mapTask,
- workUnitExecutor,
- ImmutableMap.of(),
- forComputationStateCacheFactory.apply(computationId));
- })
- .collect(toConcurrentMap(ComputationState::getComputationId, Function.identity()));
- }
-
private static BoundedQueueExecutor createWorkUnitExecutor(DataflowWorkerHarnessOptions options) {
return new BoundedQueueExecutor(
chooseMaxThreads(options),
@@ -580,10 +652,6 @@ public class StreamingDataflowWorker {
new ThreadFactoryBuilder().setNameFormat("DataflowWorkUnits-%d").setDaemon(true).build());
}
- private static MapTask parseMapTask(String input) throws IOException {
- return Transport.getJsonFactory().fromString(input, MapTask.class);
- }
-
public static void main(String[] args) throws Exception {
JvmInitializers.runOnStartup();
@@ -624,22 +692,38 @@ public class StreamingDataflowWorker {
private static WindmillServerStub createWindmillServerStub(
DataflowWorkerHarnessOptions options,
GrpcWindmillStreamFactory windmillStreamFactory,
+ GrpcDispatcherClient dispatcherClient,
Consumer<List<Windmill.ComputationHeartbeatResponse>> processHeartbeatResponses) {
if (options.getWindmillServiceEndpoint() != null
|| options.isEnableStreamingEngine()
|| options.getLocalWindmillHostport().startsWith("grpc:")) {
- try {
- windmillStreamFactory.scheduleHealthChecks(
- options.getWindmillServiceStreamingRpcHealthCheckPeriodMs());
- return GrpcWindmillServer.create(options, windmillStreamFactory, processHeartbeatResponses);
- } catch (IOException e) {
- throw new RuntimeException("Failed to create GrpcWindmillServer: ", e);
- }
+ windmillStreamFactory.scheduleHealthChecks(
+ options.getWindmillServiceStreamingRpcHealthCheckPeriodMs());
+ return GrpcWindmillServer.create(
+ options, windmillStreamFactory, dispatcherClient, processHeartbeatResponses);
} else {
return new JniWindmillApplianceServer(options.getLocalWindmillHostport());
}
}
+ private static ChannelCachingStubFactory createStubFactory(
+ DataflowWorkerHarnessOptions workerOptions) {
+ Function<WindmillServiceAddress, ManagedChannel> channelFactory =
+ serviceAddress ->
+ remoteChannel(
+ serviceAddress, workerOptions.getWindmillServiceRpcChannelAliveTimeoutSec());
+ ChannelCache channelCache =
+ ChannelCache.create(
+ serviceAddress ->
+ // IsolationChannel will create and manage separate RPC channels to the same
+ // serviceAddress via calling the channelFactory, else just directly return the
+ // RPC channel.
+ workerOptions.getUseWindmillIsolatedChannels()
+ ? IsolationChannel.create(() -> channelFactory.apply(serviceAddress))
+ : channelFactory.apply(serviceAddress));
+ return ChannelCachingRemoteStubFactory.create(workerOptions.getGcpCredential(), channelCache);
+ }
+
private static void sleep(int millis) {
Uninterruptibles.sleepUninterruptibly(millis, TimeUnit.MILLISECONDS);
}
@@ -692,18 +776,6 @@ public class StreamingDataflowWorker {
BigQuerySinkMetrics.setSupportStreamingInsertsMetrics(true);
}
- void addStateNameMappings(Map<String, String> nameMap) {
- stateNameMap.putAll(nameMap);
- }
-
- @VisibleForTesting
- public void setMaxWorkItemCommitBytes(int maxWorkItemCommitBytes) {
- if (maxWorkItemCommitBytes != this.maxWorkItemCommitBytes) {
- LOG.info("Setting maxWorkItemCommitBytes to {}", maxWorkItemCommitBytes);
- }
- this.maxWorkItemCommitBytes = maxWorkItemCommitBytes;
- }
-
@VisibleForTesting
public boolean workExecutorIsEmpty() {
return workUnitExecutor.executorQueueIsEmpty();
@@ -714,104 +786,37 @@ public class StreamingDataflowWorker {
return workCommitter.parallelism();
}
+ @VisibleForTesting
+ ComputationStateCache getComputationStateCache() {
+ return computationStateCache;
+ }
+
@SuppressWarnings("FutureReturnValueIgnored")
public void start() {
running.set(true);
- if (windmillServiceEnabled) {
- // Schedule the background getConfig thread. Blocks until windmillServer stub is ready.
- schedulePeriodicGlobalConfigRequests();
- }
+ configFetcher.start();
memoryMonitorThread.start();
dispatchThread.start();
sampler.start();
- if (options.getPeriodicStatusPageOutputDirectory() != null) {
- ScheduledExecutorService statusPageTimer = executorSupplier.apply("DumpStatusPages");
- statusPageTimer.scheduleWithFixedDelay(
- () -> {
- Collection<Capturable> pages = statusPages.getDebugCapturePages();
- if (pages.isEmpty()) {
- LOG.warn("No captured status pages.");
- }
- long timestamp = clock.get().getMillis();
- for (Capturable page : pages) {
- PrintWriter writer = null;
- try {
- File outputFile =
- new File(
- options.getPeriodicStatusPageOutputDirectory(),
- ("StreamingDataflowWorker"
- + options.getWorkerId()
- + "_"
- + page.pageName()
- + timestamp
- + ".html")
- .replaceAll("/", "_"));
- writer = new PrintWriter(outputFile, UTF_8.name());
- page.captureData(writer);
- } catch (IOException e) {
- LOG.warn("Error dumping status page.", e);
- } finally {
- if (writer != null) {
- writer.close();
- }
- }
- }
- },
- 60,
- 60,
- TimeUnit.SECONDS);
- scheduledExecutors.add(statusPageTimer);
- }
workCommitter.start();
workerStatusReporter.start();
activeWorkRefresher.start();
}
- public void startStatusPages() {
- if (debugCaptureManager != null) {
- debugCaptureManager.start();
- }
-
- if (windmillServiceEnabled) {
- ChannelzServlet channelzServlet = new ChannelzServlet(CHANNELZ_PATH, options, windmillServer);
- statusPages.addServlet(channelzServlet);
- statusPages.addCapturePage(channelzServlet);
- }
-
- statusPages.addServlet(stateCache.statusServlet());
- statusPages.addServlet(new SpecsServlet());
-
- statusPages.addStatusDataProvider("harness", "Harness", new HarnessDataProvider());
- statusPages.addStatusDataProvider("metrics", "Metrics", new MetricsDataProvider());
- statusPages.addStatusDataProvider(
- "exception", "Last Exception", new LastExceptionDataProvider());
- statusPages.addStatusDataProvider("cache", "State Cache", stateCache);
- statusPages.addStatusDataProvider("streaming", "Streaming Rpcs", windmillServer);
-
- statusPages.start();
+ /** Starts the status page server for debugging. May be omitted for lighter weight testing. */
+ private void startStatusPages() {
+ statusPages.start(options);
}
public void stop() {
try {
- for (ScheduledExecutorService timer : scheduledExecutors) {
- if (timer != null) {
- timer.shutdown();
- }
- }
- for (ScheduledExecutorService timer : scheduledExecutors) {
- if (timer != null) {
- timer.awaitTermination(300, TimeUnit.SECONDS);
- }
- }
+ configFetcher.stop();
activeWorkRefresher.stop();
statusPages.stop();
- if (debugCaptureManager != null) {
- debugCaptureManager.stop();
- }
running.set(false);
dispatchThread.interrupt();
dispatchThread.join();
@@ -820,55 +825,13 @@ public class StreamingDataflowWorker {
memoryMonitor.stop();
memoryMonitorThread.join();
workUnitExecutor.shutdown();
- for (ComputationState state : computationMap.values()) {
- state.close();
- }
+
+ computationStateCache.closeAndInvalidateAll();
workerStatusReporter.stop();
} catch (Exception e) {
LOG.warn("Exception while shutting down: ", e);
}
- setIsDone();
- }
-
- // null is the only value of type Void, but findbugs thinks
- // it violates the contract of CompletableFuture.complete
- @SuppressFBWarnings("NP_NONNULL_PARAM_VIOLATION")
- private void setIsDone() {
- isDoneFuture.complete(null);
- }
-
- private synchronized void addComputation(
- String computationId,
- MapTask originalMapTask,
- Map<String, String> transformUserNameToStateFamily) {
- // Map task instances are shared amongst multiple threads during computation hence
- // we fix the map task before we add a new computation state that would reference it.
- MapTask mapTask = fixMultiOutputInfos.apply(originalMapTask);
- if (!computationMap.containsKey(computationId)) {
- LOG.info("Adding config for {}: {}", computationId, mapTask);
- computationMap.put(
- computationId,
- new ComputationState(
- computationId,
- mapTask,
- workUnitExecutor,
- transformUserNameToStateFamily,
- stateCache.forComputation(computationId)));
- }
- }
-
- /**
- * If the computation is not yet known about, configuration for it will be fetched. This can still
- * return null if there is no configuration fetched for the computation.
- */
- private ComputationState getComputationState(String computationId) {
- ComputationState state = computationMap.get(computationId);
- if (state == null) {
- getConfig(computationId);
- state = computationMap.get(computationId);
- }
- return state;
}
private void dispatchLoop() {
@@ -891,15 +854,12 @@ public class StreamingDataflowWorker {
} while (running.get());
for (final Windmill.ComputationWorkItems computationWork : workResponse.getWorkList()) {
final String computationId = computationWork.getComputationId();
- final ComputationState computationState = getComputationState(computationId);
- if (computationState == null) {
- LOG.warn(
- "Received work for unknown computation: {}. Known computations are {}",
- computationId,
- computationMap.keySet());
+ Optional<ComputationState> maybeComputationState = computationStateCache.get(computationId);
+ if (!maybeComputationState.isPresent()) {
continue;
}
+ final ComputationState computationState = maybeComputationState.get();
final Instant inputDataWatermark =
WindmillTimeUtils.windmillToHarnessWatermark(computationWork.getInputDataWatermark());
Preconditions.checkNotNull(inputDataWatermark);
@@ -931,15 +891,19 @@ public class StreamingDataflowWorker {
Instant inputDataWatermark,
Instant synchronizedProcessingTime,
Windmill.WorkItem workItem,
- Collection<LatencyAttribution> getWorkStreamLatencies) -> {
- memoryMonitor.waitForResources("GetWork");
- scheduleWorkItem(
- getComputationState(computation),
- inputDataWatermark,
- synchronizedProcessingTime,
- workItem,
- getWorkStreamLatencies);
- });
+ Collection<LatencyAttribution> getWorkStreamLatencies) ->
+ computationStateCache
+ .get(computation)
+ .ifPresent(
+ computationState -> {
+ memoryMonitor.waitForResources("GetWork");
+ scheduleWorkItem(
+ computationState,
+ inputDataWatermark,
+ synchronizedProcessingTime,
+ workItem,
+ getWorkStreamLatencies);
+ }));
try {
// Reconnect every now and again to enable better load balancing.
// If at any point the server closes the stream, we will reconnect immediately; otherwise
@@ -1073,14 +1037,14 @@ public class StreamingDataflowWorker {
stageInfoMap.computeIfAbsent(
mapTask.getStageName(), s -> StageInfo.create(s, mapTask.getSystemName()));
- ExecutionState executionState = null;
+ @Nullable ExecutionState executionState = null;
String counterName = "dataflow_source_bytes_processed-" + mapTask.getSystemName();
try {
if (work.isFailed()) {
throw new WorkItemCancelledException(workItem.getShardingKey());
}
- executionState = computationState.getExecutionStateQueue().poll();
+ executionState = computationState.acquireExecutionState().orElse(null);
if (executionState == null) {
MutableNetwork<Node, Edge> mapTaskNetwork = mapTaskToNetwork.apply(mapTask);
if (LOG.isDebugEnabled()) {
@@ -1114,9 +1078,7 @@ public class StreamingDataflowWorker {
streamingCounters.pendingDeltaCounters(),
computationId,
readerCache,
- !computationState.getTransformUserNameToStateFamily().isEmpty()
- ? computationState.getTransformUserNameToStateFamily()
- : stateNameMap,
+ computationState.getTransformUserNameToStateFamily(),
stateCache.forComputation(computationId),
stageInfo.metricsContainerRegistry(),
executionStateTracker,
@@ -1131,7 +1093,7 @@ public class StreamingDataflowWorker {
sinkRegistry,
context,
streamingCounters.pendingDeltaCounters(),
- idGenerator);
+ ID_GENERATOR);
ReadOperation readOperation = mapTaskExecutor.getReadOperation();
// Disable progress updates since its results are unused for streaming
// and involves starting a thread.
@@ -1263,7 +1225,7 @@ public class StreamingDataflowWorker {
commitCallbacks.putAll(executionState.context().flushState());
// Release the execution state for another thread to use.
- computationState.getExecutionStateQueue().offer(executionState);
+ computationState.releaseExecutionState(executionState);
executionState = null;
// Add the output to the commit queue.
@@ -1272,7 +1234,7 @@ public class StreamingDataflowWorker {
work.getLatencyAttributions(false, work.getLatencyTrackingId(), sampler));
WorkItemCommitRequest commitRequest = outputBuilder.build();
- int byteLimit = maxWorkItemCommitBytes;
+ int byteLimit = maxWorkItemCommitBytes.get();
int commitSize = commitRequest.getSerializedSize();
int estimatedCommitSize = commitSize < 0 ? Integer.MAX_VALUE : commitSize;
@@ -1375,7 +1337,8 @@ public class StreamingDataflowWorker {
.invalidate(completeCommit.shardedKey());
}
- Optional.ofNullable(computationMap.get(completeCommit.computationId()))
+ computationStateCache
+ .getIfPresent(completeCommit.computationId())
.ifPresent(
state ->
state.completeWorkAndScheduleNextWorkForKey(
@@ -1391,183 +1354,6 @@ public class StreamingDataflowWorker {
.build());
}
- private void getConfigFromWindmill(String computation) {
- Windmill.GetConfigRequest request =
- Windmill.GetConfigRequest.newBuilder().addComputations(computation).build();
-
- Windmill.GetConfigResponse response = windmillServer.getConfig(request);
- // The max work item commit bytes should be modified to be dynamic once it is available in
- // the request.
- for (Windmill.GetConfigResponse.SystemNameToComputationIdMapEntry entry :
- response.getSystemNameToComputationIdMapList()) {
- systemNameToComputationIdMap.put(entry.getSystemName(), entry.getComputationId());
- }
-
- // Outer keys are computation ids. Outer values are map from transform username to state family.
- Map<String, Map<String, String>> transformUserNameToStateFamilyByComputationId =
- new HashMap<>();
- for (Windmill.GetConfigResponse.ComputationConfigMapEntry computationConfig :
- response.getComputationConfigMapList()) {
- Map<String, String> transformUserNameToStateFamily =
- transformUserNameToStateFamilyByComputationId.computeIfAbsent(
- computationConfig.getComputationId(), k -> new HashMap<>());
- for (Windmill.ComputationConfig.TransformUserNameToStateFamilyEntry entry :
- computationConfig.getComputationConfig().getTransformUserNameToStateFamilyList()) {
- transformUserNameToStateFamily.put(entry.getTransformUserName(), entry.getStateFamily());
- }
- }
-
- for (String serializedMapTask : response.getCloudWorksList()) {
- try {
- MapTask mapTask = parseMapTask(serializedMapTask);
- String computationId =
- systemNameToComputationIdMap.containsKey(mapTask.getSystemName())
- ? systemNameToComputationIdMap.get(mapTask.getSystemName())
- : mapTask.getSystemName();
- addComputation(
- computationId,
- mapTask,
- transformUserNameToStateFamilyByComputationId.getOrDefault(
- computationId, ImmutableMap.of()));
- } catch (IOException e) {
- LOG.warn("Parsing MapTask failed: {}", serializedMapTask);
- LOG.warn("Error: ", e);
- }
- }
- for (Windmill.GetConfigResponse.NameMapEntry entry : response.getNameMapList()) {
- stateNameMap.put(entry.getUserName(), entry.getSystemName());
- }
- }
-
- /**
- * Sends a request to get configuration from Dataflow, either for a specific computation (if
- * computation is not null) or global configuration (if computation is null).
- *
- * @throws IOException if the RPC fails.
- */
- private void getConfigFromDataflowService(@Nullable String computation) throws IOException {
- Optional<WorkItem> workItem =
- computation != null
- ? workUnitClient.getStreamingConfigWorkItem(computation)
- : workUnitClient.getGlobalStreamingConfigWorkItem();
-
- if (!workItem.isPresent()) {
- return;
- }
- StreamingConfigTask config = workItem.get().getStreamingConfigTask();
- Preconditions.checkState(config != null);
- if (config.getUserStepToStateFamilyNameMap() != null) {
- stateNameMap.putAll(config.getUserStepToStateFamilyNameMap());
- }
- if (computation == null) {
- if (config.getMaxWorkItemCommitBytes() != null
- && config.getMaxWorkItemCommitBytes() > 0
- && config.getMaxWorkItemCommitBytes() <= Integer.MAX_VALUE) {
- setMaxWorkItemCommitBytes(config.getMaxWorkItemCommitBytes().intValue());
- } else {
- setMaxWorkItemCommitBytes(180 << 20);
- }
- }
- List<StreamingComputationConfig> configs = config.getStreamingComputationConfigs();
- if (configs != null) {
- for (StreamingComputationConfig computationConfig : configs) {
- MapTask mapTask = new MapTask();
- mapTask.setSystemName(computationConfig.getSystemName());
- mapTask.setStageName(computationConfig.getStageName());
- mapTask.setInstructions(computationConfig.getInstructions());
- addComputation(
- computationConfig.getComputationId(),
- mapTask,
- Optional.ofNullable(computationConfig.getTransformUserNameToStateFamily())
- .orElseGet(ImmutableMap::of));
- }
- }
-
- if (config.getWindmillServiceEndpoint() != null
- && !config.getWindmillServiceEndpoint().isEmpty()) {
- int port = 443;
- if (config.getWindmillServicePort() != null && config.getWindmillServicePort() != 0) {
- port = config.getWindmillServicePort().intValue();
- }
- HashSet<HostAndPort> endpoints = new HashSet<>();
- for (String endpoint : Splitter.on(',').split(config.getWindmillServiceEndpoint())) {
- endpoints.add(HostAndPort.fromString(endpoint).withDefaultPort(port));
- }
- windmillServer.setWindmillServiceEndpoints(endpoints);
- }
- }
-
- /**
- * Schedules a background thread that periodically sends getConfig requests to Dataflow Service to
- * obtain the windmill service endpoint. Blocks until the windmillServer is ready.
- */
- @SuppressWarnings("FutureReturnValueIgnored")
- private void schedulePeriodicGlobalConfigRequests() {
- Preconditions.checkState(windmillServiceEnabled);
- if (!windmillServer.isReady()) {
- // Get the initial global configuration. This will initialize the windmillServer stub.
- while (true) {
- LOG.info("Sending request to get global configuration for this worker");
- getGlobalConfig();
- if (windmillServer.isReady()) {
- break;
- }
- LOG.info("windmillServerStub is not ready yet, will retry in 5 seconds");
- sleep(5000);
- }
- }
- LOG.info("windmillServerStub is now ready");
-
- // Now start a thread that periodically refreshes the windmill service endpoint.
- ScheduledExecutorService configRefreshTimer =
- executorSupplier.apply("GlobalConfigRefreshTimer");
- configRefreshTimer.scheduleWithFixedDelay(
- this::getGlobalConfig,
- 0,
- options.getGlobalConfigRefreshPeriod().getMillis(),
- TimeUnit.MILLISECONDS);
- scheduledExecutors.add(configRefreshTimer);
- }
-
- private void getGlobalConfig() {
- // No need to pass a computation since we are only interested in the global config.
- getConfig(null);
- }
-
- // Attempts to populate computationMap with an entry for the computation. Upon
- // error or shutdown computationMap may remain unchanged.
- private void getConfig(String computation) {
- BackOff backoff =
- FluentBackoff.DEFAULT
- .withInitialBackoff(Duration.millis(100))
- .withMaxBackoff(Duration.standardMinutes(1))
- .withMaxCumulativeBackoff(Duration.standardMinutes(5))
- .backoff();
- while (running.get()) {
- try {
- if (windmillServiceEnabled) {
- getConfigFromDataflowService(computation);
- } else {
- getConfigFromWindmill(computation);
- }
- return;
- } catch (IllegalArgumentException | IOException e) {
- LOG.warn("Error fetching config: ", e);
- try {
- if (!BackOffUtils.next(Sleeper.DEFAULT, backoff)) {
- return;
- }
- } catch (IOException ioe) {
- LOG.warn("Error backing off, will not retry: ", ioe);
- return;
- } catch (InterruptedException ie) {
- Thread.currentThread().interrupt();
- return;
- }
- }
- }
- }
-
@VisibleForTesting
public Iterable<CounterUpdate> buildCounters() {
return Iterables.concat(
@@ -1578,70 +1364,4 @@ public class StreamingDataflowWorker {
.pendingCumulativeCounters()
.extractUpdates(false, DataflowCounterUpdateExtractor.INSTANCE));
}
-
- private class HarnessDataProvider implements StatusDataProvider {
- @Override
- public void appendSummaryHtml(PrintWriter writer) {
- writer.println("Running: " + running.get() + "<br>");
- writer.println("ID: " + clientId + "<br>");
- }
- }
-
- private class SpecsServlet extends BaseStatusServlet {
-
- public SpecsServlet() {
- super("/specs");
- }
-
- @Override
- public void doGet(HttpServletRequest request, HttpServletResponse response) throws IOException {
- PrintWriter writer = response.getWriter();
- writer.println("<h1>Specs</h1>");
- for (Map.Entry<String, ComputationState> entry : computationMap.entrySet()) {
- writer.println("<h3>" + entry.getKey() + "</h3>");
- writer.print("<script>document.write(JSON.stringify(");
- writer.print(entry.getValue().getMapTask().toString());
- writer.println(", null, \"  \").replace(/\\n/g, \"<br>\"))</script>");
- }
- }
- }
-
- private class MetricsDataProvider implements StatusDataProvider {
-
- @Override
- public void appendSummaryHtml(PrintWriter writer) {
- writer.println(workUnitExecutor.summaryHtml());
-
- writer.print("Active commit: ");
- appendHumanizedBytes(workCommitter.currentActiveCommitBytes(), writer);
- writer.println("<br>");
-
- metricTrackingWindmillServer.printHtml(writer);
-
- writer.println("<br>");
-
- writer.println("Active Keys: <br>");
- for (Map.Entry<String, ComputationState> computationEntry : computationMap.entrySet()) {
- writer.print(computationEntry.getKey());
- writer.print(":<br>");
- computationEntry.getValue().printActiveWork(writer);
- writer.println("<br>");
- }
- }
-
- private void appendHumanizedBytes(long bytes, PrintWriter writer) {
- if (bytes < (4 << 10)) {
- writer.print(bytes);
- writer.print("B");
- } else if (bytes < (4 << 20)) {
- writer.print("~");
- writer.print(bytes >> 10);
- writer.print("KB");
- } else {
- writer.print("~");
- writer.print(bytes >> 20);
- writer.print("MB");
- }
- }
- }
}
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/status/WorkerStatusPages.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/status/WorkerStatusPages.java
index aaa6cd6d6d3..ef783a00f8e 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/status/WorkerStatusPages.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/status/WorkerStatusPages.java
@@ -80,6 +80,10 @@ public class WorkerStatusPages {
return new WorkerStatusPages(new Server(statusPort), memoryMonitor, healthyIndicator);
}
+ public static WorkerStatusPages create(int defaultStatusPort, MemoryMonitor memoryMonitor) {
+ return create(defaultStatusPort, memoryMonitor, () -> true);
+ }
+
/** Start the server. */
public void start() {
if (statusServer.isStarted()) {
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationState.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationState.java
index 33ef4950f9a..f9466f25577 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationState.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationState.java
@@ -20,12 +20,14 @@ package org.apache.beam.runners.dataflow.worker.streaming;
import com.google.api.services.dataflow.model.MapTask;
import java.io.PrintWriter;
import java.util.Map;
+import java.util.Optional;
import java.util.concurrent.ConcurrentLinkedQueue;
import javax.annotation.Nullable;
import org.apache.beam.runners.dataflow.worker.DataflowExecutionStateSampler;
import org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.HeartbeatRequest;
import org.apache.beam.runners.dataflow.worker.windmill.state.WindmillStateCache;
+import org.apache.beam.runners.dataflow.worker.windmill.work.budget.GetWorkBudget;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
@@ -38,7 +40,7 @@ import org.joda.time.Instant;
* <p>This class is synchronized, but only used from the dispatch and commit threads, so should not
* be heavily contended. Still, blocking work should not be done by it.
*/
-public class ComputationState implements AutoCloseable {
+public class ComputationState {
private final String computationId;
private final MapTask mapTask;
private final ImmutableMap<String, String> transformUserNameToStateFamily;
@@ -74,8 +76,20 @@ public class ComputationState implements AutoCloseable {
return transformUserNameToStateFamily;
}
- public ConcurrentLinkedQueue<ExecutionState> getExecutionStateQueue() {
- return executionStateQueue;
+ /**
+ * Cache the {@link ExecutionState} so that it can be re-used in future {@link
+ * #acquireExecutionState()} calls.
+ */
+ public void releaseExecutionState(ExecutionState executionState) {
+ executionStateQueue.offer(executionState);
+ }
+
+ /**
+ * Returns {@link ExecutionState} that was previously offered in {@link
+ * #releaseExecutionState(ExecutionState)} or {@link Optional#empty()} if one does not exist.
+ */
+ public Optional<ExecutionState> acquireExecutionState() {
+ return Optional.ofNullable(executionStateQueue.poll());
}
/**
@@ -134,12 +148,15 @@ public class ComputationState implements AutoCloseable {
return activeWorkState.getKeyHeartbeats(refreshDeadline, sampler);
}
+ public GetWorkBudget getActiveWorkBudget() {
+ return activeWorkState.currentActiveWorkBudget();
+ }
+
public void printActiveWork(PrintWriter writer) {
activeWorkState.printActiveWork(writer, Instant.now());
}
- @Override
- public void close() throws Exception {
+ public final void close() {
@Nullable ExecutionState executionState;
while ((executionState = executionStateQueue.poll()) != null) {
executionState.workExecutor().close();
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationStateCache.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationStateCache.java
new file mode 100644
index 00000000000..33e28075487
--- /dev/null
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationStateCache.java
@@ -0,0 +1,218 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.streaming;
+
+import static java.util.stream.Collectors.toConcurrentMap;
+import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList.toImmutableList;
+
+import com.google.api.services.dataflow.model.MapTask;
+import java.io.PrintWriter;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.ExecutionException;
+import java.util.function.Function;
+import javax.annotation.concurrent.ThreadSafe;
+import org.apache.beam.runners.dataflow.worker.apiary.FixMultiOutputInfosOnParDoInstructions;
+import org.apache.beam.runners.dataflow.worker.status.StatusDataProvider;
+import org.apache.beam.runners.dataflow.worker.streaming.config.ComputationConfig;
+import org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor;
+import org.apache.beam.runners.dataflow.worker.windmill.state.WindmillStateCache;
+import org.apache.beam.runners.dataflow.worker.windmill.work.budget.GetWorkBudget;
+import org.apache.beam.sdk.annotations.Internal;
+import org.apache.beam.sdk.fn.IdGenerator;
+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.cache.CacheBuilder;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.CacheLoader;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.LoadingCache;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.UncheckedExecutionException;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** Cache of {@link String} computationId to {@link ComputationState}. */
+@Internal
+@ThreadSafe
+public final class ComputationStateCache implements StatusDataProvider {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ComputationStateCache.class);
+
+ private final ConcurrentMap<String, String> pipelineUserNameToStateFamilyNameMap;
+ private final LoadingCache<String, ComputationState> computationCache;
+
+ /**
+ * Fix up MapTask representation because MultiOutputInfos are missing from system generated
+ * ParDoInstructions.
+ */
+ private final Function<MapTask, MapTask> fixMultiOutputInfosOnParDoInstructions;
+
+ private ComputationStateCache(
+ LoadingCache<String, ComputationState> computationCache,
+ Function<MapTask, MapTask> fixMultiOutputInfosOnParDoInstructions,
+ ConcurrentMap<String, String> pipelineUserNameToStateFamilyNameMap) {
+ this.computationCache = computationCache;
+ this.fixMultiOutputInfosOnParDoInstructions = fixMultiOutputInfosOnParDoInstructions;
+ this.pipelineUserNameToStateFamilyNameMap = pipelineUserNameToStateFamilyNameMap;
+ }
+
+ public static ComputationStateCache create(
+ ComputationConfig.Fetcher computationConfigFetcher,
+ BoundedQueueExecutor workUnitExecutor,
+ Function<String, WindmillStateCache.ForComputation> perComputationStateCacheViewFactory,
+ IdGenerator idGenerator) {
+ Function<MapTask, MapTask> fixMultiOutputInfosOnParDoInstructions =
+ new FixMultiOutputInfosOnParDoInstructions(idGenerator);
+ ConcurrentMap<String, String> pipelineUserNameToStateFamilyNameMap = new ConcurrentHashMap<>();
+ return new ComputationStateCache(
+ CacheBuilder.newBuilder()
+ .build(
+ new CacheLoader<String, ComputationState>() {
+ @Override
+ public ComputationState load(String computationId) {
+ // LoadingCache load(K key) will throw an exception if we return null here,
+ // throw ComputationStateNotFoundException to represent semantics better.
+ ComputationConfig computationConfig =
+ computationConfigFetcher
+ .fetchConfig(computationId)
+ .orElseThrow(
+ () -> new ComputationStateNotFoundException(computationId));
+ pipelineUserNameToStateFamilyNameMap.putAll(computationConfig.stateNameMap());
+ Map<String, String> transformUserNameToStateFamilyForComputation =
+ !computationConfig.userTransformToStateFamilyName().isEmpty()
+ ? computationConfig.userTransformToStateFamilyName()
+ : pipelineUserNameToStateFamilyNameMap;
+ return new ComputationState(
+ computationId,
+ fixMultiOutputInfosOnParDoInstructions.apply(computationConfig.mapTask()),
+ workUnitExecutor,
+ transformUserNameToStateFamilyForComputation,
+ perComputationStateCacheViewFactory.apply(computationId));
+ }
+ }),
+ fixMultiOutputInfosOnParDoInstructions,
+ pipelineUserNameToStateFamilyNameMap);
+ }
+
+ @VisibleForTesting
+ public static ComputationStateCache forTesting(
+ ComputationConfig.Fetcher computationConfigFetcher,
+ BoundedQueueExecutor workUnitExecutor,
+ Function<String, WindmillStateCache.ForComputation> perComputationStateCacheViewFactory,
+ IdGenerator idGenerator,
+ ConcurrentMap<String, String> pipelineUserNameToStateFamilyNameMap) {
+ ComputationStateCache cache =
+ create(
+ computationConfigFetcher,
+ workUnitExecutor,
+ perComputationStateCacheViewFactory,
+ idGenerator);
+ cache.pipelineUserNameToStateFamilyNameMap.putAll(pipelineUserNameToStateFamilyNameMap);
+ return cache;
+ }
+
+ @VisibleForTesting
+ ImmutableMap<String, String> getGlobalUsernameToStateFamilyNameMap() {
+ return ImmutableMap.copyOf(pipelineUserNameToStateFamilyNameMap);
+ }
+
+ /**
+ * Returns the {@link ComputationState} associated with the given computationId. May perform IO if
+ * a value is not present, and it is possible that after IO is performed there is no value
+ * correlated with that computationId.
+ */
+ public Optional<ComputationState> get(String computationId) {
+ try {
+ return Optional.ofNullable(computationCache.get(computationId));
+ } catch (UncheckedExecutionException
+ | ExecutionException
+ | ComputationStateNotFoundException e) {
+ if (e.getCause() instanceof ComputationStateNotFoundException
+ || e instanceof ComputationStateNotFoundException) {
+ LOG.error(
+ "Trying to fetch unknown computation={}, known computations are {}.",
+ computationId,
+ ImmutableSet.copyOf(computationCache.asMap().keySet()));
+ } else {
+ LOG.warn("Error occurred fetching computation for computationId={}", computationId, e);
+ }
+ }
+
+ return Optional.empty();
+ }
+
+ public Optional<ComputationState> getIfPresent(String computationId) {
+ return Optional.ofNullable(computationCache.getIfPresent(computationId));
+ }
+
+ /** Returns a read-only view of all computations that have been fetched by the worker. */
+ public ImmutableList<ComputationState> getAllPresentComputations() {
+ return computationCache.asMap().values().stream().collect(toImmutableList());
+ }
+
+ /** Returns the current active {@link GetWorkBudget} across all known computations. */
+ public GetWorkBudget totalCurrentActiveGetWorkBudget() {
+ return computationCache.asMap().values().stream()
+ .map(ComputationState::getActiveWorkBudget)
+ .reduce(GetWorkBudget.noBudget(), GetWorkBudget::apply);
+ }
+
+ @VisibleForTesting
+ public void loadCacheForTesting(
+ List<MapTask> mapTasks, Function<MapTask, ComputationState> createComputationStateFn) {
+ Map<String, ComputationState> computationStates =
+ mapTasks.stream()
+ .map(fixMultiOutputInfosOnParDoInstructions)
+ .map(
+ mapTask -> {
+ LOG.info("Adding config for {}: {}", mapTask.getSystemName(), mapTask);
+ return createComputationStateFn.apply(mapTask);
+ })
+ .collect(toConcurrentMap(ComputationState::getComputationId, Function.identity()));
+ computationCache.putAll(computationStates);
+ }
+
+ /**
+ * Close all {@link ComputationState}(s) present in the cache, then invalidate the entire cache.
+ */
+ public void closeAndInvalidateAll() {
+ computationCache.asMap().values().forEach(ComputationState::close);
+ computationCache.invalidateAll();
+ }
+
+ @Override
+ public void appendSummaryHtml(PrintWriter writer) {
+ writer.println("<h1>Specs</h1>");
+ for (ComputationState computationState : getAllPresentComputations()) {
+ writer.println("<h3>" + computationState.getComputationId() + "</h3>");
+ writer.print("<script>document.write(JSON.stringify(");
+ writer.print(computationState.getMapTask().toString());
+ writer.println(", null, \"  \").replace(/\\n/g, \"<br>\"))</script>");
+ }
+ }
+
+ @VisibleForTesting
+ static class ComputationStateNotFoundException extends IllegalStateException {
+ private ComputationStateNotFoundException(String computationId) {
+ super("No computation found for computationId=[ " + computationId + " ]");
+ }
+ }
+}
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/config/ComputationConfig.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/config/ComputationConfig.java
new file mode 100644
index 00000000000..fb8bcf7edbf
--- /dev/null
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/config/ComputationConfig.java
@@ -0,0 +1,59 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.streaming.config;
+
+import com.google.api.services.dataflow.model.MapTask;
+import com.google.auto.value.AutoValue;
+import java.util.Map;
+import java.util.Optional;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.annotations.Internal;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
+
+@Internal
+@AutoValue
+public abstract class ComputationConfig {
+
+ public static ComputationConfig create(
+ MapTask mapTask,
+ @Nullable Map<String, String> userTransformToStateFamilyName,
+ Map<String, String> stateNameMap) {
+ return new AutoValue_ComputationConfig(
+ mapTask,
+ Optional.ofNullable(userTransformToStateFamilyName)
+ .map(ImmutableMap::copyOf)
+ .orElseGet(ImmutableMap::of),
+ ImmutableMap.copyOf(stateNameMap));
+ }
+
+ public abstract MapTask mapTask();
+
+ public abstract ImmutableMap<String, String> userTransformToStateFamilyName();
+
+ public abstract ImmutableMap<String, String> stateNameMap();
+
+ /** Interface to fetch configurations for a specific computation. */
+ @FunctionalInterface
+ public interface Fetcher {
+ default void start() {}
+
+ default void stop() {}
+
+ Optional<ComputationConfig> fetchConfig(String computationId);
+ }
+}
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/config/StreamingApplianceComputationConfigFetcher.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/config/StreamingApplianceComputationConfigFetcher.java
new file mode 100644
index 00000000000..786ded09498
--- /dev/null
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/config/StreamingApplianceComputationConfigFetcher.java
@@ -0,0 +1,136 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.streaming.config;
+
+import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap.toImmutableMap;
+
+import com.google.api.services.dataflow.model.MapTask;
+import java.io.IOException;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.ConcurrentHashMap;
+import javax.annotation.concurrent.ThreadSafe;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.GetConfigRequest;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.GetConfigResponse;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill.GetConfigResponse.NameMapEntry;
+import org.apache.beam.sdk.annotations.Internal;
+import org.apache.beam.sdk.extensions.gcp.util.Transport;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.HashBasedTable;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Table;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** Fetches computation config from Streaming Appliance. */
+@Internal
+@ThreadSafe
+public final class StreamingApplianceComputationConfigFetcher implements ComputationConfig.Fetcher {
+
+ private static final Logger LOG =
+ LoggerFactory.getLogger(StreamingApplianceComputationConfigFetcher.class);
+
+ private final ApplianceComputationConfigFetcher applianceComputationConfigFetcher;
+ private final ConcurrentHashMap<String, String> systemNameToComputationIdMap;
+
+ public StreamingApplianceComputationConfigFetcher(
+ ApplianceComputationConfigFetcher applianceComputationConfigFetcher) {
+ this.applianceComputationConfigFetcher = applianceComputationConfigFetcher;
+ this.systemNameToComputationIdMap = new ConcurrentHashMap<>();
+ }
+
+ /** Returns a {@code Table<ComputationId, TransformUserName, StateFamilyName>} */
+ private static Table<String, String, String> transformUserNameToStateFamilyByComputationId(
+ Windmill.GetConfigResponse response) {
+ Table<String, String, String> computationIdTransformUserNameStateFamilyNameTable =
+ HashBasedTable.create();
+ for (Windmill.GetConfigResponse.ComputationConfigMapEntry computationConfig :
+ response.getComputationConfigMapList()) {
+ for (Windmill.ComputationConfig.TransformUserNameToStateFamilyEntry entry :
+ computationConfig.getComputationConfig().getTransformUserNameToStateFamilyList()) {
+ computationIdTransformUserNameStateFamilyNameTable.put(
+ computationConfig.getComputationId(),
+ entry.getTransformUserName(),
+ entry.getStateFamily());
+ }
+ }
+
+ return computationIdTransformUserNameStateFamilyNameTable;
+ }
+
+ /** Deserialize {@link MapTask} and populate MultiOutputInfos in MapTask. */
+ private Optional<MapTask> deserializeMapTask(String serializedMapTask) {
+ try {
+ return Optional.of(Transport.getJsonFactory().fromString(serializedMapTask, MapTask.class));
+ } catch (IOException e) {
+ LOG.warn("Parsing MapTask failed: {}", serializedMapTask, e);
+ }
+ return Optional.empty();
+ }
+
+ @Override
+ public Optional<ComputationConfig> fetchConfig(String computationId) {
+ Preconditions.checkArgument(
+ !computationId.isEmpty(),
+ "computationId is empty. Cannot fetch computation config without a computationId.");
+
+ GetConfigResponse response =
+ applianceComputationConfigFetcher.fetchConfig(
+ GetConfigRequest.newBuilder().addComputations(computationId).build());
+
+ if (response == null) {
+ return Optional.empty();
+ }
+
+ for (Windmill.GetConfigResponse.SystemNameToComputationIdMapEntry entry :
+ response.getSystemNameToComputationIdMapList()) {
+ systemNameToComputationIdMap.put(entry.getSystemName(), entry.getComputationId());
+ }
+
+ return createComputationConfig(
+ // We are only fetching the config for 1 computation, so we should only be getting that
+ // computation back.
+ Iterables.getOnlyElement(response.getCloudWorksList()),
+ transformUserNameToStateFamilyByComputationId(response),
+ response.getNameMapList().stream()
+ .collect(toImmutableMap(NameMapEntry::getUserName, NameMapEntry::getSystemName)));
+ }
+
+ private Optional<ComputationConfig> createComputationConfig(
+ String serializedMapTask,
+ Table<String, String, String> transformUserNameToStateFamilyByComputationId,
+ Map<String, String> stateNameMap) {
+ return deserializeMapTask(serializedMapTask)
+ .map(
+ mapTask -> {
+ String computationId =
+ systemNameToComputationIdMap.getOrDefault(
+ mapTask.getSystemName(), mapTask.getSystemName());
+ return ComputationConfig.create(
+ mapTask,
+ transformUserNameToStateFamilyByComputationId.row(computationId),
+ stateNameMap);
+ });
+ }
+
+ @FunctionalInterface
+ public interface ApplianceComputationConfigFetcher {
+ GetConfigResponse fetchConfig(GetConfigRequest getConfigRequest);
+ }
+}
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/config/StreamingEngineComputationConfigFetcher.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/config/StreamingEngineComputationConfigFetcher.java
new file mode 100644
index 00000000000..51d1507af5f
--- /dev/null
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/config/StreamingEngineComputationConfigFetcher.java
@@ -0,0 +1,283 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.streaming.config;
+
+import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet.toImmutableSet;
+import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.Uninterruptibles.sleepUninterruptibly;
+
+import com.google.api.services.dataflow.model.MapTask;
+import com.google.api.services.dataflow.model.StreamingComputationConfig;
+import com.google.api.services.dataflow.model.StreamingConfigTask;
+import com.google.api.services.dataflow.model.WorkItem;
+import java.io.IOException;
+import java.util.Optional;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import java.util.stream.StreamSupport;
+import javax.annotation.concurrent.ThreadSafe;
+import org.apache.beam.runners.dataflow.worker.WorkUnitClient;
+import org.apache.beam.sdk.annotations.Internal;
+import org.apache.beam.sdk.util.BackOff;
+import org.apache.beam.sdk.util.BackOffUtils;
+import org.apache.beam.sdk.util.FluentBackoff;
+import org.apache.beam.sdk.util.Sleeper;
+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.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Splitter;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.net.HostAndPort;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.ThreadFactoryBuilder;
+import org.joda.time.Duration;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Streaming engine implementation of {@link ComputationConfig.Fetcher}. Asynchronously fetches
+ * global pipeline config.
+ */
+@Internal
+@ThreadSafe
+public final class StreamingEngineComputationConfigFetcher implements ComputationConfig.Fetcher {
+ private static final Logger LOG =
+ LoggerFactory.getLogger(StreamingEngineComputationConfigFetcher.class);
+ private static final String CONFIG_REFRESHER_THREAD_NAME = "GlobalPipelineConfigRefresher";
+ private static final int DEFAULT_WINDMILL_SERVICE_PORT = 443;
+ private static final FluentBackoff BACKOFF_FACTORY =
+ FluentBackoff.DEFAULT
+ .withInitialBackoff(Duration.millis(100))
+ .withMaxBackoff(Duration.standardMinutes(1))
+ .withMaxCumulativeBackoff(Duration.standardMinutes(5));
+
+ private final long globalConfigRefreshPeriodMillis;
+ private final WorkUnitClient dataflowServiceClient;
+ private final ScheduledExecutorService globalConfigRefresher;
+ private final Consumer<StreamingEnginePipelineConfig> onStreamingConfig;
+ private final AtomicBoolean hasReceivedGlobalConfig;
+
+ private StreamingEngineComputationConfigFetcher(
+ boolean hasReceivedGlobalConfig,
+ long globalConfigRefreshPeriodMillis,
+ WorkUnitClient dataflowServiceClient,
+ ScheduledExecutorService globalConfigRefresher,
+ Consumer<StreamingEnginePipelineConfig> onStreamingConfig) {
+ this.globalConfigRefreshPeriodMillis = globalConfigRefreshPeriodMillis;
+ this.dataflowServiceClient = dataflowServiceClient;
+ this.globalConfigRefresher = globalConfigRefresher;
+ this.onStreamingConfig = onStreamingConfig;
+ this.hasReceivedGlobalConfig = new AtomicBoolean(hasReceivedGlobalConfig);
+ }
+
+ public static StreamingEngineComputationConfigFetcher create(
+ long globalConfigRefreshPeriodMillis,
+ WorkUnitClient dataflowServiceClient,
+ Consumer<StreamingEnginePipelineConfig> onStreamingConfig) {
+ return new StreamingEngineComputationConfigFetcher(
+ /* hasReceivedGlobalConfig= */ false,
+ globalConfigRefreshPeriodMillis,
+ dataflowServiceClient,
+ Executors.newSingleThreadScheduledExecutor(
+ new ThreadFactoryBuilder().setNameFormat(CONFIG_REFRESHER_THREAD_NAME).build()),
+ onStreamingConfig);
+ }
+
+ @VisibleForTesting
+ public static StreamingEngineComputationConfigFetcher forTesting(
+ boolean hasReceivedGlobalConfig,
+ long globalConfigRefreshPeriodMillis,
+ WorkUnitClient dataflowServiceClient,
+ Function<String, ScheduledExecutorService> executorSupplier,
+ Consumer<StreamingEnginePipelineConfig> onStreamingConfig) {
+ return new StreamingEngineComputationConfigFetcher(
+ hasReceivedGlobalConfig,
+ globalConfigRefreshPeriodMillis,
+ dataflowServiceClient,
+ executorSupplier.apply(CONFIG_REFRESHER_THREAD_NAME),
+ onStreamingConfig);
+ }
+
+ @VisibleForTesting
+ static MapTask createMapTask(StreamingComputationConfig computationConfig) {
+ return new MapTask()
+ .setSystemName(computationConfig.getSystemName())
+ .setStageName(computationConfig.getStageName())
+ .setInstructions(computationConfig.getInstructions());
+ }
+
+ private static Optional<StreamingConfigTask> fetchConfigWithRetry(
+ ThrowingFetchWorkItemFn fetchWorkItemFn) {
+ BackOff backoff = BACKOFF_FACTORY.backoff();
+ while (true) {
+ try {
+ return fetchWorkItemFn
+ .fetchWorkItem()
+ .map(
+ workItem -> {
+ StreamingConfigTask config = workItem.getStreamingConfigTask();
+ Preconditions.checkState(
+ config != null,
+ "Received invalid WorkItem without StreamingConfigTask. WorkItem={}",
+ workItem);
+ return config;
+ });
+ } catch (IllegalArgumentException | IOException e) {
+ LOG.warn("Error fetching config: ", e);
+ try {
+ if (!BackOffUtils.next(Sleeper.DEFAULT, backoff)) {
+ return Optional.empty();
+ }
+ } catch (IOException ioe) {
+ LOG.warn("Error backing off, will not retry: ", ioe);
+ return Optional.empty();
+ } catch (InterruptedException ie) {
+ Thread.currentThread().interrupt();
+ return Optional.empty();
+ }
+ }
+ }
+ }
+
+ private static StreamingEnginePipelineConfig createPipelineConfig(StreamingConfigTask config) {
+ StreamingEnginePipelineConfig.Builder pipelineConfig = StreamingEnginePipelineConfig.builder();
+ if (config.getUserStepToStateFamilyNameMap() != null) {
+ pipelineConfig.setUserStepToStateFamilyNameMap(config.getUserStepToStateFamilyNameMap());
+ }
+
+ if (config.getWindmillServiceEndpoint() != null
+ && !config.getWindmillServiceEndpoint().isEmpty()) {
+ int windmillPort =
+ Optional.ofNullable(config.getWindmillServicePort())
+ .filter(port -> port != 0)
+ .map(Long::intValue)
+ .orElse(DEFAULT_WINDMILL_SERVICE_PORT);
+
+ ImmutableSet<HostAndPort> endpoints =
+ StreamSupport.stream(
+ Splitter.on(',').split(config.getWindmillServiceEndpoint()).spliterator(),
+ /* isParallel= */ false)
+ .map(endpoint -> HostAndPort.fromString(endpoint).withDefaultPort(windmillPort))
+ .collect(toImmutableSet());
+
+ pipelineConfig.setWindmillServiceEndpoints(endpoints);
+ }
+
+ if (config.getMaxWorkItemCommitBytes() != null
+ && config.getMaxWorkItemCommitBytes() > 0
+ && config.getMaxWorkItemCommitBytes() <= Integer.MAX_VALUE) {
+ pipelineConfig.setMaxWorkItemCommitBytes(config.getMaxWorkItemCommitBytes().intValue());
+ }
+
+ return pipelineConfig.build();
+ }
+
+ private static Optional<ComputationConfig> createComputationConfig(StreamingConfigTask config) {
+ return Optional.ofNullable(config.getStreamingComputationConfigs())
+ .map(Iterables::getOnlyElement)
+ .map(
+ streamingComputationConfig ->
+ ComputationConfig.create(
+ createMapTask(streamingComputationConfig),
+ streamingComputationConfig.getTransformUserNameToStateFamily(),
+ config.getUserStepToStateFamilyNameMap() != null
+ ? config.getUserStepToStateFamilyNameMap()
+ : ImmutableMap.of()));
+ }
+
+ @Override
+ public void start() {
+ fetchInitialPipelineGlobalConfig();
+ schedulePeriodicGlobalConfigRequests();
+ }
+
+ @Override
+ public Optional<ComputationConfig> fetchConfig(String computationId) {
+ Preconditions.checkArgument(
+ !computationId.isEmpty(),
+ "computationId is empty. Cannot fetch computation config without a computationId.");
+ return fetchConfigWithRetry(
+ () -> dataflowServiceClient.getStreamingConfigWorkItem(computationId))
+ .flatMap(StreamingEngineComputationConfigFetcher::createComputationConfig);
+ }
+
+ @Override
+ public void stop() {
+ // We have already shutdown or start has not been called.
+ if (globalConfigRefresher.isShutdown() || !hasReceivedGlobalConfig.get()) {
+ return;
+ }
+
+ globalConfigRefresher.shutdown();
+ boolean isTerminated = false;
+ try {
+ isTerminated = globalConfigRefresher.awaitTermination(10, TimeUnit.SECONDS);
+ } catch (InterruptedException e) {
+ LOG.warn("Error occurred shutting down: {}", globalConfigRefresher);
+ }
+ if (!isTerminated) {
+ globalConfigRefresher.shutdownNow();
+ }
+ }
+
+ /**
+ * Initially b Schedules a background thread that periodically sends getConfig requests to
+ * Dataflow Service to obtain the windmill service endpoint.
+ */
+ @SuppressWarnings("FutureReturnValueIgnored")
+ private void schedulePeriodicGlobalConfigRequests() {
+ globalConfigRefresher.scheduleWithFixedDelay(
+ () -> fetchGlobalConfig().ifPresent(onStreamingConfig),
+ 0,
+ globalConfigRefreshPeriodMillis,
+ TimeUnit.MILLISECONDS);
+ }
+
+ /**
+ * Blocks until we have received the initial global config from Dataflow or an exception is
+ * thrown.
+ */
+ private synchronized void fetchInitialPipelineGlobalConfig() {
+ while (!hasReceivedGlobalConfig.get()) {
+ LOG.info("Sending request to get initial global configuration for this worker.");
+ Optional<StreamingEnginePipelineConfig> globalConfig = fetchGlobalConfig();
+ if (globalConfig.isPresent()) {
+ onStreamingConfig.accept(globalConfig.get());
+ hasReceivedGlobalConfig.set(true);
+ break;
+ }
+ LOG.info("Haven't received initial global configuration, will retry in 5 seconds");
+ sleepUninterruptibly(5, TimeUnit.SECONDS);
+ }
+
+ LOG.info("Initial global configuration received, harness is now ready");
+ }
+
+ private Optional<StreamingEnginePipelineConfig> fetchGlobalConfig() {
+ return fetchConfigWithRetry(dataflowServiceClient::getGlobalStreamingConfigWorkItem)
+ .map(StreamingEngineComputationConfigFetcher::createPipelineConfig);
+ }
+
+ @FunctionalInterface
+ private interface ThrowingFetchWorkItemFn {
+ Optional<WorkItem> fetchWorkItem() throws IOException;
+ }
+}
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/config/StreamingEnginePipelineConfig.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/config/StreamingEnginePipelineConfig.java
new file mode 100644
index 00000000000..b5b761ada70
--- /dev/null
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/config/StreamingEnginePipelineConfig.java
@@ -0,0 +1,57 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.streaming.config;
+
+import com.google.auto.value.AutoValue;
+import java.util.HashMap;
+import java.util.Map;
+import org.apache.beam.sdk.annotations.Internal;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.net.HostAndPort;
+
+/** Global pipeline config for pipelines running in Streaming Engine mode. */
+@AutoValue
+@Internal
+public abstract class StreamingEnginePipelineConfig {
+
+ private static final long DEFAULT_MAX_WORK_ITEM_COMMIT_BYTES = 180 << 20;
+
+ public static StreamingEnginePipelineConfig.Builder builder() {
+ return new AutoValue_StreamingEnginePipelineConfig.Builder()
+ .setMaxWorkItemCommitBytes(DEFAULT_MAX_WORK_ITEM_COMMIT_BYTES)
+ .setUserStepToStateFamilyNameMap(new HashMap<>())
+ .setWindmillServiceEndpoints(ImmutableSet.of());
+ }
+
+ public abstract long maxWorkItemCommitBytes();
+
+ public abstract Map<String, String> userStepToStateFamilyNameMap();
+
+ public abstract ImmutableSet<HostAndPort> windmillServiceEndpoints();
+
+ @AutoValue.Builder
+ public abstract static class Builder {
+ public abstract Builder setMaxWorkItemCommitBytes(long value);
+
+ public abstract Builder setUserStepToStateFamilyNameMap(Map<String, String> value);
+
+ public abstract Builder setWindmillServiceEndpoints(ImmutableSet<HostAndPort> value);
+
+ public abstract StreamingEnginePipelineConfig build();
+ }
+}
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/harness/MetricsDataProvider.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/harness/MetricsDataProvider.java
new file mode 100644
index 00000000000..901e2d235f8
--- /dev/null
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/harness/MetricsDataProvider.java
@@ -0,0 +1,83 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.streaming.harness;
+
+import java.io.PrintWriter;
+import java.util.Collection;
+import java.util.function.Consumer;
+import java.util.function.Supplier;
+import org.apache.beam.runners.dataflow.worker.status.StatusDataProvider;
+import org.apache.beam.runners.dataflow.worker.streaming.ComputationState;
+import org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor;
+import org.apache.beam.sdk.annotations.Internal;
+
+@Internal
+class MetricsDataProvider implements StatusDataProvider {
+
+ private final BoundedQueueExecutor workUnitExecutor;
+ private final Supplier<Long> currentActiveCommitBytes;
+ private final Consumer<PrintWriter> getDataStatusProvider;
+ private final Supplier<Collection<ComputationState>> allComputationStates;
+
+ MetricsDataProvider(
+ BoundedQueueExecutor workUnitExecutor,
+ Supplier<Long> currentActiveCommitBytes,
+ Consumer<PrintWriter> getDataStatusProvider,
+ Supplier<Collection<ComputationState>> allComputationStates) {
+ this.workUnitExecutor = workUnitExecutor;
+ this.currentActiveCommitBytes = currentActiveCommitBytes;
+ this.getDataStatusProvider = getDataStatusProvider;
+ this.allComputationStates = allComputationStates;
+ }
+
+ @Override
+ public void appendSummaryHtml(PrintWriter writer) {
+ writer.println(workUnitExecutor.summaryHtml());
+
+ writer.print("Active commit: ");
+ appendHumanizedBytes(currentActiveCommitBytes.get(), writer);
+ writer.println("<br>");
+
+ getDataStatusProvider.accept(writer);
+
+ writer.println("<br>");
+
+ writer.println("Active Keys: <br>");
+ for (ComputationState computationState : allComputationStates.get()) {
+ writer.print(computationState.getComputationId());
+ writer.print(":<br>");
+ computationState.printActiveWork(writer);
+ writer.println("<br>");
+ }
+ }
+
+ private void appendHumanizedBytes(long bytes, PrintWriter writer) {
+ if (bytes < (4 << 10)) {
+ writer.print(bytes);
+ writer.print("B");
+ } else if (bytes < (4 << 20)) {
+ writer.print("~");
+ writer.print(bytes >> 10);
+ writer.print("KB");
+ } else {
+ writer.print("~");
+ writer.print(bytes >> 20);
+ writer.print("MB");
+ }
+ }
+}
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
new file mode 100644
index 00000000000..d305e25af7e
--- /dev/null
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/harness/StreamingWorkerStatusPages.java
@@ -0,0 +1,261 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.streaming.harness;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+
+import com.google.auto.value.AutoBuilder;
+import java.io.File;
+import java.io.IOException;
+import java.io.PrintWriter;
+import java.util.Collection;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Consumer;
+import java.util.function.Supplier;
+import javax.servlet.http.HttpServletRequest;
+import javax.servlet.http.HttpServletResponse;
+import org.apache.beam.runners.dataflow.options.DataflowWorkerHarnessOptions;
+import org.apache.beam.runners.dataflow.worker.status.BaseStatusServlet;
+import org.apache.beam.runners.dataflow.worker.status.DebugCapture;
+import org.apache.beam.runners.dataflow.worker.status.LastExceptionDataProvider;
+import org.apache.beam.runners.dataflow.worker.status.WorkerStatusPages;
+import org.apache.beam.runners.dataflow.worker.streaming.ComputationStateCache;
+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.state.WindmillStateCache;
+import org.apache.beam.sdk.annotations.Internal;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.ThreadFactoryBuilder;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Dataflow status pages for debugging current worker and processing state.
+ *
+ * @implNote Class member state should only be accessed, not modified.
+ */
+@Internal
+public final class StreamingWorkerStatusPages {
+ private static final Logger LOG = LoggerFactory.getLogger(StreamingWorkerStatusPages.class);
+ private static final String DUMP_STATUS_PAGES_EXECUTOR = "DumpStatusPages";
+ private static final long DEFAULT_STATUS_PAGE_DUMP_PERIOD_SECONDS = 60;
+
+ private final Supplier<Instant> clock;
+ private final long clientId;
+ private final AtomicBoolean isRunning;
+ private final WorkerStatusPages statusPages;
+ private final WindmillStateCache stateCache;
+ private final ComputationStateCache computationStateCache;
+ private final Supplier<Long> currentActiveCommitBytes;
+ private final Consumer<PrintWriter> getDataStatusProvider;
+ private final BoundedQueueExecutor workUnitExecutor;
+ private final ScheduledExecutorService statusPageDumper;
+
+ // StreamingEngine status providers.
+ private final @Nullable GrpcWindmillStreamFactory windmillStreamFactory;
+ private final DebugCapture.@Nullable Manager debugCapture;
+ private final @Nullable ChannelzServlet channelzServlet;
+
+ StreamingWorkerStatusPages(
+ Supplier<Instant> clock,
+ long clientId,
+ AtomicBoolean isRunning,
+ WorkerStatusPages statusPages,
+ DebugCapture.@Nullable Manager debugCapture,
+ @Nullable ChannelzServlet channelzServlet,
+ WindmillStateCache stateCache,
+ ComputationStateCache computationStateCache,
+ Supplier<Long> currentActiveCommitBytes,
+ @Nullable GrpcWindmillStreamFactory windmillStreamFactory,
+ Consumer<PrintWriter> getDataStatusProvider,
+ BoundedQueueExecutor workUnitExecutor,
+ ScheduledExecutorService statusPageDumper) {
+ this.clock = clock;
+ this.clientId = clientId;
+ this.isRunning = isRunning;
+ this.statusPages = statusPages;
+ this.debugCapture = debugCapture;
+ this.channelzServlet = channelzServlet;
+ this.stateCache = stateCache;
+ this.computationStateCache = computationStateCache;
+ this.currentActiveCommitBytes = currentActiveCommitBytes;
+ this.windmillStreamFactory = windmillStreamFactory;
+ this.getDataStatusProvider = getDataStatusProvider;
+ this.workUnitExecutor = workUnitExecutor;
+ this.statusPageDumper = statusPageDumper;
+ }
+
+ public static StreamingWorkerStatusPages.Builder builder() {
+ return new AutoBuilder_StreamingWorkerStatusPages_Builder()
+ .setStatusPageDumper(
+ Executors.newSingleThreadScheduledExecutor(
+ new ThreadFactoryBuilder().setNameFormat(DUMP_STATUS_PAGES_EXECUTOR).build()));
+ }
+
+ public void start(DataflowWorkerHarnessOptions options) {
+ statusPages.addServlet(stateCache.statusServlet());
+ statusPages.addServlet(newSpecServlet());
+ statusPages.addStatusDataProvider(
+ "harness",
+ "Harness",
+ writer -> {
+ writer.println("Running: " + isRunning.get() + "<br>");
+ writer.println("ID: " + clientId + "<br>");
+ });
+
+ statusPages.addStatusDataProvider(
+ "metrics",
+ "Metrics",
+ new MetricsDataProvider(
+ workUnitExecutor,
+ currentActiveCommitBytes,
+ getDataStatusProvider,
+ computationStateCache::getAllPresentComputations));
+ statusPages.addStatusDataProvider(
+ "exception", "Last Exception", new LastExceptionDataProvider());
+ statusPages.addStatusDataProvider("cache", "State Cache", stateCache);
+
+ if (isStreamingEngine()) {
+ addStreamingEngineStatusPages();
+ }
+
+ statusPages.start();
+ scheduleStatusPageDump(options);
+ }
+
+ private void addStreamingEngineStatusPages() {
+ Preconditions.checkNotNull(debugCapture).start();
+ statusPages.addServlet(Preconditions.checkNotNull(channelzServlet));
+ statusPages.addCapturePage(Preconditions.checkNotNull(channelzServlet));
+ statusPages.addStatusDataProvider(
+ "streaming", "Streaming RPCs", Preconditions.checkNotNull(windmillStreamFactory));
+ }
+
+ private boolean isStreamingEngine() {
+ return debugCapture != null && channelzServlet != null && windmillStreamFactory != null;
+ }
+
+ public void stop() {
+ statusPages.stop();
+ if (debugCapture != null) {
+ debugCapture.stop();
+ }
+ statusPageDumper.shutdown();
+ try {
+ statusPageDumper.awaitTermination(300, TimeUnit.SECONDS);
+ } catch (InterruptedException e) {
+ LOG.warn("Error occurred shutting down periodic status page dumper", e);
+ }
+ statusPageDumper.shutdownNow();
+ }
+
+ @SuppressWarnings("FutureReturnValueIgnored")
+ private void scheduleStatusPageDump(DataflowWorkerHarnessOptions options) {
+ if (options.getPeriodicStatusPageOutputDirectory() != null) {
+ LOG.info(
+ "Scheduling status page dump every {} seconds", DEFAULT_STATUS_PAGE_DUMP_PERIOD_SECONDS);
+ statusPageDumper.scheduleWithFixedDelay(
+ () ->
+ dumpStatusPages(
+ options.getPeriodicStatusPageOutputDirectory(), options.getWorkerId()),
+ DEFAULT_STATUS_PAGE_DUMP_PERIOD_SECONDS,
+ DEFAULT_STATUS_PAGE_DUMP_PERIOD_SECONDS,
+ TimeUnit.SECONDS);
+ } else {
+ LOG.info(
+ "Status page output directory was not set, "
+ + "status pages will not be periodically dumped. "
+ + "If this was not intended check pipeline options.");
+ }
+ }
+
+ private void dumpStatusPages(String getPeriodicStatusPageOutputDirectory, String workerId) {
+ Collection<DebugCapture.Capturable> pages = statusPages.getDebugCapturePages();
+ if (pages.isEmpty()) {
+ LOG.warn("No captured status pages.");
+ }
+ long timestamp = clock.get().getMillis();
+ for (DebugCapture.Capturable page : pages) {
+ PrintWriter writer = null;
+ try {
+ File outputFile =
+ new File(
+ getPeriodicStatusPageOutputDirectory,
+ ("StreamingDataflowWorker" + workerId + "_" + page.pageName() + timestamp + ".html")
+ .replaceAll("/", "_"));
+ writer = new PrintWriter(outputFile, UTF_8.name());
+ page.captureData(writer);
+ } catch (IOException e) {
+ LOG.warn("Error dumping status page.", e);
+ } finally {
+ if (writer != null) {
+ writer.close();
+ }
+ }
+ }
+ }
+
+ private BaseStatusServlet newSpecServlet() {
+ return new BaseStatusServlet("/spec") {
+ @Override
+ protected void doGet(HttpServletRequest request, HttpServletResponse response)
+ throws IOException {
+ PrintWriter writer = response.getWriter();
+ computationStateCache.appendSummaryHtml(writer);
+ }
+ };
+ }
+
+ @Internal
+ @AutoBuilder(ofClass = StreamingWorkerStatusPages.class)
+ public interface Builder {
+ Builder setClock(Supplier<Instant> clock);
+
+ Builder setClientId(long clientId);
+
+ Builder setIsRunning(AtomicBoolean isRunning);
+
+ Builder setStatusPages(WorkerStatusPages statusPages);
+
+ Builder setDebugCapture(DebugCapture.Manager debugCapture);
+
+ Builder setChannelzServlet(ChannelzServlet channelzServlet);
+
+ Builder setStateCache(WindmillStateCache stateCache);
+
+ Builder setComputationStateCache(ComputationStateCache computationStateCache);
+
+ Builder setCurrentActiveCommitBytes(Supplier<Long> currentActiveCommitBytes);
+
+ Builder setWindmillStreamFactory(GrpcWindmillStreamFactory windmillStreamFactory);
+
+ Builder setGetDataStatusProvider(Consumer<PrintWriter> getDataStatusProvider);
+
+ Builder setWorkUnitExecutor(BoundedQueueExecutor workUnitExecutor);
+
+ Builder setStatusPageDumper(ScheduledExecutorService statusPageDumper);
+
+ StreamingWorkerStatusPages build();
+ }
+}
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/WorkExecutor.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/WorkExecutor.java
index 2f836dfea3b..b7170c80ced 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/WorkExecutor.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/WorkExecutor.java
@@ -61,7 +61,7 @@ public interface WorkExecutor extends AutoCloseable {
}
@Override
- default void close() throws Exception {
+ default void close() {
// By default, nothing to close or shut down.
}
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/WindmillServerBase.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/WindmillServerBase.java
index a1160b4f98d..0785ae96626 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/WindmillServerBase.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/WindmillServerBase.java
@@ -50,7 +50,7 @@ public class WindmillServerBase extends WindmillServerStub {
}
@Override
- public void setWindmillServiceEndpoints(Set<HostAndPort> endpoints) throws IOException {
+ public void setWindmillServiceEndpoints(Set<HostAndPort> endpoints) {
// This class is used for windmill appliance and local runner tests.
}
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/WindmillServerStub.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/WindmillServerStub.java
index 34461ab471f..7d0c4f5aba3 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/WindmillServerStub.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/WindmillServerStub.java
@@ -17,7 +17,6 @@
*/
package org.apache.beam.runners.dataflow.worker.windmill;
-import java.io.IOException;
import java.io.PrintWriter;
import java.util.Set;
import org.apache.beam.runners.dataflow.worker.status.StatusDataProvider;
@@ -39,7 +38,7 @@ public abstract class WindmillServerStub implements StatusDataProvider {
* subsequent calls, if endpoints are different from previous values new stubs are created,
* replacing the previous ones.
*/
- public abstract void setWindmillServiceEndpoints(Set<HostAndPort> endpoints) throws IOException;
+ public abstract void setWindmillServiceEndpoints(Set<HostAndPort> endpoints);
/*
* Returns the windmill service endpoints set by setWindmillServiceEndpoints
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 edc193ff99c..824b40fa42b 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
@@ -125,7 +125,7 @@ public class GrpcDispatcherClient {
return dispatcherStubs.get().hasInitializedEndpoints();
}
- synchronized void consumeWindmillDispatcherEndpoints(
+ public synchronized void consumeWindmillDispatcherEndpoints(
ImmutableSet<HostAndPort> dispatcherEndpoints) {
ImmutableSet<HostAndPort> currentDispatcherEndpoints =
dispatcherStubs.get().dispatcherEndpoints();
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetDataStream.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetDataStream.java
index 992c9fb589b..feb15c2ac83 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetDataStream.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetDataStream.java
@@ -72,7 +72,7 @@ public final class GrpcGetDataStream
// If true, then active work refreshes will be sent as KeyedGetDataRequests. Otherwise, use the
// newer ComputationHeartbeatRequests.
private final boolean sendKeyedGetDataRequests;
- private Consumer<List<ComputationHeartbeatResponse>> processHeartbeatResponses;
+ private final Consumer<List<ComputationHeartbeatResponse>> processHeartbeatResponses;
private GrpcGetDataStream(
Function<StreamObserver<StreamingGetDataResponse>, StreamObserver<StreamingGetDataRequest>>
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 52de3ef7bc0..69807c523ed 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
@@ -20,7 +20,6 @@ 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.WindmillChannelFactory.remoteChannel;
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.io.IOException;
@@ -30,7 +29,6 @@ import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.function.Consumer;
-import java.util.function.Function;
import java.util.function.Supplier;
import javax.annotation.Nullable;
import org.apache.beam.runners.dataflow.DataflowRunner;
@@ -54,13 +52,9 @@ import org.apache.beam.runners.dataflow.worker.windmill.Windmill.ReportStatsRequ
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.ReportStatsResponse;
import org.apache.beam.runners.dataflow.worker.windmill.WindmillApplianceGrpc;
import org.apache.beam.runners.dataflow.worker.windmill.WindmillServerStub;
-import org.apache.beam.runners.dataflow.worker.windmill.WindmillServiceAddress;
import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStream.CommitWorkStream;
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.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.IsolationChannel;
import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs.WindmillStubFactory;
import org.apache.beam.runners.dataflow.worker.windmill.client.throttling.StreamingEngineThrottleTimers;
import org.apache.beam.runners.dataflow.worker.windmill.work.WorkItemReceiver;
@@ -154,29 +148,11 @@ public final class GrpcWindmillServer extends WindmillServerStub {
public static GrpcWindmillServer create(
DataflowWorkerHarnessOptions workerOptions,
GrpcWindmillStreamFactory grpcWindmillStreamFactory,
- Consumer<List<Windmill.ComputationHeartbeatResponse>> processHeartbeatResponses)
- throws IOException {
- Function<WindmillServiceAddress, ManagedChannel> channelFactory =
- serviceAddress ->
- remoteChannel(
- serviceAddress, workerOptions.getWindmillServiceRpcChannelAliveTimeoutSec());
- ChannelCache channelCache =
- ChannelCache.create(
- serviceAddress ->
- // IsolationChannel will create and manage separate RPC channels to the same
- // serviceAddress via calling the channelFactory, else just directly return the
- // RPC channel.
- workerOptions.getUseWindmillIsolatedChannels()
- ? IsolationChannel.create(() -> channelFactory.apply(serviceAddress))
- : channelFactory.apply(serviceAddress));
- WindmillStubFactory stubFactory =
- ChannelCachingRemoteStubFactory.create(workerOptions.getGcpCredential(), channelCache);
+ GrpcDispatcherClient dispatcherClient,
+ Consumer<List<Windmill.ComputationHeartbeatResponse>> processHeartbeatResponses) {
GrpcWindmillServer grpcWindmillServer =
new GrpcWindmillServer(
- workerOptions,
- grpcWindmillStreamFactory,
- GrpcDispatcherClient.create(stubFactory),
- processHeartbeatResponses);
+ workerOptions, grpcWindmillStreamFactory, dispatcherClient, processHeartbeatResponses);
if (workerOptions.getWindmillServiceEndpoint() != null) {
grpcWindmillServer.configureWindmillServiceEndpoints();
} else if (!workerOptions.isEnableStreamingEngine()
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 4fe5fa00c21..148adedbafc 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
@@ -57,6 +57,7 @@ import com.google.api.services.dataflow.model.StreamingConfigTask;
import com.google.api.services.dataflow.model.WorkItem;
import com.google.api.services.dataflow.model.WorkItemStatus;
import com.google.api.services.dataflow.model.WriteInstruction;
+import com.google.auto.value.AutoValue;
import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
@@ -72,7 +73,6 @@ import java.util.Optional;
import java.util.PriorityQueue;
import java.util.concurrent.Callable;
import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Executors;
@@ -88,8 +88,8 @@ import java.util.concurrent.atomic.AtomicLong;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.Supplier;
+import javax.annotation.Nullable;
import org.apache.beam.runners.dataflow.internal.CustomSources;
-import org.apache.beam.runners.dataflow.options.DataflowPipelineDebugOptions;
import org.apache.beam.runners.dataflow.options.DataflowPipelineOptions;
import org.apache.beam.runners.dataflow.options.DataflowWorkerHarnessOptions;
import org.apache.beam.runners.dataflow.util.CloudObject;
@@ -97,6 +97,7 @@ import org.apache.beam.runners.dataflow.util.CloudObjects;
import org.apache.beam.runners.dataflow.util.PropertyNames;
import org.apache.beam.runners.dataflow.util.Structs;
import org.apache.beam.runners.dataflow.worker.streaming.ComputationState;
+import org.apache.beam.runners.dataflow.worker.streaming.ComputationStateCache;
import org.apache.beam.runners.dataflow.worker.streaming.ShardedKey;
import org.apache.beam.runners.dataflow.worker.streaming.Work;
import org.apache.beam.runners.dataflow.worker.testing.RestoreDataflowLoggingMDC;
@@ -179,6 +180,7 @@ import org.hamcrest.Matcher;
import org.hamcrest.Matchers;
import org.joda.time.Duration;
import org.joda.time.Instant;
+import org.junit.After;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Ignore;
@@ -200,7 +202,7 @@ import org.slf4j.LoggerFactory;
@RunWith(Parameterized.class)
// TODO(https://github.com/apache/beam/issues/21230): Remove when new version of errorprone is
// released (2.11.0)
-@SuppressWarnings("unused")
+@SuppressWarnings({"unused", "deprecation"})
public class StreamingDataflowWorkerTest {
private static final Logger LOG = LoggerFactory.getLogger(StreamingDataflowWorkerTest.class);
private static final IntervalWindow DEFAULT_WINDOW =
@@ -271,9 +273,10 @@ public class StreamingDataflowWorkerTest {
WorkUnitClient mockWorkUnitClient = mock(WorkUnitClient.class);
HotKeyLogger hotKeyLogger = mock(HotKeyLogger.class);
- private final ConcurrentMap<String, ComputationState> computationMap = new ConcurrentHashMap<>();
+ private @Nullable ComputationStateCache computationStateCache = null;
private final FakeWindmillServer server =
- new FakeWindmillServer(errorCollector, id -> Optional.ofNullable(computationMap.get(id)));
+ new FakeWindmillServer(
+ errorCollector, computationId -> computationStateCache.get(computationId));
public StreamingDataflowWorkerTest(Boolean streamingEngine) {
this.streamingEngine = streamingEngine;
@@ -295,10 +298,15 @@ public class StreamingDataflowWorkerTest {
@Before
public void setUp() {
- computationMap.clear();
server.clearCommitsReceived();
}
+ @After
+ public void cleanUp() {
+ Optional.ofNullable(computationStateCache)
+ .ifPresent(ComputationStateCache::closeAndInvalidateAll);
+ }
+
static Work createMockWork(long workToken) {
return createMockWork(workToken, work -> {});
}
@@ -316,6 +324,11 @@ public class StreamingDataflowWorkerTest {
DEFAULT_WINDOW_COLLECTION_CODER, Collections.singletonList(window));
}
+ private StreamingDataflowWorkerTestParams.Builder defaultWorkerParams(String... options) {
+ return StreamingDataflowWorkerTestParams.builder()
+ .setOptions(createTestingPipelineOptions(options));
+ }
+
private String keyStringForIndex(int index) {
return DEFAULT_KEY_STRING + index;
}
@@ -786,80 +799,44 @@ public class StreamingDataflowWorkerTest {
options.setProject("test_project");
options.setWorkerId("test_worker");
options.setStreaming(true);
- options.setActiveWorkRefreshPeriodMillis(0);
+
+ // If activeWorkRefreshPeriodMillis is the default value, override it.
+ if (options.getActiveWorkRefreshPeriodMillis() == 10000) {
+ options.setActiveWorkRefreshPeriodMillis(0);
+ }
+
return options;
}
private StreamingDataflowWorker makeWorker(
- List<ParallelInstruction> instructions,
- DataflowWorkerHarnessOptions options,
- boolean publishCounters,
- Supplier<Instant> clock,
- Function<String, ScheduledExecutorService> executorSupplier,
- int localRetryTimeoutMs) {
+ StreamingDataflowWorkerTestParams streamingDataflowWorkerTestParams) {
StreamingDataflowWorker worker =
StreamingDataflowWorker.forTesting(
- computationMap,
+ streamingDataflowWorkerTestParams.stateNameMappings(),
server,
- Collections.singletonList(defaultMapTask(instructions)),
+ Collections.singletonList(
+ defaultMapTask(streamingDataflowWorkerTestParams.instructions())),
IntrinsicMapTaskExecutorFactory.defaultFactory(),
mockWorkUnitClient,
- options,
- publishCounters,
+ streamingDataflowWorkerTestParams.options(),
+ streamingDataflowWorkerTestParams.publishCounters(),
hotKeyLogger,
- clock,
- executorSupplier,
- localRetryTimeoutMs);
- worker.addStateNameMappings(
- ImmutableMap.of(DEFAULT_PARDO_USER_NAME, DEFAULT_PARDO_STATE_FAMILY));
+ streamingDataflowWorkerTestParams.clock(),
+ streamingDataflowWorkerTestParams.executorSupplier(),
+ streamingDataflowWorkerTestParams.localRetryTimeoutMs(),
+ streamingDataflowWorkerTestParams.maxWorkItemCommitBytes());
+ this.computationStateCache = worker.getComputationStateCache();
return worker;
}
- private StreamingDataflowWorker makeWorker(
- List<ParallelInstruction> instructions,
- DataflowWorkerHarnessOptions options,
- boolean publishCounters,
- Supplier<Instant> clock,
- Function<String, ScheduledExecutorService> executorSupplier) {
- return makeWorker(instructions, options, publishCounters, clock, executorSupplier, -1);
- }
-
- private StreamingDataflowWorker makeWorker(
- List<ParallelInstruction> instructions,
- DataflowWorkerHarnessOptions options,
- boolean publishCounters) {
- return makeWorker(
- instructions,
- options,
- publishCounters,
- Instant::now,
- (threadName) -> Executors.newSingleThreadScheduledExecutor(),
- -1);
- }
-
- private StreamingDataflowWorker makeWorker(
- List<ParallelInstruction> instructions,
- DataflowWorkerHarnessOptions options,
- boolean publishCounters,
- int localRetryTimeoutMs) {
- return makeWorker(
- instructions,
- options,
- publishCounters,
- Instant::now,
- (threadName) -> Executors.newSingleThreadScheduledExecutor(),
- localRetryTimeoutMs);
- }
-
@Test
public void testBasicHarness() throws Exception {
List<ParallelInstruction> instructions =
Arrays.asList(
makeSourceInstruction(StringUtf8Coder.of()),
makeSinkInstruction(StringUtf8Coder.of(), 0));
-
StreamingDataflowWorker worker =
- makeWorker(instructions, createTestingPipelineOptions(), true /* publishCounters */);
+ makeWorker(defaultWorkerParams().setInstructions(instructions).publishCounters().build());
worker.start();
final int numIters = 2000;
@@ -894,9 +871,8 @@ public class StreamingDataflowWorkerTest {
WorkItem workItem = new WorkItem();
workItem.setStreamingConfigTask(streamingConfig);
when(mockWorkUnitClient.getGlobalStreamingConfigWorkItem()).thenReturn(Optional.of(workItem));
-
StreamingDataflowWorker worker =
- makeWorker(instructions, createTestingPipelineOptions(), true /* publishCounters */);
+ makeWorker(defaultWorkerParams().setInstructions(instructions).publishCounters().build());
worker.start();
final int numIters = 2000;
@@ -944,12 +920,12 @@ public class StreamingDataflowWorkerTest {
WorkItem workItem = new WorkItem();
workItem.setStreamingConfigTask(streamingConfig);
when(mockWorkUnitClient.getGlobalStreamingConfigWorkItem()).thenReturn(Optional.of(workItem));
-
StreamingDataflowWorker worker =
makeWorker(
- instructions,
- createTestingPipelineOptions("--hotKeyLoggingEnabled=true"),
- true /* publishCounters */);
+ defaultWorkerParams("--hotKeyLoggingEnabled=true")
+ .setInstructions(instructions)
+ .publishCounters()
+ .build());
worker.start();
final int numIters = 2000;
@@ -985,7 +961,7 @@ public class StreamingDataflowWorkerTest {
when(mockWorkUnitClient.getGlobalStreamingConfigWorkItem()).thenReturn(Optional.of(workItem));
StreamingDataflowWorker worker =
- makeWorker(instructions, createTestingPipelineOptions(), true /* publishCounters */);
+ makeWorker(defaultWorkerParams().setInstructions(instructions).publishCounters().build());
worker.start();
final int numIters = 2000;
@@ -1009,9 +985,8 @@ public class StreamingDataflowWorkerTest {
makeSourceInstruction(StringUtf8Coder.of()),
makeDoFnInstruction(blockingFn, 0, StringUtf8Coder.of()),
makeSinkInstruction(StringUtf8Coder.of(), 0));
-
StreamingDataflowWorker worker =
- makeWorker(instructions, createTestingPipelineOptions(), true /* publishCounters */);
+ makeWorker(defaultWorkerParams().setInstructions(instructions).publishCounters().build());
worker.start();
for (int i = 0; i < numIters; ++i) {
@@ -1132,10 +1107,12 @@ public class StreamingDataflowWorkerTest {
makeDoFnInstruction(blockingFn, 0, StringUtf8Coder.of()),
makeSinkInstruction(StringUtf8Coder.of(), 0));
- DataflowWorkerHarnessOptions options = createTestingPipelineOptions();
- options.setNumberOfWorkerHarnessThreads(expectedNumberOfThreads);
-
- StreamingDataflowWorker worker = makeWorker(instructions, options, true /* publishCounters */);
+ StreamingDataflowWorker worker =
+ makeWorker(
+ defaultWorkerParams("--numberOfWorkerHarnessThreads=" + expectedNumberOfThreads)
+ .setInstructions(instructions)
+ .publishCounters()
+ .build());
worker.start();
for (int i = 0; i < expectedNumberOfThreads * 2; ++i) {
@@ -1180,7 +1157,7 @@ public class StreamingDataflowWorkerTest {
.thenReturn(makeInput(0, 0, DEFAULT_KEY_STRING, DEFAULT_SHARDING_KEY));
StreamingDataflowWorker worker =
- makeWorker(instructions, createTestingPipelineOptions(), true /* publishCounters */);
+ makeWorker(defaultWorkerParams().setInstructions(instructions).publishCounters().build());
worker.start();
server.waitForEmptyWorkQueue();
@@ -1211,8 +1188,12 @@ public class StreamingDataflowWorkerTest {
server.setExpectedExceptionCount(1);
StreamingDataflowWorker worker =
- makeWorker(instructions, createTestingPipelineOptions(), true /* publishCounters */);
- worker.setMaxWorkItemCommitBytes(1000);
+ makeWorker(
+ defaultWorkerParams()
+ .setInstructions(instructions)
+ .setMaxWorkItemCommitBytes(1000)
+ .publishCounters()
+ .build());
worker.start();
server
@@ -1290,7 +1271,7 @@ public class StreamingDataflowWorkerTest {
}
StreamingDataflowWorker worker =
- makeWorker(instructions, createTestingPipelineOptions(), true /* publishCounters */);
+ makeWorker(defaultWorkerParams().setInstructions(instructions).publishCounters().build());
worker.start();
Map<Long, Windmill.WorkItemCommitRequest> result = server.waitForAndGetCommits(4);
@@ -1366,7 +1347,7 @@ public class StreamingDataflowWorkerTest {
Collections.singletonList(DEFAULT_WINDOW))));
StreamingDataflowWorker worker =
- makeWorker(instructions, createTestingPipelineOptions(), true /* publishCounters */);
+ makeWorker(defaultWorkerParams().setInstructions(instructions).publishCounters().build());
worker.start();
server.waitForEmptyWorkQueue();
@@ -1471,7 +1452,7 @@ public class StreamingDataflowWorkerTest {
.thenReturn(makeInput(timestamp2, timestamp2));
StreamingDataflowWorker worker =
- makeWorker(instructions, createTestingPipelineOptions(), false /* publishCounters */);
+ makeWorker(defaultWorkerParams().setInstructions(instructions).build());
worker.start();
Map<Long, Windmill.WorkItemCommitRequest> result = server.waitForAndGetCommits(2);
@@ -1589,10 +1570,11 @@ public class StreamingDataflowWorkerTest {
makeSinkInstruction(groupedCoder, 1));
StreamingDataflowWorker worker =
- makeWorker(instructions, createTestingPipelineOptions(), false /* publishCounters */);
- Map<String, String> nameMap = new HashMap<>();
- nameMap.put("MergeWindowsStep", "MergeWindows");
- worker.addStateNameMappings(nameMap);
+ makeWorker(
+ defaultWorkerParams()
+ .setInstructions(instructions)
+ .addStateNameMapping("MergeWindowsStep", "MergeWindows")
+ .build());
worker.start();
server
@@ -1876,10 +1858,11 @@ public class StreamingDataflowWorkerTest {
makeSinkInstruction(groupedCoder, 2));
StreamingDataflowWorker worker =
- makeWorker(instructions, createTestingPipelineOptions(), false /* publishCounters */);
- Map<String, String> nameMap = new HashMap<>();
- nameMap.put("MergeWindowsStep", "MergeWindows");
- worker.addStateNameMappings(nameMap);
+ makeWorker(
+ defaultWorkerParams()
+ .setInstructions(instructions)
+ .addStateNameMapping("MergeWindowsStep", "MergeWindows")
+ .build());
worker.start();
server
@@ -2170,12 +2153,12 @@ public class StreamingDataflowWorkerTest {
makeWindowingSourceInstruction(kvCoder),
mergeWindowsInstruction,
makeSinkInstruction(groupedCoder, 1));
-
StreamingDataflowWorker worker =
- makeWorker(instructions, createTestingPipelineOptions(), false /* publishCounters */);
- Map<String, String> nameMap = new HashMap<>();
- nameMap.put("MergeWindowsStep", "MergeWindows");
- worker.addStateNameMappings(nameMap);
+ makeWorker(
+ defaultWorkerParams()
+ .setInstructions(instructions)
+ .addStateNameMapping("MergeWindowsStep", "MergeWindows")
+ .build());
worker.start();
// Respond to any GetData requests with empty state.
@@ -2320,12 +2303,8 @@ public class StreamingDataflowWorkerTest {
public void testUnboundedSources() throws Exception {
List<Integer> finalizeTracker = Lists.newArrayList();
TestCountingSource.setFinalizeTracker(finalizeTracker);
-
StreamingDataflowWorker worker =
- makeWorker(
- makeUnboundedSourcePipeline(),
- createTestingPipelineOptions(),
- false /* publishCounters */);
+ makeWorker(defaultWorkerParams().setInstructions(makeUnboundedSourcePipeline()).build());
worker.start();
// Test new key.
@@ -2489,9 +2468,10 @@ public class StreamingDataflowWorkerTest {
StreamingDataflowWorker worker =
makeWorker(
- makeUnboundedSourcePipeline(),
- createTestingPipelineOptions(),
- true /* publishCounters */);
+ defaultWorkerParams()
+ .setInstructions(makeUnboundedSourcePipeline())
+ .publishCounters()
+ .build());
worker.start();
// Test new key.
@@ -2600,10 +2580,12 @@ public class StreamingDataflowWorkerTest {
List<Integer> finalizeTracker = Lists.newArrayList();
TestCountingSource.setFinalizeTracker(finalizeTracker);
- DataflowWorkerHarnessOptions options = createTestingPipelineOptions();
- options.setWorkerCacheMb(0); // Disable state cache so it doesn't detect retry.
StreamingDataflowWorker worker =
- makeWorker(makeUnboundedSourcePipeline(), options, false /* publishCounters */);
+ makeWorker(
+ defaultWorkerParams(
+ "--workerCacheMb=0") // Disable state cache so it doesn't detect retry.
+ .setInstructions(makeUnboundedSourcePipeline())
+ .build());
worker.start();
// Test new key.
@@ -3119,8 +3101,7 @@ public class StreamingDataflowWorkerTest {
DataflowPipelineOptions options = createTestingPipelineOptions();
options.setNumWorkers(1);
- DataflowPipelineDebugOptions debugOptions = options.as(DataflowPipelineDebugOptions.class);
- debugOptions.setUnboundedReaderMaxElements(1);
+ options.setUnboundedReaderMaxElements(1);
CloudObject codec =
CloudObjects.asCloudObject(
@@ -3158,10 +3139,12 @@ public class StreamingDataflowWorkerTest {
StreamingDataflowWorker worker =
makeWorker(
- instructions,
- options.as(DataflowWorkerHarnessOptions.class),
- true /* publishCounters */,
- 100);
+ defaultWorkerParams()
+ .setInstructions(instructions)
+ .setOptions(options.as(DataflowWorkerHarnessOptions.class))
+ .setLocalRetryTimeoutMs(100)
+ .publishCounters()
+ .build());
worker.start();
// Three GetData requests
@@ -3300,9 +3283,8 @@ public class StreamingDataflowWorkerTest {
makeSourceInstruction(StringUtf8Coder.of()),
makeDoFnInstruction(new FanoutFn(), 0, StringUtf8Coder.of()),
makeSinkInstruction(StringUtf8Coder.of(), 0));
-
StreamingDataflowWorker worker =
- makeWorker(instructions, createTestingPipelineOptions(), true /* publishCounters */);
+ makeWorker(defaultWorkerParams().setInstructions(instructions).publishCounters().build());
worker.start();
server.whenGetWorkCalled().thenReturn(makeInput(0, TimeUnit.MILLISECONDS.toMicros(0)));
@@ -3319,9 +3301,12 @@ public class StreamingDataflowWorkerTest {
makeDoFnInstruction(new SlowDoFn(), 0, StringUtf8Coder.of()),
makeSinkInstruction(StringUtf8Coder.of(), 0));
- DataflowWorkerHarnessOptions options = createTestingPipelineOptions();
- options.setActiveWorkRefreshPeriodMillis(100);
- StreamingDataflowWorker worker = makeWorker(instructions, options, true /* publishCounters */);
+ StreamingDataflowWorker worker =
+ makeWorker(
+ defaultWorkerParams("--activeWorkRefreshPeriodMillis=100")
+ .setInstructions(instructions)
+ .publishCounters()
+ .build());
worker.start();
server.whenGetWorkCalled().thenReturn(makeInput(0, TimeUnit.MILLISECONDS.toMicros(0)));
@@ -3342,9 +3327,12 @@ public class StreamingDataflowWorkerTest {
makeDoFnInstruction(blockingFn, 0, StringUtf8Coder.of()),
makeSinkInstruction(StringUtf8Coder.of(), 0));
- DataflowWorkerHarnessOptions options = createTestingPipelineOptions();
- options.setActiveWorkRefreshPeriodMillis(100);
- StreamingDataflowWorker worker = makeWorker(instructions, options, true /* publishCounters */);
+ StreamingDataflowWorker worker =
+ makeWorker(
+ defaultWorkerParams("--activeWorkRefreshPeriodMillis=100")
+ .setInstructions(instructions)
+ .publishCounters()
+ .build());
worker.start();
GetWorkResponse workItem =
@@ -3444,18 +3432,17 @@ public class StreamingDataflowWorkerTest {
new FakeSlowDoFn(clock, Duration.millis(1000)), 0, StringUtf8Coder.of()),
makeSinkInstruction(StringUtf8Coder.of(), 0));
- DataflowWorkerHarnessOptions options = createTestingPipelineOptions();
- options.setActiveWorkRefreshPeriodMillis(100);
- // A single-threaded worker processes work sequentially, leaving a second work item in state
- // QUEUED until the first work item is committed.
- options.setNumberOfWorkerHarnessThreads(1);
StreamingDataflowWorker worker =
makeWorker(
- instructions,
- options,
- false /* publishCounters */,
- clock,
- clock::newFakeScheduledExecutor);
+ defaultWorkerParams(
+ "--activeWorkRefreshPeriodMillis=100",
+ // A single-threaded worker processes work sequentially, leaving a second work
+ // item in state QUEUED until the first work item is committed.
+ "--numberOfWorkerHarnessThreads=1")
+ .setInstructions(instructions)
+ .setClock(clock)
+ .setExecutorSupplier(clock::newFakeScheduledExecutor)
+ .build());
worker.start();
ActiveWorkRefreshSink awrSink = new ActiveWorkRefreshSink(EMPTY_DATA_RESPONDER);
@@ -3486,15 +3473,13 @@ public class StreamingDataflowWorkerTest {
new FakeSlowDoFn(clock, Duration.millis(1000)), 0, StringUtf8Coder.of()),
makeSinkInstruction(StringUtf8Coder.of(), 0));
- DataflowWorkerHarnessOptions options = createTestingPipelineOptions();
- options.setActiveWorkRefreshPeriodMillis(100);
StreamingDataflowWorker worker =
makeWorker(
- instructions,
- options,
- false /* publishCounters */,
- clock,
- clock::newFakeScheduledExecutor);
+ defaultWorkerParams("--activeWorkRefreshPeriodMillis=100")
+ .setInstructions(instructions)
+ .setClock(clock)
+ .setExecutorSupplier(clock::newFakeScheduledExecutor)
+ .build());
worker.start();
ActiveWorkRefreshSink awrSink = new ActiveWorkRefreshSink(EMPTY_DATA_RESPONDER);
@@ -3519,15 +3504,13 @@ public class StreamingDataflowWorkerTest {
makeDoFnInstruction(new ReadingDoFn(), 0, StringUtf8Coder.of()),
makeSinkInstruction(StringUtf8Coder.of(), 0));
- DataflowWorkerHarnessOptions options = createTestingPipelineOptions();
- options.setActiveWorkRefreshPeriodMillis(100);
StreamingDataflowWorker worker =
makeWorker(
- instructions,
- options,
- false /* publishCounters */,
- clock,
- clock::newFakeScheduledExecutor);
+ defaultWorkerParams("--activeWorkRefreshPeriodMillis=100")
+ .setInstructions(instructions)
+ .setClock(clock)
+ .setExecutorSupplier(clock::newFakeScheduledExecutor)
+ .build());
worker.start();
// Inject latency on the fake clock when the server receives a GetData call that isn't
@@ -3567,15 +3550,14 @@ public class StreamingDataflowWorkerTest {
clock.sleep(Duration.millis(1000));
return Windmill.CommitWorkResponse.getDefaultInstance();
});
- DataflowWorkerHarnessOptions options = createTestingPipelineOptions();
- options.setActiveWorkRefreshPeriodMillis(100);
+
StreamingDataflowWorker worker =
makeWorker(
- instructions,
- options,
- false /* publishCounters */,
- clock,
- clock::newFakeScheduledExecutor);
+ defaultWorkerParams("--activeWorkRefreshPeriodMillis=100")
+ .setInstructions(instructions)
+ .setClock(clock)
+ .setExecutorSupplier(clock::newFakeScheduledExecutor)
+ .build());
worker.start();
ActiveWorkRefreshSink awrSink = new ActiveWorkRefreshSink(EMPTY_DATA_RESPONDER);
@@ -3602,16 +3584,13 @@ public class StreamingDataflowWorkerTest {
new FakeSlowDoFn(clock, Duration.millis(dofnWaitTimeMs)), 0, StringUtf8Coder.of()),
makeSinkInstruction(StringUtf8Coder.of(), 0));
- DataflowWorkerHarnessOptions options = createTestingPipelineOptions();
- options.setActiveWorkRefreshPeriodMillis(100);
- options.setNumberOfWorkerHarnessThreads(1);
StreamingDataflowWorker worker =
makeWorker(
- instructions,
- options,
- false /* publishCounters */,
- clock,
- clock::newFakeScheduledExecutor);
+ defaultWorkerParams("--activeWorkRefreshPeriodMillis=100")
+ .setInstructions(instructions)
+ .setClock(clock)
+ .setExecutorSupplier(clock::newFakeScheduledExecutor)
+ .build());
worker.start();
ActiveWorkRefreshSink awrSink = new ActiveWorkRefreshSink(EMPTY_DATA_RESPONDER);
@@ -3654,9 +3633,12 @@ public class StreamingDataflowWorkerTest {
makeDoFnInstruction(new SlowDoFn(), 0, StringUtf8Coder.of()),
makeSinkInstruction(StringUtf8Coder.of(), 0));
- DataflowWorkerHarnessOptions options = createTestingPipelineOptions();
- options.setActiveWorkRefreshPeriodMillis(100);
- StreamingDataflowWorker worker = makeWorker(instructions, options, true /* publishCounters */);
+ StreamingDataflowWorker worker =
+ makeWorker(
+ defaultWorkerParams("--activeWorkRefreshPeriodMillis=100")
+ .setInstructions(instructions)
+ .publishCounters()
+ .build());
worker.start();
server.whenGetWorkCalled().thenReturn(makeInput(0, TimeUnit.MILLISECONDS.toMicros(0)));
@@ -3690,9 +3672,12 @@ public class StreamingDataflowWorkerTest {
makeDoFnInstruction(new SlowDoFn(), 0, StringUtf8Coder.of()),
makeSinkInstruction(StringUtf8Coder.of(), 0));
- DataflowWorkerHarnessOptions options = createTestingPipelineOptions();
- options.setActiveWorkRefreshPeriodMillis(10);
- StreamingDataflowWorker worker = makeWorker(instructions, options, true /* publishCounters */);
+ StreamingDataflowWorker worker =
+ makeWorker(
+ defaultWorkerParams("--activeWorkRefreshPeriodMillis=10")
+ .setInstructions(instructions)
+ .publishCounters()
+ .build());
worker.start();
server.whenGetWorkCalled().thenReturn(makeInput(0, TimeUnit.MILLISECONDS.toMicros(0)));
@@ -3728,10 +3713,11 @@ public class StreamingDataflowWorkerTest {
StreamingDataflowWorker worker =
makeWorker(
- makeUnboundedSourcePipeline(
- numMessagesInCustomSourceShard, new InflateDoFn(inflatedSizePerMessage)),
- createTestingPipelineOptions(),
- false /* publishCounters */);
+ defaultWorkerParams()
+ .setInstructions(
+ makeUnboundedSourcePipeline(
+ numMessagesInCustomSourceShard, new InflateDoFn(inflatedSizePerMessage)))
+ .build());
worker.start();
// Test new key.
@@ -3811,9 +3797,8 @@ public class StreamingDataflowWorkerTest {
StringUtf8Coder.of(),
1,
GlobalWindow.Coder.INSTANCE));
-
StreamingDataflowWorker worker =
- makeWorker(instructions, createTestingPipelineOptions(), true /* publishCounters */);
+ makeWorker(defaultWorkerParams().setInstructions(instructions).publishCounters().build());
worker.start();
// Test new key.
@@ -3879,9 +3864,12 @@ public class StreamingDataflowWorkerTest {
makeSourceInstruction(StringUtf8Coder.of()),
makeSinkInstruction(StringUtf8Coder.of(), 0));
- DataflowWorkerHarnessOptions options = createTestingPipelineOptions();
- options.setStuckCommitDurationMillis(2000);
- StreamingDataflowWorker worker = makeWorker(instructions, options, true /* publishCounters */);
+ StreamingDataflowWorker worker =
+ makeWorker(
+ defaultWorkerParams("--stuckCommitDurationMillis=2000")
+ .setInstructions(instructions)
+ .publishCounters()
+ .build());
worker.start();
// Prevent commit callbacks from being called to simulate a stuck commit.
server.setDropStreamingCommits(true);
@@ -3917,9 +3905,13 @@ public class StreamingDataflowWorkerTest {
Arrays.asList(
makeSourceInstruction(StringUtf8Coder.of()),
makeSinkInstruction(StringUtf8Coder.of(), 0));
- DataflowWorkerHarnessOptions options = createTestingPipelineOptions();
- options.setWindmillServiceCommitThreads(configNumCommitThreads);
- StreamingDataflowWorker worker = makeWorker(instructions, options, true /* publishCounters */);
+
+ StreamingDataflowWorker worker =
+ makeWorker(
+ defaultWorkerParams("--windmillServiceCommitThreads=" + configNumCommitThreads)
+ .setInstructions(instructions)
+ .publishCounters()
+ .build());
worker.start();
assertEquals(expectedNumCommitThreads, worker.numCommitThreads());
worker.stop();
@@ -4412,4 +4404,71 @@ public class StreamingDataflowWorkerTest {
c.output(new String(chars));
}
}
+
+ @AutoValue
+ abstract static class StreamingDataflowWorkerTestParams {
+
+ private static StreamingDataflowWorkerTestParams.Builder builder() {
+ return new AutoValue_StreamingDataflowWorkerTest_StreamingDataflowWorkerTestParams.Builder()
+ .addStateNameMapping(DEFAULT_PARDO_USER_NAME, DEFAULT_PARDO_STATE_FAMILY)
+ .setExecutorSupplier(ignored -> Executors.newSingleThreadScheduledExecutor())
+ .setLocalRetryTimeoutMs(-1)
+ .setPublishCounters(false)
+ .setClock(Instant::now)
+ .setMaxWorkItemCommitBytes(Integer.MAX_VALUE);
+ }
+
+ abstract ImmutableMap<String, String> stateNameMappings();
+
+ abstract List<ParallelInstruction> instructions();
+
+ abstract DataflowWorkerHarnessOptions options();
+
+ abstract boolean publishCounters();
+
+ abstract Supplier<Instant> clock();
+
+ abstract Function<String, ScheduledExecutorService> executorSupplier();
+
+ abstract int localRetryTimeoutMs();
+
+ abstract int maxWorkItemCommitBytes();
+
+ @AutoValue.Builder
+ abstract static class Builder {
+ abstract Builder setStateNameMappings(ImmutableMap<String, String> value);
+
+ abstract ImmutableMap.Builder<String, String> stateNameMappingsBuilder();
+
+ final Builder addStateNameMapping(String username, String stateFamilyName) {
+ stateNameMappingsBuilder().put(username, stateFamilyName);
+ return this;
+ }
+
+ final Builder addAllStateNameMappings(Map<String, String> stateNameMappings) {
+ stateNameMappingsBuilder().putAll(stateNameMappings);
+ return this;
+ }
+
+ abstract Builder setInstructions(List<ParallelInstruction> value);
+
+ abstract Builder setOptions(DataflowWorkerHarnessOptions value);
+
+ abstract Builder setPublishCounters(boolean value);
+
+ final Builder publishCounters() {
+ return setPublishCounters(true);
+ }
+
+ abstract Builder setClock(Supplier<Instant> value);
+
+ abstract Builder setExecutorSupplier(Function<String, ScheduledExecutorService> value);
+
+ abstract Builder setLocalRetryTimeoutMs(int value);
+
+ abstract Builder setMaxWorkItemCommitBytes(int maxWorkItemCommitBytes);
+
+ abstract StreamingDataflowWorkerTestParams build();
+ }
+ }
}
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationStateCacheTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationStateCacheTest.java
new file mode 100644
index 00000000000..3eb434c3903
--- /dev/null
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationStateCacheTest.java
@@ -0,0 +1,282 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.streaming;
+
+import static com.google.common.truth.Truth.assertThat;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoInteractions;
+import static org.mockito.Mockito.when;
+
+import com.google.api.services.dataflow.model.MapTask;
+import java.util.Collections;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.stream.Collectors;
+import org.apache.beam.runners.dataflow.worker.streaming.config.ComputationConfig;
+import org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill;
+import org.apache.beam.runners.dataflow.worker.windmill.state.WindmillStateCache;
+import org.apache.beam.runners.dataflow.worker.windmill.work.budget.GetWorkBudget;
+import org.apache.beam.sdk.fn.IdGenerators;
+import org.apache.beam.vendor.grpc.v1p60p1.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet;
+import org.joda.time.Instant;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class ComputationStateCacheTest {
+ private final BoundedQueueExecutor workExecutor = mock(BoundedQueueExecutor.class);
+ private final WindmillStateCache.ForComputation stateCache =
+ mock(WindmillStateCache.ForComputation.class);
+ private final ComputationConfig.Fetcher configFetcher = mock(ComputationConfig.Fetcher.class);
+ private ComputationStateCache computationStateCache;
+
+ private static Work createWork(long workToken, long cacheToken) {
+ return Work.create(
+ Windmill.WorkItem.newBuilder()
+ .setKey(ByteString.copyFromUtf8(""))
+ .setShardingKey(1)
+ .setWorkToken(workToken)
+ .setCacheToken(cacheToken)
+ .build(),
+ Instant::now,
+ Collections.emptyList(),
+ unused -> {});
+ }
+
+ @Before
+ public void setUp() {
+ computationStateCache =
+ ComputationStateCache.create(
+ configFetcher, workExecutor, ignored -> stateCache, IdGenerators.decrementingLongs());
+ }
+
+ @Test
+ public void testGet_computationStateNotCached() {
+ String computationId = "computationId";
+ MapTask mapTask = new MapTask().setStageName("stageName").setSystemName("systemName");
+ Map<String, String> userTransformToStateFamilyName =
+ ImmutableMap.of("userTransformName", "stateFamilyName");
+ ComputationConfig computationConfig =
+ ComputationConfig.create(mapTask, userTransformToStateFamilyName, ImmutableMap.of());
+ when(configFetcher.fetchConfig(eq(computationId))).thenReturn(Optional.of(computationConfig));
+ Optional<ComputationState> computationState = computationStateCache.get(computationId);
+ assertTrue(computationState.isPresent());
+ assertThat(computationState.get().getComputationId()).isEqualTo(computationId);
+ assertThat(computationState.get().getMapTask()).isEqualTo(mapTask);
+ assertThat(computationState.get().getTransformUserNameToStateFamily())
+ .isEqualTo(userTransformToStateFamilyName);
+ }
+
+ @Test
+ public void testGet_computationStateCached() {
+ String computationId = "computationId";
+ MapTask mapTask = new MapTask().setStageName("stageName").setSystemName("systemName");
+ Map<String, String> userTransformToStateFamilyName =
+ ImmutableMap.of("userTransformName", "stateFamilyName");
+ ComputationConfig computationConfig =
+ ComputationConfig.create(mapTask, userTransformToStateFamilyName, ImmutableMap.of());
+ when(configFetcher.fetchConfig(eq(computationId))).thenReturn(Optional.of(computationConfig));
+ computationStateCache.get(computationId);
+ Optional<ComputationState> computationState = computationStateCache.get(computationId);
+ assertTrue(computationState.isPresent());
+ assertThat(computationState.get().getComputationId()).isEqualTo(computationId);
+ assertThat(computationState.get().getMapTask()).isEqualTo(mapTask);
+ assertThat(computationState.get().getTransformUserNameToStateFamily())
+ .isEqualTo(userTransformToStateFamilyName);
+ }
+
+ @Test
+ public void testGet_computationStateNotCachedOrFetchable() {
+ String computationId = "computationId";
+ when(configFetcher.fetchConfig(eq(computationId))).thenReturn(Optional.empty());
+ Optional<ComputationState> computationState = computationStateCache.get(computationId);
+ assertFalse(computationState.isPresent());
+ // Fetch again to make sure we call configFetcher and nulls/empty are not cached.
+ Optional<ComputationState> computationState2 = computationStateCache.get(computationId);
+ assertFalse(computationState2.isPresent());
+ verify(configFetcher, times(2)).fetchConfig(eq(computationId));
+ }
+
+ @Test
+ public void testGet_computationStateConfigFetcherFailure() {
+ String computationId = "computationId";
+ when(configFetcher.fetchConfig(eq(computationId))).thenThrow(new RuntimeException());
+ Optional<ComputationState> computationState = computationStateCache.get(computationId);
+ assertFalse(computationState.isPresent());
+ // Fetch again to make sure we call configFetcher and nulls/empty are not cached.
+ Optional<ComputationState> computationState2 = computationStateCache.get(computationId);
+ assertFalse(computationState2.isPresent());
+ verify(configFetcher, times(2)).fetchConfig(eq(computationId));
+ }
+
+ @Test
+ public void testGet_usesUserTransformToStateFamilyNameIfNotEmpty() {
+ String computationId = "computationId";
+ MapTask mapTask = new MapTask().setStageName("stageName").setSystemName("systemName");
+ Map<String, String> userTransformToStateFamilyName =
+ ImmutableMap.of("userTransformName", "stateFamilyName");
+ ComputationConfig computationConfig =
+ ComputationConfig.create(
+ mapTask, userTransformToStateFamilyName, ImmutableMap.of("stateName1", "stateName"));
+ when(configFetcher.fetchConfig(eq(computationId))).thenReturn(Optional.of(computationConfig));
+ Optional<ComputationState> computationState = computationStateCache.get(computationId);
+ assertTrue(computationState.isPresent());
+ assertThat(computationState.get().getComputationId()).isEqualTo(computationId);
+ assertThat(computationState.get().getMapTask()).isEqualTo(mapTask);
+ assertThat(computationState.get().getTransformUserNameToStateFamily())
+ .isEqualTo(userTransformToStateFamilyName);
+ }
+
+ @Test
+ public void testGet_defaultsToStateNameMapWhenUserTransformToStateFamilyNameEmpty() {
+ String computationId = "computationId";
+ MapTask mapTask = new MapTask().setStageName("stageName").setSystemName("systemName");
+ ComputationConfig computationConfig =
+ ComputationConfig.create(
+ mapTask, ImmutableMap.of(), ImmutableMap.of("stateName1", "stateName"));
+ when(configFetcher.fetchConfig(eq(computationId))).thenReturn(Optional.of(computationConfig));
+ Optional<ComputationState> computationState = computationStateCache.get(computationId);
+ assertTrue(computationState.isPresent());
+ assertThat(computationState.get().getComputationId()).isEqualTo(computationId);
+ assertThat(computationState.get().getMapTask()).isEqualTo(mapTask);
+ assertThat(computationState.get().getTransformUserNameToStateFamily())
+ .isEqualTo(computationStateCache.getGlobalUsernameToStateFamilyNameMap());
+ }
+
+ @Test
+ public void testGet_buildsStateNameMap() {
+ String computationId = "computationId";
+ MapTask mapTask = new MapTask().setStageName("stageName").setSystemName("systemName");
+ Map<String, String> stateNameMap = ImmutableMap.of("stateName1", "stateName");
+ ComputationConfig computationConfig = ComputationConfig.create(mapTask, null, stateNameMap);
+ when(configFetcher.fetchConfig(eq(computationId))).thenReturn(Optional.of(computationConfig));
+ computationStateCache.get(computationId);
+ assertThat(computationStateCache.getGlobalUsernameToStateFamilyNameMap())
+ .containsExactlyEntriesIn(stateNameMap);
+ }
+
+ @Test
+ public void testGetIfPresent_computationStateCached() {
+ String computationId = "computationId";
+ MapTask mapTask = new MapTask().setStageName("stageName").setSystemName("systemName");
+ Map<String, String> userTransformToStateFamilyName =
+ ImmutableMap.of("userTransformName", "stateFamilyName");
+ ComputationConfig computationConfig =
+ ComputationConfig.create(mapTask, userTransformToStateFamilyName, ImmutableMap.of());
+ when(configFetcher.fetchConfig(eq(computationId))).thenReturn(Optional.of(computationConfig));
+ computationStateCache.get(computationId);
+ Optional<ComputationState> computationState = computationStateCache.getIfPresent(computationId);
+ assertTrue(computationState.isPresent());
+ assertThat(computationState.get().getComputationId()).isEqualTo(computationId);
+ assertThat(computationState.get().getMapTask()).isEqualTo(mapTask);
+ assertThat(computationState.get().getTransformUserNameToStateFamily())
+ .isEqualTo(userTransformToStateFamilyName);
+ }
+
+ @Test
+ public void testGetIfPresent_computationStateNotCached() {
+ Optional<ComputationState> computationState =
+ computationStateCache.getIfPresent("computationId");
+ assertFalse(computationState.isPresent());
+ verifyNoInteractions(configFetcher);
+ }
+
+ @Test
+ public void testGetAllPresentComputations() {
+ String computationId1 = "computationId1";
+ String computationId2 = "computationId2";
+ MapTask mapTask = new MapTask().setStageName("stageName").setSystemName("systemName");
+ Map<String, String> userTransformToStateFamilyName =
+ ImmutableMap.of("userTransformName", "stateFamilyName");
+ ComputationConfig computationConfig =
+ ComputationConfig.create(mapTask, userTransformToStateFamilyName, ImmutableMap.of());
+ when(configFetcher.fetchConfig(eq(computationId1))).thenReturn(Optional.of(computationConfig));
+ when(configFetcher.fetchConfig(eq(computationId2))).thenReturn(Optional.of(computationConfig));
+
+ computationStateCache.get(computationId1);
+ computationStateCache.get(computationId2);
+ Set<String> expectedComputationIds = ImmutableSet.of(computationId1, computationId2);
+ Set<String> actualComputationIds =
+ computationStateCache.getAllPresentComputations().stream()
+ .map(ComputationState::getComputationId)
+ .collect(Collectors.toSet());
+ assertThat(actualComputationIds).containsExactlyElementsIn(expectedComputationIds);
+ computationStateCache
+ .getAllPresentComputations()
+ .forEach(
+ computationState -> {
+ assertThat(expectedComputationIds).contains(computationState.getComputationId());
+ assertThat(computationState.getMapTask()).isEqualTo(mapTask);
+ assertThat(computationState.getTransformUserNameToStateFamily())
+ .isEqualTo(userTransformToStateFamilyName);
+ });
+ }
+
+ @Test
+ public void testTotalCurrentActiveGetWorkBudget() {
+ String computationId = "computationId";
+ String computationId2 = "computationId2";
+ MapTask mapTask = new MapTask().setStageName("stageName").setSystemName("systemName");
+ Map<String, String> userTransformToStateFamilyName =
+ ImmutableMap.of("userTransformName", "stateFamilyName");
+ ComputationConfig computationConfig =
+ ComputationConfig.create(mapTask, userTransformToStateFamilyName, ImmutableMap.of());
+ when(configFetcher.fetchConfig(eq(computationId))).thenReturn(Optional.of(computationConfig));
+ when(configFetcher.fetchConfig(eq(computationId2))).thenReturn(Optional.of(computationConfig));
+ Work work1 = createWork(1, 1);
+ Work work2 = createWork(2, 2);
+ Work work3 = createWork(3, 3);
+
+ // Activate 1 Work for computationId
+ Optional<ComputationState> maybeComputationState = computationStateCache.get(computationId);
+ assertTrue(maybeComputationState.isPresent());
+ ComputationState computationState = maybeComputationState.get();
+ ShardedKey shardedKey = ShardedKey.create(ByteString.EMPTY, 1);
+ computationState.activateWork(shardedKey, work1);
+
+ // Activate 2 Work(s) for computationId2
+ Optional<ComputationState> maybeComputationState2 = computationStateCache.get(computationId);
+ assertTrue(maybeComputationState2.isPresent());
+ ComputationState computationState2 = maybeComputationState2.get();
+ ShardedKey shardedKey2 = ShardedKey.create(ByteString.EMPTY, 2);
+ computationState2.activateWork(shardedKey2, work2);
+ computationState2.activateWork(shardedKey2, work3);
+
+ // GetWorkBudget should have 3 items. 1 from computationId, 2 from computationId2.
+ assertThat(computationStateCache.totalCurrentActiveGetWorkBudget())
+ .isEqualTo(
+ GetWorkBudget.builder()
+ .setItems(3)
+ .setBytes(
+ work1.getWorkItem().getSerializedSize()
+ + work2.getWorkItem().getSerializedSize()
+ + work3.getWorkItem().getSerializedSize())
+ .build());
+ }
+}
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/config/StreamingApplianceComputationConfigFetcherTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/config/StreamingApplianceComputationConfigFetcherTest.java
new file mode 100644
index 00000000000..f39c98c61b1
--- /dev/null
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/config/StreamingApplianceComputationConfigFetcherTest.java
@@ -0,0 +1,142 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.streaming.config;
+
+import static com.google.common.truth.Truth.assertThat;
+import static java.util.stream.Collectors.toMap;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import com.google.api.services.dataflow.model.MapTask;
+import java.io.IOException;
+import java.util.List;
+import java.util.NoSuchElementException;
+import java.util.Optional;
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill;
+import org.apache.beam.runners.dataflow.worker.windmill.WindmillServerStub;
+import org.apache.beam.sdk.extensions.gcp.util.Transport;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class StreamingApplianceComputationConfigFetcherTest {
+ private final WindmillServerStub mockWindmillServer = mock(WindmillServerStub.class);
+
+ @Test
+ public void testGetComputationConfig() throws IOException {
+ List<Windmill.GetConfigResponse.NameMapEntry> nameMapEntries =
+ ImmutableList.of(
+ Windmill.GetConfigResponse.NameMapEntry.newBuilder()
+ .setUserName("userName1")
+ .setSystemName("systemName1")
+ .build(),
+ Windmill.GetConfigResponse.NameMapEntry.newBuilder()
+ .setUserName("userName2")
+ .setSystemName("userName2")
+ .build());
+ String serializedMapTask =
+ Transport.getJsonFactory()
+ .toString(
+ new MapTask()
+ .setSystemName("systemName")
+ .setStageName("stageName")
+ .setInstructions(ImmutableList.of()));
+ Windmill.GetConfigResponse getConfigResponse =
+ Windmill.GetConfigResponse.newBuilder()
+ .addAllNameMap(nameMapEntries)
+ .addComputationConfigMap(
+ Windmill.GetConfigResponse.ComputationConfigMapEntry.newBuilder()
+ .setComputationId("systemName")
+ .setComputationConfig(
+ Windmill.ComputationConfig.newBuilder()
+ .addTransformUserNameToStateFamily(
+ Windmill.ComputationConfig.TransformUserNameToStateFamilyEntry
+ .newBuilder()
+ .setStateFamily("stateFamilyName")
+ .setTransformUserName("transformUserName")
+ .build())
+ .build())
+ .build())
+ .addCloudWorks(serializedMapTask)
+ .build();
+ ComputationConfig expectedConfig =
+ ComputationConfig.create(
+ Transport.getJsonFactory().fromString(serializedMapTask, MapTask.class),
+ getConfigResponse.getComputationConfigMapList().stream()
+ .map(Windmill.GetConfigResponse.ComputationConfigMapEntry::getComputationConfig)
+ .flatMap(
+ computationConfig ->
+ computationConfig.getTransformUserNameToStateFamilyList().stream())
+ .collect(
+ toMap(
+ Windmill.ComputationConfig.TransformUserNameToStateFamilyEntry
+ ::getTransformUserName,
+ Windmill.ComputationConfig.TransformUserNameToStateFamilyEntry
+ ::getStateFamily)),
+ nameMapEntries.stream()
+ .collect(
+ toMap(
+ Windmill.GetConfigResponse.NameMapEntry::getUserName,
+ Windmill.GetConfigResponse.NameMapEntry::getSystemName)));
+ StreamingApplianceComputationConfigFetcher configLoader =
+ createStreamingApplianceConfigLoader();
+ when(mockWindmillServer.getConfig(any())).thenReturn(getConfigResponse);
+ Optional<ComputationConfig> config = configLoader.fetchConfig("someComputationId");
+ assertTrue(config.isPresent());
+ assertThat(config.get()).isEqualTo(expectedConfig);
+ }
+
+ @Test
+ public void testGetComputationConfig_whenNoConfigReturned() {
+ StreamingApplianceComputationConfigFetcher configLoader =
+ createStreamingApplianceConfigLoader();
+ when(mockWindmillServer.getConfig(any())).thenReturn(null);
+ Optional<ComputationConfig> configResponse = configLoader.fetchConfig("someComputationId");
+ assertFalse(configResponse.isPresent());
+ }
+
+ @Test
+ public void testGetComputationConfig_errorOnNoComputationConfig() {
+ StreamingApplianceComputationConfigFetcher configLoader =
+ createStreamingApplianceConfigLoader();
+ when(mockWindmillServer.getConfig(any()))
+ .thenReturn(Windmill.GetConfigResponse.newBuilder().build());
+ assertThrows(NoSuchElementException.class, () -> configLoader.fetchConfig("someComputationId"));
+ }
+
+ @Test
+ public void testGetComputationConfig_onFetchConfigError() {
+ StreamingApplianceComputationConfigFetcher configLoader =
+ createStreamingApplianceConfigLoader();
+ RuntimeException e = new RuntimeException("something bad happened.");
+ when(mockWindmillServer.getConfig(any())).thenThrow(e);
+ Throwable fetchConfigError =
+ assertThrows(RuntimeException.class, () -> configLoader.fetchConfig("someComputationId"));
+ assertThat(fetchConfigError).isSameInstanceAs(e);
+ }
+
+ private StreamingApplianceComputationConfigFetcher createStreamingApplianceConfigLoader() {
+ return new StreamingApplianceComputationConfigFetcher(mockWindmillServer::getConfig);
+ }
+}
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/config/StreamingEngineComputationConfigFetcherTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/config/StreamingEngineComputationConfigFetcherTest.java
new file mode 100644
index 00000000000..59fd092adcb
--- /dev/null
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/config/StreamingEngineComputationConfigFetcherTest.java
@@ -0,0 +1,219 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.streaming.config;
+
+import static com.google.common.truth.Truth.assertThat;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import com.google.api.services.dataflow.model.StreamingComputationConfig;
+import com.google.api.services.dataflow.model.StreamingConfigTask;
+import com.google.api.services.dataflow.model.WorkItem;
+import java.io.IOException;
+import java.util.HashSet;
+import java.util.Optional;
+import java.util.Set;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.Executors;
+import java.util.function.Consumer;
+import org.apache.beam.runners.dataflow.worker.WorkUnitClient;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
+import org.joda.time.Duration;
+import org.junit.After;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.internal.stubbing.answers.Returns;
+
+@RunWith(JUnit4.class)
+public class StreamingEngineComputationConfigFetcherTest {
+ private final WorkUnitClient mockDataflowServiceClient =
+ mock(WorkUnitClient.class, new Returns(Optional.empty()));
+ private StreamingEngineComputationConfigFetcher streamingEngineConfigFetcher;
+
+ private StreamingEngineComputationConfigFetcher createConfigFetcher(
+ boolean waitForInitialConfig,
+ long globalConfigRefreshPeriod,
+ Consumer<StreamingEnginePipelineConfig> onPipelineConfig) {
+ return StreamingEngineComputationConfigFetcher.forTesting(
+ !waitForInitialConfig,
+ globalConfigRefreshPeriod,
+ mockDataflowServiceClient,
+ ignored -> Executors.newSingleThreadScheduledExecutor(),
+ onPipelineConfig);
+ }
+
+ @After
+ public void cleanUp() {
+ streamingEngineConfigFetcher.stop();
+ }
+
+ @Test
+ public void testStart_requiresInitialConfig() throws IOException, InterruptedException {
+ WorkItem initialConfig =
+ new WorkItem()
+ .setJobId("job")
+ .setStreamingConfigTask(new StreamingConfigTask().setMaxWorkItemCommitBytes(10L));
+ CountDownLatch waitForInitialConfig = new CountDownLatch(1);
+ Set<StreamingEnginePipelineConfig> receivedPipelineConfig = new HashSet<>();
+ when(mockDataflowServiceClient.getGlobalStreamingConfigWorkItem())
+ .thenReturn(Optional.of(initialConfig));
+ streamingEngineConfigFetcher =
+ createConfigFetcher(
+ /* waitForInitialConfig= */ true,
+ 0,
+ config -> {
+ try {
+ receivedPipelineConfig.add(config);
+ waitForInitialConfig.await();
+ } catch (InterruptedException e) {
+ throw new RuntimeException(e);
+ }
+ });
+ Thread asyncStartConfigLoader = new Thread(streamingEngineConfigFetcher::start);
+ asyncStartConfigLoader.start();
+ waitForInitialConfig.countDown();
+ asyncStartConfigLoader.join();
+ assertThat(receivedPipelineConfig)
+ .containsExactly(
+ StreamingEnginePipelineConfig.builder()
+ .setMaxWorkItemCommitBytes(
+ initialConfig.getStreamingConfigTask().getMaxWorkItemCommitBytes())
+ .build());
+ }
+
+ @Test
+ public void testStart_startsPeriodicConfigRequests() throws IOException, InterruptedException {
+ WorkItem firstConfig =
+ new WorkItem()
+ .setJobId("job")
+ .setStreamingConfigTask(new StreamingConfigTask().setMaxWorkItemCommitBytes(10L));
+ WorkItem secondConfig =
+ new WorkItem()
+ .setJobId("job")
+ .setStreamingConfigTask(new StreamingConfigTask().setMaxWorkItemCommitBytes(15L));
+ WorkItem thirdConfig =
+ new WorkItem()
+ .setJobId("job")
+ .setStreamingConfigTask(new StreamingConfigTask().setMaxWorkItemCommitBytes(100L));
+ CountDownLatch numExpectedRefreshes = new CountDownLatch(3);
+ Set<StreamingEnginePipelineConfig> receivedPipelineConfig = new HashSet<>();
+ when(mockDataflowServiceClient.getGlobalStreamingConfigWorkItem())
+ .thenReturn(Optional.of(firstConfig))
+ .thenReturn(Optional.of(secondConfig))
+ // ConfigFetcher should still fetch subsequent configs on error.
+ .thenThrow(new IOException("something bad happened."))
+ .thenReturn(Optional.of(thirdConfig))
+ // ConfigFetcher should not do anything with a config that doesn't contain a
+ // StreamingConfigTask.
+ .thenReturn(Optional.of(new WorkItem().setJobId("jobId")));
+
+ streamingEngineConfigFetcher =
+ createConfigFetcher(
+ /* waitForInitialConfig= */ true,
+ Duration.millis(100).getMillis(),
+ config -> {
+ receivedPipelineConfig.add(config);
+ numExpectedRefreshes.countDown();
+ });
+
+ Thread asyncStartConfigLoader = new Thread(streamingEngineConfigFetcher::start);
+ asyncStartConfigLoader.start();
+ numExpectedRefreshes.await();
+ asyncStartConfigLoader.join();
+ assertThat(receivedPipelineConfig)
+ .containsExactly(
+ StreamingEnginePipelineConfig.builder()
+ .setMaxWorkItemCommitBytes(
+ firstConfig.getStreamingConfigTask().getMaxWorkItemCommitBytes())
+ .build(),
+ StreamingEnginePipelineConfig.builder()
+ .setMaxWorkItemCommitBytes(
+ secondConfig.getStreamingConfigTask().getMaxWorkItemCommitBytes())
+ .build(),
+ StreamingEnginePipelineConfig.builder()
+ .setMaxWorkItemCommitBytes(
+ thirdConfig.getStreamingConfigTask().getMaxWorkItemCommitBytes())
+ .build());
+ }
+
+ @Test
+ public void testGetComputationConfig() throws IOException {
+ streamingEngineConfigFetcher =
+ createConfigFetcher(/* waitForInitialConfig= */ false, 0, ignored -> {});
+ String computationId = "computationId";
+ String stageName = "stageName";
+ String systemName = "systemName";
+ StreamingComputationConfig pipelineConfig =
+ new StreamingComputationConfig()
+ .setComputationId(computationId)
+ .setStageName(stageName)
+ .setSystemName(systemName)
+ .setInstructions(ImmutableList.of());
+
+ WorkItem workItem =
+ new WorkItem()
+ .setStreamingConfigTask(
+ new StreamingConfigTask()
+ .setStreamingComputationConfigs(ImmutableList.of(pipelineConfig)));
+
+ when(mockDataflowServiceClient.getStreamingConfigWorkItem(anyString()))
+ .thenReturn(Optional.of(workItem));
+ Optional<ComputationConfig> actualPipelineConfig =
+ streamingEngineConfigFetcher.fetchConfig(computationId);
+
+ assertTrue(actualPipelineConfig.isPresent());
+ assertThat(actualPipelineConfig.get())
+ .isEqualTo(
+ ComputationConfig.create(
+ StreamingEngineComputationConfigFetcher.createMapTask(pipelineConfig),
+ ImmutableMap.of(),
+ ImmutableMap.of()));
+ }
+
+ @Test
+ public void testGetComputationConfig_noComputationPresent() throws IOException {
+ Set<StreamingEnginePipelineConfig> receivedPipelineConfig = new HashSet<>();
+ streamingEngineConfigFetcher =
+ createConfigFetcher(/* waitForInitialConfig= */ false, 0, receivedPipelineConfig::add);
+ when(mockDataflowServiceClient.getStreamingConfigWorkItem(anyString()))
+ .thenReturn(Optional.empty());
+ Optional<ComputationConfig> pipelineConfig =
+ streamingEngineConfigFetcher.fetchConfig("someComputationId");
+ assertFalse(pipelineConfig.isPresent());
+ assertThat(receivedPipelineConfig).isEmpty();
+ }
+
+ @Test
+ public void testGetComputationConfig_fetchConfigFromDataflowError() throws IOException {
+ streamingEngineConfigFetcher =
+ createConfigFetcher(/* waitForInitialConfig= */ false, 0, ignored -> {});
+ RuntimeException e = new RuntimeException("something bad happened.");
+ when(mockDataflowServiceClient.getStreamingConfigWorkItem(anyString())).thenThrow(e);
+ Throwable fetchConfigError =
+ assertThrows(
+ RuntimeException.class,
+ () -> streamingEngineConfigFetcher.fetchConfig("someComputationId"));
+ assertThat(fetchConfigError).isSameInstanceAs(e);
+ }
+}