Google Git
Sign in
apache / beam / d140b908701fda9fc4c6569f338db68adf8ea23b / . / runners / google-cloud-dataflow-java / worker / src / main / java / org / apache / beam / runners / dataflow / worker / StreamingDataflowWorker.java
blob: f2b0b27943ae9bc23b328e4b5c7766f002563727 [file] [log] [blame]
/*
* 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;
import static java.nio.charset.StandardCharsets.UTF_8;
import static org.apache.beam.runners.dataflow.DataflowRunner.hasExperiment;
import static org.apache.beam.runners.dataflow.worker.DataflowSystemMetrics.THROTTLING_MSECS_METRIC_NAME;
import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
import com.google.api.services.dataflow.model.CounterStructuredName;
import com.google.api.services.dataflow.model.CounterUpdate;
import com.google.api.services.dataflow.model.MapTask;
import com.google.api.services.dataflow.model.Status;
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 com.google.api.services.dataflow.model.WorkItemStatus;
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.io.File;
import java.io.IOException;
import java.io.PrintWriter;
import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.Deque;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Queue;
import java.util.Random;
import java.util.Timer;
import java.util.TimerTask;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.Semaphore;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.Function;
import javax.annotation.Nullable;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.RemoteGrpcPort;
import org.apache.beam.model.pipeline.v1.RunnerApi;
import org.apache.beam.runners.core.metrics.ExecutionStateSampler;
import org.apache.beam.runners.core.metrics.ExecutionStateTracker;
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.DataflowSystemMetrics.StreamingPerStageSystemCounterNames;
import org.apache.beam.runners.dataflow.worker.DataflowSystemMetrics.StreamingSystemCounterNames;
import org.apache.beam.runners.dataflow.worker.SdkHarnessRegistry.SdkWorkerHarness;
import org.apache.beam.runners.dataflow.worker.StreamingDataflowWorker.Work.State;
import org.apache.beam.runners.dataflow.worker.StreamingModeExecutionContext.StreamingModeExecutionStateRegistry;
import org.apache.beam.runners.dataflow.worker.apiary.FixMultiOutputInfosOnParDoInstructions;
import org.apache.beam.runners.dataflow.worker.counters.Counter;
import org.apache.beam.runners.dataflow.worker.counters.CounterSet;
import org.apache.beam.runners.dataflow.worker.counters.CounterUpdateAggregators;
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.CloneAmbiguousFlattensFunction;
import org.apache.beam.runners.dataflow.worker.graph.CreateRegisterFnOperationFunction;
import org.apache.beam.runners.dataflow.worker.graph.DeduceFlattenLocationsFunction;
import org.apache.beam.runners.dataflow.worker.graph.DeduceNodeLocationsFunction;
import org.apache.beam.runners.dataflow.worker.graph.Edges.Edge;
import org.apache.beam.runners.dataflow.worker.graph.InsertFetchAndFilterStreamingSideInputNodes;
import org.apache.beam.runners.dataflow.worker.graph.LengthPrefixUnknownCoders;
import org.apache.beam.runners.dataflow.worker.graph.MapTaskToNetworkFunction;
import org.apache.beam.runners.dataflow.worker.graph.Networks;
import org.apache.beam.runners.dataflow.worker.graph.Nodes.InstructionOutputNode;
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.graph.Nodes.RemoteGrpcPortNode;
import org.apache.beam.runners.dataflow.worker.graph.RegisterNodeFunction;
import org.apache.beam.runners.dataflow.worker.graph.ReplacePgbkWithPrecombineFunction;
import org.apache.beam.runners.dataflow.worker.logging.DataflowWorkerLoggingMDC;
import org.apache.beam.runners.dataflow.worker.options.StreamingDataflowWorkerOptions;
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.util.BoundedQueueExecutor;
import org.apache.beam.runners.dataflow.worker.util.MemoryMonitor;
import org.apache.beam.runners.dataflow.worker.util.common.worker.OutputObjectAndByteCounter;
import org.apache.beam.runners.dataflow.worker.util.common.worker.ReadOperation;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill;
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.WindmillServerStub.CommitWorkStream;
import org.apache.beam.runners.dataflow.worker.windmill.WindmillServerStub.GetWorkStream;
import org.apache.beam.runners.dataflow.worker.windmill.WindmillServerStub.StreamPool;
import org.apache.beam.sdk.coders.Coder;
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;
import org.apache.beam.sdk.io.FileSystems;
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.UserCodeException;
import org.apache.beam.sdk.util.WindowedValue.WindowedValueCoder;
import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.TextFormat;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.MoreObjects;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Optional;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Splitter;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.EvictingQueue;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ListMultimap;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.MultimapBuilder;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.graph.MutableNetwork;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.net.HostAndPort;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.util.concurrent.Uninterruptibles;
import org.joda.time.Duration;
import org.joda.time.Instant;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
/** Implements a Streaming Dataflow worker. */
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);
/**
* Function which converts map tasks to their network representation for execution.
*
* <ul>
* <li>Translate the map task to a network representation.
* <li>Remove flatten instructions by rewiring edges.
* </ul>
*/
private static final Function<MapTask, MutableNetwork<Node, Edge>> mapTaskToBaseNetwork =
new MapTaskToNetworkFunction(idGenerator);
private static Random clientIdGenerator = new Random();
// Maximum number of threads for processing. Currently each thread processes one key at a time.
static final int MAX_PROCESSING_THREADS = 300;
static final long THREAD_EXPIRATION_TIME_SEC = 60;
// Maximum work units retrieved from Windmill and queued before processing. Limiting this delays
// retrieving extra work from Windmill without working on it, leading to better
// prioritization / utilization.
static final int MAX_WORK_UNITS_QUEUED = 100;
static final long TARGET_COMMIT_BUNDLE_BYTES = 32 << 20;
static final int MAX_COMMIT_QUEUE_BYTES = 500 << 20; // 500MB
static final int NUM_COMMIT_STREAMS = 1;
static final int GET_WORK_STREAM_TIMEOUT_MINUTES = 3;
static final Duration COMMIT_STREAM_TIMEOUT = Duration.standardMinutes(1);
private static final int DEFAULT_STATUS_PORT = 8081;
// Maximum size of the result of a GetWork request.
private static final long MAX_GET_WORK_FETCH_BYTES = 64L << 20; // 64m
// Reserved ID for counter updates.
// Matches kWindmillCounterUpdate in workflow_worker_service_multi_hubs.cc.
private static final String WINDMILL_COUNTER_UPDATE_WORK_ID = "3";
/** Maximum number of items to return in a GetWork request. */
private static final long MAX_GET_WORK_ITEMS = MAX_WORK_UNITS_QUEUED + MAX_PROCESSING_THREADS;
/** Maximum number of failure stacktraces to report in each update sent to backend. */
private static final int MAX_FAILURES_TO_REPORT_IN_UPDATE = 1000;
private final AtomicLong counterAggregationErrorCount = new AtomicLong();
/** Returns whether an exception was caused by a {@link OutOfMemoryError}. */
private static boolean isOutOfMemoryError(Throwable t) {
while (t != null) {
if (t instanceof OutOfMemoryError) {
return true;
}
t = t.getCause();
}
return false;
}
private static class KeyCommitTooLargeException extends Exception {
public static KeyCommitTooLargeException causedBy(
String computationId, long byteLimit, WorkItemCommitRequest request) {
StringBuilder message = new StringBuilder();
message.append("Commit request for stage ");
message.append(computationId);
message.append(" and key ");
message.append(request.getKey().toStringUtf8());
if (request.getSerializedSize() > 0) {
message.append(
" has size "
+ request.getSerializedSize()
+ " which is more than the limit of "
+ byteLimit);
} else {
message.append(" is larger than 2GB and cannot be processed");
}
message.append(
". This may be caused by grouping a very "
+ "large amount of data in a single window without using Combine,"
+ " or by producing a large amount of data from a single input element.");
return new KeyCommitTooLargeException(message.toString());
}
private KeyCommitTooLargeException(String message) {
super(message);
}
}
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();
DataflowWorkerHarnessHelper.initializeLogging(StreamingDataflowWorker.class);
DataflowWorkerHarnessOptions options =
DataflowWorkerHarnessHelper.initializeGlobalStateAndPipelineOptions(
StreamingDataflowWorker.class);
DataflowWorkerHarnessHelper.configureLogging(options);
checkArgument(
options.isStreaming(),
"%s instantiated with options indicating batch use",
StreamingDataflowWorker.class.getName());
checkArgument(
!DataflowRunner.hasExperiment(options, "beam_fn_api"),
"%s cannot be main() class with beam_fn_api enabled",
StreamingDataflowWorker.class.getSimpleName());
// Create a non fnapi registry
SdkHarnessRegistry sdkHarnessRegistry = SdkHarnessRegistries.emptySdkHarnessRegistry();
StreamingDataflowWorker worker =
StreamingDataflowWorker.fromDataflowWorkerHarnessOptions(options, sdkHarnessRegistry);
JvmInitializers.runBeforeProcessing(options);
worker.startStatusPages();
worker.start();
}
/** Bounded set of queues, with a maximum total weight. */
private static class WeightedBoundedQueue<V> {
private final LinkedBlockingQueue<V> queue = new LinkedBlockingQueue<>();
private final int maxWeight;
private final Semaphore limit;
private final Function<V, Integer> weigher;
public WeightedBoundedQueue(int maxWeight, Function<V, Integer> weigher) {
this.maxWeight = maxWeight;
this.limit = new Semaphore(maxWeight, true);
this.weigher = weigher;
}
/**
* Adds the value to the queue, blocking if this would cause the overall weight to exceed the
* limit.
*/
public void put(V value) {
limit.acquireUninterruptibly(weigher.apply(value));
queue.add(value);
}
/** Returns and removes the next value, or null if there is no such value. */
@Nullable
public V poll() {
V result = queue.poll();
if (result != null) {
limit.release(weigher.apply(result));
}
return result;
}
/**
* Retrieves and removes the head of this queue, waiting up to the specified wait time if
* necessary for an element to become available.
*
* @param timeout how long to wait before giving up, in units of {@code unit}
* @param unit a {@code TimeUnit} determining how to interpret the {@code timeout} parameter
* @return the head of this queue, or {@code null} if the specified waiting time elapses before
* an element is available
* @throws InterruptedException if interrupted while waiting
*/
@Nullable
public V poll(long timeout, TimeUnit unit) throws InterruptedException {
V result = queue.poll(timeout, unit);
if (result != null) {
limit.release(weigher.apply(result));
}
return result;
}
/** Returns and removes the next value, or blocks until one is available. */
@Nullable
public V take() throws InterruptedException {
V result = queue.take();
limit.release(weigher.apply(result));
return result;
}
/** Returns the current weight of the queue. */
public int weight() {
return maxWeight - limit.availablePermits();
}
public int size() {
return queue.size();
}
}
// Value class for a queued commit.
static class Commit {
private Windmill.WorkItemCommitRequest request;
private ComputationState computationState;
private Work work;
public Commit(
Windmill.WorkItemCommitRequest request, ComputationState computationState, Work work) {
this.request = request;
assert request.getSerializedSize() > 0;
this.computationState = computationState;
this.work = work;
}
public Windmill.WorkItemCommitRequest getRequest() {
return request;
}
public ComputationState getComputationState() {
return computationState;
}
public Work getWork() {
return work;
}
public int getSize() {
return request.getSerializedSize();
}
}
// Maps from computation ids to per-computation state.
private final ConcurrentMap<String, ComputationState> computationMap = new ConcurrentHashMap<>();
private final WeightedBoundedQueue<Commit> commitQueue =
new WeightedBoundedQueue<>(
MAX_COMMIT_QUEUE_BYTES, commit -> Math.min(MAX_COMMIT_QUEUE_BYTES, commit.getSize()));
// Map of tokens to commit callbacks.
private final ConcurrentMap<Long, Runnable> commitCallbacks = new ConcurrentHashMap<>();
// 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 WindmillStateCache stateCache = new WindmillStateCache();
private final ThreadFactory threadFactory;
private DataflowMapTaskExecutorFactory mapTaskExecutorFactory;
private final BoundedQueueExecutor workUnitExecutor;
private final WindmillServerStub windmillServer;
private final Thread dispatchThread;
private final Thread commitThread;
private final AtomicLong activeCommitBytes = new AtomicLong();
private final AtomicBoolean running = new AtomicBoolean();
private final StateFetcher stateFetcher;
private final StreamingDataflowWorkerOptions options;
private final boolean windmillServiceEnabled;
private final long clientId;
private final MetricTrackingWindmillServerStub metricTrackingWindmillServer;
private final CounterSet pendingDeltaCounters = new CounterSet();
private final CounterSet pendingCumulativeCounters = new CounterSet();
private final java.util.concurrent.ConcurrentLinkedQueue<CounterUpdate> pendingMonitoringInfos =
new ConcurrentLinkedQueue<>();
// Map from stage name to StageInfo containing metrics container registry and per stage counters.
private final ConcurrentMap<String, StageInfo> stageInfoMap = new ConcurrentHashMap();
// Built-in delta counters.
private final Counter<Long, Long> windmillShuffleBytesRead;
private final Counter<Long, Long> windmillStateBytesRead;
private final Counter<Long, Long> windmillStateBytesWritten;
private final Counter<Long, Long> windmillQuotaThrottling;
// Built-in cumulative counters.
private final Counter<Long, Long> javaHarnessUsedMemory;
private final Counter<Long, Long> javaHarnessMaxMemory;
private final Counter<Integer, Integer> windmillMaxObservedWorkItemCommitBytes;
private final Counter<Integer, Integer> memoryThrashing;
private Timer refreshWorkTimer;
private Timer statusPageTimer;
private final boolean publishCounters;
private Timer globalWorkerUpdatesTimer;
private int retryLocallyDelayMs = 10000;
// Periodically fires a global config request to dataflow service. Only used when windmill service
// is enabled.
private Timer globalConfigRefreshTimer;
private final MemoryMonitor memoryMonitor;
private final Thread memoryMonitorThread;
private final WorkerStatusPages statusPages;
// Periodic sender of debug information to the debug capture service.
private DebugCapture.Manager debugCaptureManager = null;
// Limit on bytes sinked (committed) in a work item.
private final long maxSinkBytes; // = MAX_SINK_BYTES unless disabled in options.
// Possibly overridden by streaming engine config.
private int maxWorkItemCommitBytes = Integer.MAX_VALUE;
private final EvictingQueue<String> pendingFailuresToReport =
EvictingQueue.<String>create(MAX_FAILURES_TO_REPORT_IN_UPDATE);
private final ReaderCache readerCache = new ReaderCache();
private final WorkUnitClient workUnitClient;
private final CompletableFuture<Void> isDoneFuture;
private final SdkHarnessRegistry sdkHarnessRegistry;
private final Function<MapTask, MutableNetwork<Node, Edge>> mapTaskToNetwork;
/**
* Sinks are marked 'full' in {@link StreamingModeExecutionContext} once the amount of data sinked
* (across all the sinks, if there are more than one) reaches this limit. This serves as hint for
* readers to stop producing more. This can be disabled with 'disable_limiting_bundle_sink_bytes'
* experiment.
*/
static final int MAX_SINK_BYTES = 10_000_000;
private final ReaderRegistry readerRegistry = ReaderRegistry.defaultRegistry();
private final SinkRegistry sinkRegistry = SinkRegistry.defaultRegistry();
private HotKeyLogger hotKeyLogger;
/** Contains a few of the stage specific fields. E.g. metrics container registry, counters etc. */
private static class StageInfo {
final String stageName;
final String systemName;
final MetricsContainerRegistry<StreamingStepMetricsContainer> metricsContainerRegistry;
final StreamingModeExecutionStateRegistry executionStateRegistry;
final CounterSet deltaCounters;
final Counter<Long, Long> throttledMsecs;
final Counter<Long, Long> totalProcessingMsecs;
final Counter<Long, Long> timerProcessingMsecs;
StageInfo(String stageName, String systemName, StreamingDataflowWorker worker) {
this.stageName = stageName;
this.systemName = systemName;
metricsContainerRegistry = StreamingStepMetricsContainer.createRegistry();
executionStateRegistry = new StreamingModeExecutionStateRegistry(worker);
NameContext nameContext = NameContext.create(stageName, null, systemName, null);
deltaCounters = new CounterSet();
throttledMsecs =
deltaCounters.longSum(
StreamingPerStageSystemCounterNames.THROTTLED_MSECS.counterName(nameContext));
totalProcessingMsecs =
deltaCounters.longSum(
StreamingPerStageSystemCounterNames.TOTAL_PROCESSING_MSECS.counterName(nameContext));
timerProcessingMsecs =
deltaCounters.longSum(
StreamingPerStageSystemCounterNames.TIMER_PROCESSING_MSECS.counterName(nameContext));
}
List<CounterUpdate> extractCounterUpdates() {
List<CounterUpdate> counterUpdates = new ArrayList<>();
Iterables.addAll(
counterUpdates,
StreamingStepMetricsContainer.extractMetricUpdates(metricsContainerRegistry));
Iterables.addAll(counterUpdates, executionStateRegistry.extractUpdates(false));
for (CounterUpdate counterUpdate : counterUpdates) {
translateKnownStepCounters(counterUpdate);
}
counterUpdates.addAll(
deltaCounters.extractModifiedDeltaUpdates(DataflowCounterUpdateExtractor.INSTANCE));
return counterUpdates;
}
// Checks if the step counter affects any per-stage counters. Currently 'throttled_millis'
// is the only counter updated.
private void translateKnownStepCounters(CounterUpdate stepCounterUpdate) {
CounterStructuredName structuredName =
stepCounterUpdate.getStructuredNameAndMetadata().getName();
if (THROTTLING_MSECS_METRIC_NAME.getNamespace().equals(structuredName.getOriginNamespace())
&& THROTTLING_MSECS_METRIC_NAME.getName().equals(structuredName.getName())) {
long msecs = DataflowCounterUpdateExtractor.splitIntToLong(stepCounterUpdate.getInteger());
if (msecs > 0) {
throttledMsecs.addValue(msecs);
}
}
}
}
public static StreamingDataflowWorker forStreamingFnWorkerHarness(
List<MapTask> mapTasks,
WorkUnitClient workUnitClient,
DataflowWorkerHarnessOptions options,
@Nullable RunnerApi.Pipeline pipeline,
SdkHarnessRegistry sdkHarnessRegistry)
throws IOException {
return new StreamingDataflowWorker(
mapTasks,
BeamFnMapTaskExecutorFactory.defaultFactory(),
workUnitClient,
options.as(StreamingDataflowWorkerOptions.class),
pipeline,
sdkHarnessRegistry,
true,
new HotKeyLogger());
}
public static StreamingDataflowWorker fromDataflowWorkerHarnessOptions(
DataflowWorkerHarnessOptions options, SdkHarnessRegistry sdkHarnessRegistry)
throws IOException {
return new StreamingDataflowWorker(
Collections.emptyList(),
IntrinsicMapTaskExecutorFactory.defaultFactory(),
new DataflowWorkUnitClient(options, LOG),
options.as(StreamingDataflowWorkerOptions.class),
null,
sdkHarnessRegistry,
true,
new HotKeyLogger());
}
@VisibleForTesting
StreamingDataflowWorker(
List<MapTask> mapTasks,
DataflowMapTaskExecutorFactory mapTaskExecutorFactory,
WorkUnitClient workUnitClient,
StreamingDataflowWorkerOptions options,
@Nullable RunnerApi.Pipeline pipeline,
SdkHarnessRegistry sdkHarnessRegistry,
boolean publishCounters,
HotKeyLogger hotKeyLogger)
throws IOException {
this.mapTaskExecutorFactory = mapTaskExecutorFactory;
this.workUnitClient = workUnitClient;
this.options = options;
this.sdkHarnessRegistry = sdkHarnessRegistry;
this.hotKeyLogger = hotKeyLogger;
this.windmillServiceEnabled = options.isEnableStreamingEngine();
this.memoryMonitor = MemoryMonitor.fromOptions(options);
this.statusPages =
WorkerStatusPages.create(
DEFAULT_STATUS_PORT, memoryMonitor, sdkHarnessRegistry::sdkHarnessesAreHealthy);
if (windmillServiceEnabled) {
this.debugCaptureManager =
new DebugCapture.Manager(options, statusPages.getDebugCapturePages());
}
this.windmillShuffleBytesRead =
pendingDeltaCounters.longSum(
StreamingSystemCounterNames.WINDMILL_SHUFFLE_BYTES_READ.counterName());
this.windmillStateBytesRead =
pendingDeltaCounters.longSum(
StreamingSystemCounterNames.WINDMILL_STATE_BYTES_READ.counterName());
this.windmillStateBytesWritten =
pendingDeltaCounters.longSum(
StreamingSystemCounterNames.WINDMILL_STATE_BYTES_WRITTEN.counterName());
this.windmillQuotaThrottling =
pendingDeltaCounters.longSum(
StreamingSystemCounterNames.WINDMILL_QUOTA_THROTTLING.counterName());
this.javaHarnessUsedMemory =
pendingCumulativeCounters.longSum(
StreamingSystemCounterNames.JAVA_HARNESS_USED_MEMORY.counterName());
this.javaHarnessMaxMemory =
pendingCumulativeCounters.longSum(
StreamingSystemCounterNames.JAVA_HARNESS_MAX_MEMORY.counterName());
this.windmillMaxObservedWorkItemCommitBytes =
pendingCumulativeCounters.intMax(
StreamingSystemCounterNames.WINDMILL_MAX_WORK_ITEM_COMMIT_BYTES.counterName());
this.memoryThrashing =
pendingCumulativeCounters.intSum(
StreamingSystemCounterNames.MEMORY_THRASHING.counterName());
this.isDoneFuture = new CompletableFuture<>();
this.threadFactory =
r -> {
Thread t = new Thread(r);
t.setDaemon(true);
return t;
};
this.workUnitExecutor =
new BoundedQueueExecutor(
chooseMaximumNumberOfThreads(),
THREAD_EXPIRATION_TIME_SEC,
TimeUnit.SECONDS,
MAX_WORK_UNITS_QUEUED,
threadFactory);
maxSinkBytes =
hasExperiment(options, "disable_limiting_bundle_sink_bytes")
? Long.MAX_VALUE
: MAX_SINK_BYTES;
memoryMonitorThread = new Thread(memoryMonitor);
memoryMonitorThread.setPriority(Thread.MIN_PRIORITY);
memoryMonitorThread.setName("MemoryMonitor");
dispatchThread =
threadFactory.newThread(
new Runnable() {
@Override
public void run() {
LOG.info("Dispatch starting");
if (windmillServiceEnabled) {
streamingDispatchLoop();
} else {
dispatchLoop();
}
LOG.info("Dispatch done");
}
});
dispatchThread.setPriority(Thread.MIN_PRIORITY);
dispatchThread.setName("DispatchThread");
commitThread =
threadFactory.newThread(
new Runnable() {
@Override
public void run() {
if (windmillServiceEnabled) {
streamingCommitLoop();
} else {
commitLoop();
}
}
});
commitThread.setPriority(Thread.MAX_PRIORITY);
commitThread.setName("CommitThread");
this.publishCounters = publishCounters;
this.windmillServer = options.getWindmillServerStub();
this.metricTrackingWindmillServer =
new MetricTrackingWindmillServerStub(windmillServer, memoryMonitor, windmillServiceEnabled);
this.metricTrackingWindmillServer.start();
this.stateFetcher = new StateFetcher(metricTrackingWindmillServer);
this.clientId = clientIdGenerator.nextLong();
for (MapTask mapTask : mapTasks) {
addComputation(mapTask.getSystemName(), mapTask, ImmutableMap.of());
}
// Register standard file systems.
FileSystems.setDefaultPipelineOptions(options);
if (hasExperiment(options, "beam_fn_api")) {
Function<MutableNetwork<Node, Edge>, Node> sdkFusedStage =
pipeline == null
? RegisterNodeFunction.withoutPipeline(
idGenerator, sdkHarnessRegistry.beamFnStateApiServiceDescriptor())
: RegisterNodeFunction.forPipeline(
pipeline, idGenerator, sdkHarnessRegistry.beamFnStateApiServiceDescriptor());
Function<MutableNetwork<Node, Edge>, MutableNetwork<Node, Edge>> lengthPrefixUnknownCoders =
LengthPrefixUnknownCoders::forSdkNetwork;
Function<MutableNetwork<Node, Edge>, MutableNetwork<Node, Edge>>
addStreamingSideInputHandlers =
InsertFetchAndFilterStreamingSideInputNodes.with(pipeline)::forNetwork;
Function<MutableNetwork<Node, Edge>, MutableNetwork<Node, Edge>> transformToRunnerNetwork =
new CreateRegisterFnOperationFunction(
idGenerator,
this::createPortNode,
lengthPrefixUnknownCoders.andThen(sdkFusedStage),
false);
mapTaskToNetwork =
mapTaskToBaseNetwork
.andThen(new ReplacePgbkWithPrecombineFunction())
.andThen(new DeduceNodeLocationsFunction())
.andThen(new DeduceFlattenLocationsFunction())
.andThen(new CloneAmbiguousFlattensFunction())
.andThen(addStreamingSideInputHandlers)
.andThen(transformToRunnerNetwork)
.andThen(LengthPrefixUnknownCoders::andReplaceForRunnerNetwork);
} else {
mapTaskToNetwork = mapTaskToBaseNetwork;
}
LOG.debug("windmillServiceEnabled: {}", windmillServiceEnabled);
LOG.debug("WindmillServiceEndpoint: {}", options.getWindmillServiceEndpoint());
LOG.debug("WindmillServicePort: {}", options.getWindmillServicePort());
LOG.debug("LocalWindmillHostport: {}", options.getLocalWindmillHostport());
LOG.debug("maxWorkItemCommitBytes: {}", maxWorkItemCommitBytes);
}
private Node createPortNode() {
return RemoteGrpcPortNode.create(
RemoteGrpcPort.newBuilder()
.setApiServiceDescriptor(sdkHarnessRegistry.beamFnDataApiServiceDescriptor())
.build(),
idGenerator.getId());
}
private int chooseMaximumNumberOfThreads() {
if (options.getNumberOfWorkerHarnessThreads() != 0) {
return options.getNumberOfWorkerHarnessThreads();
}
return MAX_PROCESSING_THREADS;
}
void addStateNameMappings(Map<String, String> nameMap) {
stateNameMap.putAll(nameMap);
}
@VisibleForTesting
public void setRetryLocallyDelayMs(int retryLocallyDelayMs) {
this.retryLocallyDelayMs = retryLocallyDelayMs;
}
@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.getQueue().isEmpty();
}
public void start() {
running.set(true);
if (windmillServiceEnabled) {
// Schedule the background getConfig thread. Blocks until windmillServer stub is ready.
schedulePeriodicGlobalConfigRequests();
}
memoryMonitorThread.start();
dispatchThread.start();
commitThread.start();
ExecutionStateSampler.instance().start();
// Periodically report workers counters and other updates.
globalWorkerUpdatesTimer = new Timer("GlobalWorkerUpdatesTimer");
globalWorkerUpdatesTimer.schedule(
new TimerTask() {
@Override
public void run() {
reportPeriodicWorkerUpdates();
}
},
0,
options.getWindmillHarnessUpdateReportingPeriod().getMillis());
refreshWorkTimer = new Timer("RefreshWork");
if (options.getActiveWorkRefreshPeriodMillis() > 0) {
refreshWorkTimer.schedule(
new TimerTask() {
@Override
public void run() {
refreshActiveWork();
}
},
options.getActiveWorkRefreshPeriodMillis(),
options.getActiveWorkRefreshPeriodMillis());
}
if (windmillServiceEnabled && options.getStuckCommitDurationMillis() > 0) {
int periodMillis = Math.max(options.getStuckCommitDurationMillis() / 10, 100);
refreshWorkTimer.schedule(
new TimerTask() {
@Override
public void run() {
invalidateStuckCommits();
}
},
periodMillis,
periodMillis);
}
if (options.getPeriodicStatusPageOutputDirectory() != null) {
statusPageTimer = new Timer("DumpStatusPages");
statusPageTimer.schedule(
new TimerTask() {
@Override
public void run() {
Collection<Capturable> pages = statusPages.getDebugCapturePages();
if (pages.isEmpty()) {
LOG.warn("No captured status pages.");
}
Long timestamp = Instant.now().getMillis();
for (Capturable page : pages) {
PrintWriter writer = null;
try {
File outputFile =
new File(
options.getPeriodicStatusPageOutputDirectory(),
("StreamingDataflowWorker"
+ options.getWorkerId()
+ "_"
+ page.pageName()
+ timestamp.toString())
.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 * 1000,
60 * 1000);
}
reportHarnessStartup();
}
public void startStatusPages() {
if (debugCaptureManager != null) {
debugCaptureManager.start();
}
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();
}
public void addWorkerStatusPage(BaseStatusServlet page) {
statusPages.addServlet(page);
if (page instanceof Capturable) {
statusPages.addCapturePage((Capturable) page);
}
}
public void stop() {
try {
if (globalConfigRefreshTimer != null) {
globalConfigRefreshTimer.cancel();
}
globalWorkerUpdatesTimer.cancel();
if (refreshWorkTimer != null) {
refreshWorkTimer.cancel();
}
if (statusPageTimer != null) {
statusPageTimer.cancel();
}
statusPages.stop();
if (debugCaptureManager != null) {
debugCaptureManager.stop();
}
running.set(false);
dispatchThread.interrupt();
dispatchThread.join();
// We need to interrupt the commitThread in case it is blocking on pulling
// from the commitQueue.
commitThread.interrupt();
commitThread.join();
memoryMonitor.stop();
memoryMonitorThread.join();
workUnitExecutor.shutdown();
if (!workUnitExecutor.awaitTermination(5, TimeUnit.MINUTES)) {
throw new RuntimeException("Work executor did not terminate within 5 minutes");
}
for (ComputationState state : computationMap.values()) {
state.close();
}
// one last send
reportPeriodicWorkerUpdates();
} 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);
}
public void waitTillExecutionFinishes() throws InterruptedException {
try {
isDoneFuture.get();
} catch (ExecutionException e) {
throw new IllegalStateException(e);
}
}
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)));
}
}
private static void sleep(int millis) {
Uninterruptibles.sleepUninterruptibly(millis, TimeUnit.MILLISECONDS);
}
/**
* 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() {
while (running.get()) {
memoryMonitor.waitForResources("GetWork");
int backoff = 1;
Windmill.GetWorkResponse workResponse = null;
do {
try {
workResponse = getWork();
if (workResponse.getWorkCount() > 0) {
break;
}
} catch (WindmillServerStub.RpcException e) {
LOG.warn("GetWork failed, retrying:", e);
}
sleep(backoff);
backoff = Math.min(1000, backoff * 2);
} 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());
continue;
}
final Instant inputDataWatermark =
WindmillTimeUtils.windmillToHarnessWatermark(computationWork.getInputDataWatermark());
Preconditions.checkNotNull(inputDataWatermark);
@Nullable
final Instant synchronizedProcessingTime =
WindmillTimeUtils.windmillToHarnessWatermark(
computationWork.getDependentRealtimeInputWatermark());
for (final Windmill.WorkItem workItem : computationWork.getWorkList()) {
scheduleWorkItem(
computationState, inputDataWatermark, synchronizedProcessingTime, workItem);
}
}
}
}
void streamingDispatchLoop() {
while (running.get()) {
GetWorkStream stream =
windmillServer.getWorkStream(
Windmill.GetWorkRequest.newBuilder()
.setClientId(clientId)
.setMaxItems(MAX_GET_WORK_ITEMS)
.setMaxBytes(MAX_GET_WORK_FETCH_BYTES)
.build(),
(String computation,
@Nullable Instant inputDataWatermark,
Instant synchronizedProcessingTime,
Windmill.WorkItem workItem) -> {
memoryMonitor.waitForResources("GetWork");
scheduleWorkItem(
getComputationState(computation),
inputDataWatermark,
synchronizedProcessingTime,
workItem);
});
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
// we half-close the stream after some time and create a new one.
if (!stream.awaitTermination(GET_WORK_STREAM_TIMEOUT_MINUTES, TimeUnit.MINUTES)) {
stream.close();
}
} catch (InterruptedException e) {
// Continue processing until !running.get()
}
}
}
private void scheduleWorkItem(
final ComputationState computationState,
final Instant inputDataWatermark,
final Instant synchronizedProcessingTime,
final Windmill.WorkItem workItem) {
Preconditions.checkNotNull(inputDataWatermark);
// May be null if output watermark not yet known.
@Nullable
final Instant outputDataWatermark =
WindmillTimeUtils.windmillToHarnessWatermark(workItem.getOutputDataWatermark());
Preconditions.checkState(
outputDataWatermark == null || !outputDataWatermark.isAfter(inputDataWatermark));
SdkWorkerHarness worker = sdkHarnessRegistry.getAvailableWorkerAndAssignWork();
if (workItem.hasHotKeyInfo()) {
Windmill.HotKeyInfo hotKeyInfo = workItem.getHotKeyInfo();
Duration hotKeyAge = Duration.millis(hotKeyInfo.getHotKeyAgeUsec() / 1000);
// The MapTask instruction is ordered by dependencies, such that the first element is
// always going to be the shuffle task.
String stepName = computationState.getMapTask().getInstructions().get(0).getName();
hotKeyLogger.logHotKeyDetection(stepName, hotKeyAge);
}
Work work =
new Work(workItem) {
@Override
public void run() {
try {
process(
worker,
computationState,
inputDataWatermark,
outputDataWatermark,
synchronizedProcessingTime,
this);
} finally {
// Reduce the work associated with the worker
sdkHarnessRegistry.completeWork(worker);
}
}
};
if (!computationState.activateWork(workItem.getKey(), work)) {
// Free worker if the work was not activated.
// This can happen if it's duplicate work or some other reason.
sdkHarnessRegistry.completeWork(worker);
}
}
abstract static class Work implements Runnable {
enum State {
QUEUED,
PROCESSING,
COMMIT_QUEUED,
COMMITTING
}
private final Windmill.WorkItem workItem;
private final Instant startTime;
private Instant stateStartTime;
private State state;
public Work(Windmill.WorkItem workItem) {
this.workItem = workItem;
this.startTime = this.stateStartTime = Instant.now();
this.state = State.QUEUED;
}
public Windmill.WorkItem getWorkItem() {
return workItem;
}
public Instant getStartTime() {
return startTime;
}
public State getState() {
return state;
}
public void setState(State state) {
this.state = state;
this.stateStartTime = Instant.now();
}
public Instant getStateStartTime() {
return stateStartTime;
}
}
/**
* Extracts the userland key coder, if any, from the coder used in the initial read step of a
* stage. This encodes many assumptions about how the streaming execution context works.
*/
@Nullable
private Coder<?> extractKeyCoder(Coder<?> readCoder) {
if (!(readCoder instanceof WindowedValueCoder)) {
throw new RuntimeException(
String.format(
"Expected coder for streaming read to be %s, but received %s",
WindowedValueCoder.class.getSimpleName(), readCoder));
}
// Note that TimerOrElementCoder is a backwards-compatibility class
// that is really a FakeKeyedWorkItemCoder
Coder<?> valueCoder = ((WindowedValueCoder<?>) readCoder).getValueCoder();
if (!(valueCoder instanceof WindmillKeyedWorkItem.FakeKeyedWorkItemCoder<?, ?>)) {
return null;
}
return ((WindmillKeyedWorkItem.FakeKeyedWorkItemCoder<?, ?>) valueCoder).getKeyCoder();
}
private void callFinalizeCallbacks(Windmill.WorkItem work) {
for (Long callbackId : work.getSourceState().getFinalizeIdsList()) {
final Runnable callback = commitCallbacks.remove(callbackId);
// NOTE: It is possible the same callback id may be removed twice if
// windmill restarts.
// TODO: It is also possible for an earlier finalized id to be lost.
// We should automatically discard all older callbacks for the same computation and key.
if (callback != null) {
workUnitExecutor.forceExecute(
() -> {
try {
callback.run();
} catch (Throwable t) {
LOG.error("Source checkpoint finalization failed:", t);
}
});
}
}
}
private Windmill.WorkItemCommitRequest.Builder initializeOutputBuilder(
final ByteString key, final Windmill.WorkItem workItem) {
return Windmill.WorkItemCommitRequest.newBuilder()
.setKey(key)
.setShardingKey(workItem.getShardingKey())
.setWorkToken(workItem.getWorkToken())
.setCacheToken(workItem.getCacheToken());
}
private void process(
final SdkWorkerHarness worker,
final ComputationState computationState,
final Instant inputDataWatermark,
@Nullable final Instant outputDataWatermark,
@Nullable final Instant synchronizedProcessingTime,
final Work work) {
final Windmill.WorkItem workItem = work.getWorkItem();
final String computationId = computationState.getComputationId();
final ByteString key = workItem.getKey();
work.setState(State.PROCESSING);
DataflowWorkerLoggingMDC.setWorkId(
TextFormat.escapeBytes(key) + "-" + Long.toString(workItem.getWorkToken()));
DataflowWorkerLoggingMDC.setStageName(computationId);
LOG.debug("Starting processing for {}:\n{}", computationId, work);
Windmill.WorkItemCommitRequest.Builder outputBuilder = initializeOutputBuilder(key, workItem);
// Before any processing starts, call any pending OnCommit callbacks. Nothing that requires
// cleanup should be done before this, since we might exit early here.
callFinalizeCallbacks(workItem);
if (workItem.getSourceState().getOnlyFinalize()) {
outputBuilder.setSourceStateUpdates(Windmill.SourceState.newBuilder().setOnlyFinalize(true));
work.setState(State.COMMIT_QUEUED);
commitQueue.put(new Commit(outputBuilder.build(), computationState, work));
return;
}
long processingStartTimeNanos = System.nanoTime();
final MapTask mapTask = computationState.getMapTask();
StageInfo stageInfo =
stageInfoMap.computeIfAbsent(
mapTask.getStageName(), s -> new StageInfo(s, mapTask.getSystemName(), this));
ExecutionState executionState = null;
try {
executionState = computationState.getExecutionStateQueue(worker).poll();
if (executionState == null) {
MutableNetwork<Node, Edge> mapTaskNetwork = mapTaskToNetwork.apply(mapTask);
if (LOG.isDebugEnabled()) {
LOG.debug("Network as Graphviz .dot: {}", Networks.toDot(mapTaskNetwork));
}
ParallelInstructionNode readNode =
(ParallelInstructionNode)
Iterables.find(
mapTaskNetwork.nodes(),
node ->
node instanceof ParallelInstructionNode
&& ((ParallelInstructionNode) node).getParallelInstruction().getRead()
!= null);
InstructionOutputNode readOutputNode =
(InstructionOutputNode) Iterables.getOnlyElement(mapTaskNetwork.successors(readNode));
DataflowExecutionContext.DataflowExecutionStateTracker executionStateTracker =
new DataflowExecutionContext.DataflowExecutionStateTracker(
ExecutionStateSampler.instance(),
stageInfo.executionStateRegistry.getState(
NameContext.forStage(mapTask.getStageName()),
"other",
null,
ScopedProfiler.INSTANCE.emptyScope()),
stageInfo.deltaCounters,
options,
computationId);
StreamingModeExecutionContext context =
new StreamingModeExecutionContext(
pendingDeltaCounters,
computationId,
readerCache,
!computationState.getTransformUserNameToStateFamily().isEmpty()
? computationState.getTransformUserNameToStateFamily()
: stateNameMap,
stateCache.forComputation(computationId),
stageInfo.metricsContainerRegistry,
executionStateTracker,
stageInfo.executionStateRegistry,
maxSinkBytes);
DataflowMapTaskExecutor mapTaskExecutor =
mapTaskExecutorFactory.create(
worker.getControlClientHandler(),
worker.getGrpcDataFnServer(),
sdkHarnessRegistry.beamFnDataApiServiceDescriptor(),
worker.getGrpcStateFnServer(),
mapTaskNetwork,
options,
mapTask.getStageName(),
readerRegistry,
sinkRegistry,
context,
pendingDeltaCounters,
idGenerator);
ReadOperation readOperation = mapTaskExecutor.getReadOperation();
// Disable progress updates since its results are unused for streaming
// and involves starting a thread.
readOperation.setProgressUpdatePeriodMs(ReadOperation.DONT_UPDATE_PERIODICALLY);
Preconditions.checkState(
mapTaskExecutor.supportsRestart(),
"Streaming runner requires all operations support restart.");
Coder<?> readCoder;
readCoder =
CloudObjects.coderFromCloudObject(
CloudObject.fromSpec(readOutputNode.getInstructionOutput().getCodec()));
Coder<?> keyCoder = extractKeyCoder(readCoder);
// If using a custom source, count bytes read for autoscaling.
if (CustomSources.class
.getName()
.equals(
readNode.getParallelInstruction().getRead().getSource().getSpec().get("@type"))) {
NameContext nameContext =
NameContext.create(
mapTask.getStageName(),
readNode.getParallelInstruction().getOriginalName(),
readNode.getParallelInstruction().getSystemName(),
readNode.getParallelInstruction().getName());
readOperation.receivers[0].addOutputCounter(
new OutputObjectAndByteCounter(
new IntrinsicMapTaskExecutorFactory.ElementByteSizeObservableCoder<>(
readCoder),
mapTaskExecutor.getOutputCounters(),
nameContext)
.setSamplingPeriod(100)
.countBytes("dataflow_input_size-" + mapTask.getSystemName()));
}
executionState =
new ExecutionState(mapTaskExecutor, context, keyCoder, executionStateTracker);
}
WindmillStateReader stateReader =
new WindmillStateReader(
metricTrackingWindmillServer,
computationId,
key,
workItem.getShardingKey(),
workItem.getWorkToken());
StateFetcher localStateFetcher = stateFetcher.byteTrackingView();
// If the read output KVs, then we can decode Windmill's byte key into a userland
// key object and provide it to the execution context for use with per-key state.
// Otherwise, we pass null.
//
// The coder type that will be present is:
// WindowedValueCoder(TimerOrElementCoder(KvCoder))
@Nullable Coder<?> keyCoder = executionState.getKeyCoder();
@Nullable
Object executionKey =
keyCoder == null ? null : keyCoder.decode(key.newInput(), Coder.Context.OUTER);
executionState
.getContext()
.start(
executionKey,
workItem,
inputDataWatermark,
outputDataWatermark,
synchronizedProcessingTime,
stateReader,
localStateFetcher,
outputBuilder);
// Blocks while executing work.
executionState.getWorkExecutor().execute();
Iterables.addAll(
this.pendingMonitoringInfos, executionState.getWorkExecutor().extractMetricUpdates());
commitCallbacks.putAll(executionState.getContext().flushState());
// Release the execution state for another thread to use.
computationState.getExecutionStateQueue(worker).offer(executionState);
executionState = null;
// Add the output to the commit queue.
work.setState(State.COMMIT_QUEUED);
WorkItemCommitRequest commitRequest = outputBuilder.build();
int byteLimit = maxWorkItemCommitBytes;
int commitSize = commitRequest.getSerializedSize();
int estimatedCommitSize = commitSize < 0 ? Integer.MAX_VALUE : commitSize;
// Detect overflow of integer serialized size or if the byte limit was exceeded.
windmillMaxObservedWorkItemCommitBytes.addValue(estimatedCommitSize);
if (estimatedCommitSize > byteLimit) {
KeyCommitTooLargeException e =
KeyCommitTooLargeException.causedBy(computationId, byteLimit, commitRequest);
reportFailure(computationId, workItem, e);
LOG.error(e.toString());
// Drop the current request in favor of a new, minimal one requesting truncation.
// Messages, timers, counters, and other commit content will not be used by the service
// so we're purposefully dropping them here
commitRequest = buildWorkItemTruncationRequest(key, workItem, estimatedCommitSize);
}
commitQueue.put(new Commit(commitRequest, computationState, work));
// Compute shuffle and state byte statistics these will be flushed asynchronously.
long stateBytesWritten = outputBuilder.clearOutputMessages().build().getSerializedSize();
long shuffleBytesRead = 0;
for (Windmill.InputMessageBundle bundle : workItem.getMessageBundlesList()) {
for (Windmill.Message message : bundle.getMessagesList()) {
shuffleBytesRead += message.getSerializedSize();
}
}
long stateBytesRead = stateReader.getBytesRead() + localStateFetcher.getBytesRead();
windmillShuffleBytesRead.addValue(shuffleBytesRead);
windmillStateBytesRead.addValue(stateBytesRead);
windmillStateBytesWritten.addValue(stateBytesWritten);
LOG.debug("Processing done for work token: {}", workItem.getWorkToken());
} catch (Throwable t) {
if (executionState != null) {
try {
executionState.getContext().invalidateCache();
executionState.getWorkExecutor().close();
} catch (Exception e) {
LOG.warn("Failed to close map task executor: ", e);
} finally {
// Release references to potentially large objects early.
executionState = null;
}
}
t = t instanceof UserCodeException ? t.getCause() : t;
boolean retryLocally = false;
if (KeyTokenInvalidException.isKeyTokenInvalidException(t)) {
LOG.debug(
"Execution of work for {} for key {} failed due to token expiration. "
+ "Will not retry locally.",
computationId,
key.toStringUtf8());
} else {
LOG.error("Uncaught exception: ", t);
LastExceptionDataProvider.reportException(t);
LOG.debug("Failed work: {}", work);
if (!reportFailure(computationId, workItem, t)) {
LOG.error(
"Execution of work for {} for key {} failed, and Windmill "
+ "indicated not to retry locally.",
computationId,
key.toStringUtf8());
} else if (isOutOfMemoryError(t)) {
File heapDump = memoryMonitor.tryToDumpHeap();
LOG.error(
"Execution of work for {} for key {} failed with out-of-memory. "
+ "Will not retry locally. Heap dump {}.",
computationId,
key.toStringUtf8(),
heapDump == null ? "not written" : ("written to '" + heapDump + "'"));
} else {
LOG.error(
"Execution of work for {} for key {} failed. Will retry locally.",
computationId,
key.toStringUtf8());
retryLocally = true;
}
}
if (retryLocally) {
// Try again after some delay and at the end of the queue to avoid a tight loop.
sleep(retryLocallyDelayMs);
workUnitExecutor.forceExecute(work);
} else {
// Consider the item invalid. It will eventually be retried by Windmill if it still needs to
// be processed.
computationState.completeWork(key, workItem.getWorkToken());
}
} finally {
// Update total processing time counters. Updating in finally clause ensures that
// work items causing exceptions are also accounted in time spent.
long processingTimeMsecs =
TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - processingStartTimeNanos);
stageInfo.totalProcessingMsecs.addValue(processingTimeMsecs);
// Attribute all the processing to timers if the work item contains any timers.
// Tests show that work items rarely contain both timers and message bundles. It should
// be a fairly close approximation.
// Another option: Derive time split between messages and timers based on recent totals.
// either here or in DFE.
if (work.getWorkItem().hasTimers()) {
stageInfo.timerProcessingMsecs.addValue(processingTimeMsecs);
}
DataflowWorkerLoggingMDC.setWorkId(null);
DataflowWorkerLoggingMDC.setStageName(null);
}
}
private WorkItemCommitRequest buildWorkItemTruncationRequest(
final ByteString key, final Windmill.WorkItem workItem, final int estimatedCommitSize) {
Windmill.WorkItemCommitRequest.Builder outputBuilder = initializeOutputBuilder(key, workItem);
outputBuilder.setExceedsMaxWorkItemCommitBytes(true);
outputBuilder.setEstimatedWorkItemCommitBytes(estimatedCommitSize);
return outputBuilder.build();
}
private void commitLoop() {
Map<ComputationState, Windmill.ComputationCommitWorkRequest.Builder> computationRequestMap =
new HashMap<>();
while (running.get()) {
computationRequestMap.clear();
Windmill.CommitWorkRequest.Builder commitRequestBuilder =
Windmill.CommitWorkRequest.newBuilder();
long commitBytes = 0;
// Block until we have a commit, then batch with additional commits.
Commit commit = null;
try {
commit = commitQueue.take();
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
continue;
}
while (commit != null) {
ComputationState computationState = commit.getComputationState();
commit.getWork().setState(State.COMMITTING);
Windmill.ComputationCommitWorkRequest.Builder computationRequestBuilder =
computationRequestMap.get(computationState);
if (computationRequestBuilder == null) {
computationRequestBuilder = commitRequestBuilder.addRequestsBuilder();
computationRequestBuilder.setComputationId(computationState.getComputationId());
computationRequestMap.put(computationState, computationRequestBuilder);
}
computationRequestBuilder.addRequests(commit.getRequest());
// Send the request if we've exceeded the bytes or there is no more
// pending work. commitBytes is a long, so this cannot overflow.
commitBytes += commit.getSize();
if (commitBytes >= TARGET_COMMIT_BUNDLE_BYTES) {
break;
}
commit = commitQueue.poll();
}
Windmill.CommitWorkRequest commitRequest = commitRequestBuilder.build();
LOG.trace("Commit: {}", commitRequest);
activeCommitBytes.set(commitBytes);
windmillServer.commitWork(commitRequest);
activeCommitBytes.set(0);
for (Map.Entry<ComputationState, Windmill.ComputationCommitWorkRequest.Builder> entry :
computationRequestMap.entrySet()) {
ComputationState computationState = entry.getKey();
for (Windmill.WorkItemCommitRequest workRequest : entry.getValue().getRequestsList()) {
computationState.completeWork(workRequest.getKey(), workRequest.getWorkToken());
}
}
}
}
// Adds the commit to the commitStream if it fits, returning true iff it is consumed.
private boolean addCommitToStream(Commit commit, CommitWorkStream commitStream) {
Preconditions.checkNotNull(commit);
final ComputationState state = commit.getComputationState();
final Windmill.WorkItemCommitRequest request = commit.getRequest();
final int size = commit.getSize();
commit.getWork().setState(State.COMMITTING);
activeCommitBytes.addAndGet(size);
if (commitStream.commitWorkItem(
state.computationId,
request,
(Windmill.CommitStatus status) -> {
if (status != Windmill.CommitStatus.OK) {
stateCache.forComputation(state.computationId).invalidate(request.getKey());
}
activeCommitBytes.addAndGet(-size);
// This may throw an exception if the commit was not active, which is possible if it
// was deemed stuck.
state.completeWork(request.getKey(), request.getWorkToken());
})) {
return true;
} else {
// Back out the stats changes since the commit wasn't consumed.
commit.getWork().setState(State.COMMIT_QUEUED);
activeCommitBytes.addAndGet(-size);
return false;
}
}
// Helper to batch additional commits into the commit stream as long as they fit.
// Returns a commit that was removed from the queue but not consumed or null.
private Commit batchCommitsToStream(CommitWorkStream commitStream) {
int commits = 1;
while (running.get()) {
Commit commit;
try {
if (commits < 5) {
commit = commitQueue.poll(10 - 2 * commits, TimeUnit.MILLISECONDS);
} else {
commit = commitQueue.poll();
}
} catch (InterruptedException e) {
// Continue processing until !running.get()
continue;
}
if (commit == null || !addCommitToStream(commit, commitStream)) {
return commit;
}
commits++;
}
return null;
}
private void streamingCommitLoop() {
StreamPool<CommitWorkStream> streamPool =
new StreamPool<>(
NUM_COMMIT_STREAMS, COMMIT_STREAM_TIMEOUT, windmillServer::commitWorkStream);
Commit initialCommit = null;
while (running.get()) {
if (initialCommit == null) {
try {
initialCommit = commitQueue.take();
} catch (InterruptedException e) {
continue;
}
}
// We initialize the commit stream only after we have a commit to make sure it is fresh.
CommitWorkStream commitStream = streamPool.getStream();
if (!addCommitToStream(initialCommit, commitStream)) {
throw new AssertionError("Initial commit on flushed stream should always be accepted.");
}
// Batch additional commits to the stream and possibly make an un-batched commit the next
// initial commit.
initialCommit = batchCommitsToStream(commitStream);
commitStream.flush();
streamPool.releaseStream(commitStream);
}
}
private Windmill.GetWorkResponse getWork() {
return windmillServer.getWork(
Windmill.GetWorkRequest.newBuilder()
.setClientId(clientId)
.setMaxItems(MAX_GET_WORK_ITEMS)
.setMaxBytes(MAX_GET_WORK_FETCH_BYTES)
.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.get(computationId));
} 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;
if (computation != null) {
workItem = workUnitClient.getStreamingConfigWorkItem(computation);
} else {
workItem = workUnitClient.getGlobalStreamingConfigWorkItem();
}
if (workItem == null || !workItem.isPresent() || workItem.get() == null) {
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,
computationConfig.getTransformUserNameToStateFamily());
}
}
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.
*/
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.
globalConfigRefreshTimer = new Timer("GlobalConfigRefreshTimer");
globalConfigRefreshTimer.schedule(
new TimerTask() {
@Override
public void run() {
getGlobalConfig();
}
},
0,
options.getGlobalConfigRefreshPeriod().getMillis());
}
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(
pendingDeltaCounters.extractModifiedDeltaUpdates(DataflowCounterUpdateExtractor.INSTANCE),
pendingCumulativeCounters.extractUpdates(false, DataflowCounterUpdateExtractor.INSTANCE));
}
private Windmill.Exception buildExceptionReport(Throwable t) {
Windmill.Exception.Builder builder = Windmill.Exception.newBuilder();
builder.addStackFrames(t.toString());
for (StackTraceElement frame : t.getStackTrace()) {
builder.addStackFrames(frame.toString());
}
if (t.getCause() != null) {
builder.setCause(buildExceptionReport(t.getCause()));
}
return builder.build();
}
private String buildExceptionStackTrace(Throwable t, final int maxDepth) {
StringBuilder builder = new StringBuilder(1024);
Throwable cur = t;
for (int depth = 0; cur != null && depth < maxDepth; cur = cur.getCause()) {
if (depth > 0) {
builder.append("\nCaused by: ");
}
builder.append(cur);
depth++;
for (StackTraceElement frame : cur.getStackTrace()) {
if (depth < maxDepth) {
builder.append("\n ");
builder.append(frame);
depth++;
}
}
}
if (cur != null) {
builder.append("\nStack trace truncated. Please see Cloud Logging for the entire trace.");
}
return builder.toString();
}
// Returns true if reporting the exception is successful and the work should be retried.
private boolean reportFailure(String computation, Windmill.WorkItem work, Throwable t) {
// Enqueue the errors to be sent to DFE in periodic updates
addFailure(buildExceptionStackTrace(t, options.getMaxStackTraceDepthToReport()));
if (windmillServiceEnabled) {
return true;
} else {
Windmill.ReportStatsResponse response =
windmillServer.reportStats(
Windmill.ReportStatsRequest.newBuilder()
.setComputationId(computation)
.setKey(work.getKey())
.setShardingKey(work.getShardingKey())
.setWorkToken(work.getWorkToken())
.build());
return !response.getFailed();
}
}
/**
* Adds the given failure message to the queue of messages to be reported to DFE in periodic
* updates.
*/
public void addFailure(String failureMessage) {
synchronized (pendingFailuresToReport) {
pendingFailuresToReport.add(failureMessage);
}
}
private void reportHarnessStartup() {
DataflowWorkerLoggingMDC.setStageName("startup");
CounterSet restartCounter = new CounterSet();
restartCounter
.longSum(StreamingSystemCounterNames.JAVA_HARNESS_RESTARTS.counterName())
.addValue(1L);
try {
// Sending a one time update. Use empty counter set for cumulativeCounters (2nd arg).
sendWorkerUpdatesToDataflowService(restartCounter, new CounterSet());
} catch (IOException e) {
LOG.warn("Failed to send harness startup counter", e);
}
}
/** Updates VM metrics like memory and CPU utilization. */
private void updateVMMetrics() {
Runtime rt = Runtime.getRuntime();
long usedMemory = rt.totalMemory() - rt.freeMemory();
long maxMemory = rt.maxMemory();
javaHarnessUsedMemory.getAndReset();
javaHarnessUsedMemory.addValue(usedMemory);
javaHarnessMaxMemory.getAndReset();
javaHarnessMaxMemory.addValue(maxMemory);
}
@VisibleForTesting
public void reportPeriodicWorkerUpdates() {
updateVMMetrics();
try {
sendWorkerUpdatesToDataflowService(pendingDeltaCounters, pendingCumulativeCounters);
} catch (IOException e) {
LOG.warn("Failed to send periodic counter updates", e);
} catch (Exception e) {
LOG.error("Unexpected exception while trying to send counter updates", e);
}
}
/**
* Returns key for a counter update. It is a String in case of legacy counter and
* CounterStructuredName in the case of a structured counter.
*/
private Object getCounterUpdateKey(CounterUpdate counterUpdate) {
Object key = null;
if (counterUpdate.getNameAndKind() != null) {
key = counterUpdate.getNameAndKind().getName();
} else if (counterUpdate.getStructuredNameAndMetadata() != null) {
key = counterUpdate.getStructuredNameAndMetadata().getName();
}
checkArgument(key != null, "Could not find name for CounterUpdate: %s", counterUpdate);
return key;
}
/** Sends counter updates to Dataflow backend. */
private void sendWorkerUpdatesToDataflowService(
CounterSet deltaCounters, CounterSet cumulativeCounters) throws IOException {
// Throttle time is tracked by the windmillServer but is reported to DFE here.
windmillQuotaThrottling.addValue(windmillServer.getAndResetThrottleTime());
if (memoryMonitor.isThrashing()) {
memoryThrashing.addValue(1);
}
List<CounterUpdate> counterUpdates = new ArrayList<>(128);
if (publishCounters) {
stageInfoMap.values().forEach(s -> counterUpdates.addAll(s.extractCounterUpdates()));
counterUpdates.addAll(
cumulativeCounters.extractUpdates(false, DataflowCounterUpdateExtractor.INSTANCE));
counterUpdates.addAll(
deltaCounters.extractModifiedDeltaUpdates(DataflowCounterUpdateExtractor.INSTANCE));
if (hasExperiment(options, "beam_fn_api")) {
Map<Object, List<CounterUpdate>> fnApiCounters = new HashMap<>();
while (!this.pendingMonitoringInfos.isEmpty()) {
final CounterUpdate item = this.pendingMonitoringInfos.poll();
// This change will treat counter as delta.
// This is required because we receive cumulative results from FnAPI harness,
// while streaming job is expected to receive delta updates to counters on same
// WorkItem.
if (item.getCumulative()) {
item.setCumulative(false);
// Group counterUpdates by counterUpdateKey so they can be aggregated before sending to
// dataflow service.
fnApiCounters
.computeIfAbsent(getCounterUpdateKey(item), k -> new ArrayList<>())
.add(item);
} else {
// In current world all counters coming from FnAPI are cumulative.
// This is a safety check in case new counter type appears in FnAPI.
throw new UnsupportedOperationException(
"FnApi counters are expected to provide cumulative values."
+ " Please, update conversion to delta logic"
+ " if non-cumulative counter type is required.");
}
}
// Aggregates counterUpdates with same counterUpdateKey to single CounterUpdate if possible
// so we can avoid excessive I/Os for reporting to dataflow service.
for (List<CounterUpdate> counterUpdateList : fnApiCounters.values()) {
if (counterUpdateList.isEmpty()) {
continue;
}
List<CounterUpdate> aggregatedCounterUpdateList =
CounterUpdateAggregators.aggregate(counterUpdateList);
// Log a warning message if encountered enough non-aggregatable counter updates since this
// can lead to a severe performance penalty if dataflow service can not handle the
// updates.
if (aggregatedCounterUpdateList.size() > 10) {
CounterUpdate head = aggregatedCounterUpdateList.get(0);
this.counterAggregationErrorCount.getAndIncrement();
// log warning message only when error count is the power of 2 to avoid spamming.
if (this.counterAggregationErrorCount.get() > 10
&& Long.bitCount(this.counterAggregationErrorCount.get()) == 1) {
LOG.warn(
"Found non-aggregated counter updates of size {} with kind {}, this will likely "
+ "cause performance degradation and excessive GC if size is large.",
counterUpdateList.size(),
MoreObjects.firstNonNull(
head.getNameAndKind(), head.getStructuredNameAndMetadata()));
}
}
counterUpdates.addAll(aggregatedCounterUpdateList);
}
}
}
// Handle duplicate counters from different stages. Store all the counters in a multi-map and
// send the counters that appear multiple times in separate RPCs. Same logical counter could
// appear in multiple stages if a step runs in multiple stages (as with flatten-unzipped stages)
// especially if the counter definition does not set execution_step_name.
ListMultimap<Object, CounterUpdate> counterMultimap =
MultimapBuilder.hashKeys(counterUpdates.size()).linkedListValues().build();
boolean hasDuplicates = false;
for (CounterUpdate c : counterUpdates) {
Object key = getCounterUpdateKey(c);
if (counterMultimap.containsKey(key)) {
hasDuplicates = true;
}
counterMultimap.put(key, c);
}
// Clears counterUpdates and enqueues unique counters from counterMultimap. If a counter
// appears more than once, one of them is extracted leaving the remaining in the map.
Runnable extractUniqueCounters =
() -> {
counterUpdates.clear();
for (Iterator<Object> iter = counterMultimap.keySet().iterator(); iter.hasNext(); ) {
List<CounterUpdate> counters = counterMultimap.get(iter.next());
counterUpdates.add(counters.get(0));
if (counters.size() == 1) {
// There is single value. Remove the entry through the iterator.
iter.remove();
} else {
// Otherwise remove the first value.
counters.remove(0);
}
}
};
if (hasDuplicates) {
extractUniqueCounters.run();
} else { // Common case: no duplicates. We can just send counterUpdates, empty the multimap.
counterMultimap.clear();
}
List<Status> errors;
synchronized (pendingFailuresToReport) {
errors = new ArrayList<>(pendingFailuresToReport.size());
for (String stackTrace : pendingFailuresToReport) {
errors.add(
new Status()
.setCode(2) // rpc.Code.UNKNOWN
.setMessage(stackTrace));
}
pendingFailuresToReport.clear(); // Best effort only, no need to wait till successfully sent.
}
WorkItemStatus workItemStatus =
new WorkItemStatus()
.setWorkItemId(WINDMILL_COUNTER_UPDATE_WORK_ID)
.setErrors(errors)
.setCounterUpdates(counterUpdates);
workUnitClient.reportWorkItemStatus(workItemStatus);
// Send any counters appearing more than once in subsequent RPCs:
while (!counterMultimap.isEmpty()) {
extractUniqueCounters.run();
workUnitClient.reportWorkItemStatus(
new WorkItemStatus()
.setWorkItemId(WINDMILL_COUNTER_UPDATE_WORK_ID)
.setCounterUpdates(counterUpdates));
}
}
/**
* Sends a GetData request to Windmill for all sufficiently old active work.
*
* <p>This informs Windmill that processing is ongoing and the work should not be retried. The age
* threshold is determined by {@link
* StreamingDataflowWorkerOptions#getActiveWorkRefreshPeriodMillis}.
*/
private void refreshActiveWork() {
Map<String, List<Windmill.KeyedGetDataRequest>> active = new HashMap<>();
Instant refreshDeadline =
Instant.now().minus(Duration.millis(options.getActiveWorkRefreshPeriodMillis()));
for (Map.Entry<String, ComputationState> entry : computationMap.entrySet()) {
active.put(entry.getKey(), entry.getValue().getKeysToRefresh(refreshDeadline));
}
metricTrackingWindmillServer.refreshActiveWork(active);
}
private void invalidateStuckCommits() {
Instant stuckCommitDeadline =
Instant.now().minus(Duration.millis(options.getStuckCommitDurationMillis()));
for (Map.Entry<String, ComputationState> entry : computationMap.entrySet()) {
entry.getValue().invalidateStuckCommits(stuckCommitDeadline);
}
}
/**
* Class representing the state of a computation.
*
* <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.
*/
static class ComputationState implements AutoCloseable {
private final String computationId;
private final MapTask mapTask;
private final ImmutableMap<String, String> transformUserNameToStateFamily;
// Map from key to work for the key. The first item in the queue is
// actively processing. Synchronized by itself.
private final Map<ByteString, Deque<Work>> activeWork = new HashMap<>();
private final BoundedQueueExecutor executor;
private final ConcurrentMap<SdkWorkerHarness, ConcurrentLinkedQueue<ExecutionState>>
executionStateQueues = new ConcurrentHashMap<>();
private final WindmillStateCache.ForComputation computationStateCache;
public ComputationState(
String computationId,
MapTask mapTask,
BoundedQueueExecutor executor,
Map<String, String> transformUserNameToStateFamily,
WindmillStateCache.ForComputation computationStateCache) {
this.computationId = computationId;
this.mapTask = mapTask;
this.executor = executor;
this.transformUserNameToStateFamily =
transformUserNameToStateFamily != null
? ImmutableMap.copyOf(transformUserNameToStateFamily)
: ImmutableMap.of();
this.computationStateCache = computationStateCache;
Preconditions.checkNotNull(mapTask.getStageName());
Preconditions.checkNotNull(mapTask.getSystemName());
}
public String getComputationId() {
return computationId;
}
public MapTask getMapTask() {
return mapTask;
}
public ImmutableMap<String, String> getTransformUserNameToStateFamily() {
return transformUserNameToStateFamily;
}
public ConcurrentLinkedQueue<ExecutionState> getExecutionStateQueue(
SdkWorkerHarness sdkWorkerHarness) {
return executionStateQueues.computeIfAbsent(
sdkWorkerHarness, key -> new ConcurrentLinkedQueue<>());
}
/** Mark the given key and work as active. */
public boolean activateWork(ByteString key, Work work) {
synchronized (activeWork) {
Deque<Work> queue = activeWork.get(key);
if (queue != null) {
Preconditions.checkState(!queue.isEmpty());
// Ensure we don't already have this work token queueud.
for (Work queuedWork : queue) {
if (queuedWork.getWorkItem().getWorkToken() == work.getWorkItem().getWorkToken()) {
return false;
}
}
// Queue the work for later processing.
queue.addLast(work);
return true;
} else {
queue = new ArrayDeque<>();
queue.addLast(work);
activeWork.put(key, queue);
// Fall through to execute without the lock held.
}
}
executor.execute(work);
return true;
}
/** Marks the work for a the given key as complete. Schedules queued work for the key if any. */
public void completeWork(ByteString key, long workToken) {
Work nextWork;
synchronized (activeWork) {
Queue<Work> queue = activeWork.get(key);
Preconditions.checkNotNull(queue);
Work completedWork = queue.peek();
// avoid Preconditions.checkNotNull and checkState here to prevent eagerly evaluating the
// format string parameters for the error message.
if (completedWork == null) {
throw new NullPointerException(
String.format(
"No active state for key %s, expected token %s",
TextFormat.escapeBytes(key), workToken));
}
if (completedWork.getWorkItem().getWorkToken() != workToken) {
throw new IllegalStateException(
String.format(
"Token mismatch for key %s: %s and %s",
TextFormat.escapeBytes(key),
completedWork.getWorkItem().getWorkToken(),
workToken));
}
queue.remove(); // We consumed the matching work item.
nextWork = queue.peek();
if (nextWork == null) {
Preconditions.checkState(queue == activeWork.remove(key));
}
}
if (nextWork != null) {
executor.forceExecute(nextWork);
}
}
public void invalidateStuckCommits(Instant stuckCommitDeadline) {
synchronized (activeWork) {
// Determine the stuck commit keys but complete them outside of iterating over
// activeWork as completeWork may delete the entry from activeWork.
Map<ByteString, Long> stuckCommits = new HashMap<>();
for (Map.Entry<ByteString, Deque<Work>> entry : activeWork.entrySet()) {
ByteString key = entry.getKey();
Work work = entry.getValue().peek();
if (work.getState() == State.COMMITTING
&& work.getStateStartTime().isBefore(stuckCommitDeadline)) {
LOG.error(
"Detected key with sharding key {} stuck in COMMITTING state, completing it with error. Key ",
work.workItem.getShardingKey());
stuckCommits.put(key, work.getWorkItem().getWorkToken());
}
}
for (Map.Entry<ByteString, Long> stuckCommit : stuckCommits.entrySet()) {
computationStateCache.invalidate(stuckCommit.getKey());
completeWork(stuckCommit.getKey(), stuckCommit.getValue());
}
}
}
/** Adds any work started before the refreshDeadline to the GetDataRequest builder. */
public List<Windmill.KeyedGetDataRequest> getKeysToRefresh(Instant refreshDeadline) {
List<Windmill.KeyedGetDataRequest> result = new ArrayList<>();
synchronized (activeWork) {
for (Map.Entry<ByteString, Deque<Work>> entry : activeWork.entrySet()) {
ByteString key = entry.getKey();
for (Work work : entry.getValue()) {
if (work.getStartTime().isBefore(refreshDeadline)) {
result.add(
Windmill.KeyedGetDataRequest.newBuilder()
.setKey(key)
.setShardingKey(work.getWorkItem().getShardingKey())
.setWorkToken(work.getWorkItem().getWorkToken())
.build());
}
}
}
}
return result;
}
private String elapsedString(Instant start, Instant end) {
Duration activeFor = new Duration(start, end);
// Duration's toString always starts with "PT"; remove that here.
return activeFor.toString().substring(2);
}
public void printActiveWork(PrintWriter writer) {
final Instant now = Instant.now();
// The max number of keys in COMMITTING or COMMIT_QUEUED status to be shown.
final int maxCommitPending = 50;
int commitPendingCount = 0;
writer.println(
"<table border=\"1\" "
+ "style=\"border-collapse:collapse;padding:5px;border-spacing:5px;border:1px\">");
writer.println(
"<tr><th>Key</th><th>Token</th><th>Queued</th><th>Active For</th><th>State</th><th>State Active For</th></tr>");
// We use a StringBuilder in the synchronized section to buffer writes since the provided
// PrintWriter may block when flushing.
StringBuilder builder = new StringBuilder();
synchronized (activeWork) {
for (Map.Entry<ByteString, Deque<Work>> entry : activeWork.entrySet()) {
Queue<Work> queue = entry.getValue();
Preconditions.checkNotNull(queue);
Work work = queue.peek();
Preconditions.checkNotNull(work);
Windmill.WorkItem workItem = work.getWorkItem();
State state = work.getState();
if (state == State.COMMITTING || state == State.COMMIT_QUEUED) {
if (++commitPendingCount >= maxCommitPending) {
continue;
}
}
builder.append("<tr>");
builder.append("<td>");
builder.append(String.format("%016x", workItem.getShardingKey()));
builder.append("</td><td>");
builder.append(workItem.getWorkToken());
builder.append("</td><td>");
builder.append(queue.size() - 1);
builder.append("</td><td>");
builder.append(elapsedString(work.getStartTime(), now));
builder.append("</td><td>");
builder.append(state);
builder.append("</td></tr>\n");
builder.append(elapsedString(work.getStateStartTime(), now));
builder.append("</td></tr>\n");
}
}
writer.print(builder.toString());
writer.println("</table>");
if (commitPendingCount >= maxCommitPending) {
writer.println("<br>");
writer.print("Skipped keys in COMMITTING/COMMIT_QUEUED: ");
writer.println(commitPendingCount - maxCommitPending);
writer.println("<br>");
}
}
@Override
public void close() throws Exception {
ExecutionState executionState;
for (ConcurrentLinkedQueue<ExecutionState> queue : executionStateQueues.values()) {
while ((executionState = queue.poll()) != null) {
executionState.getWorkExecutor().close();
}
}
executionStateQueues.clear();
}
}
private static class ExecutionState {
public final DataflowWorkExecutor workExecutor;
public final StreamingModeExecutionContext context;
@Nullable public final Coder<?> keyCoder;
private final ExecutionStateTracker executionStateTracker;
public ExecutionState(
DataflowWorkExecutor workExecutor,
StreamingModeExecutionContext context,
Coder<?> keyCoder,
ExecutionStateTracker executionStateTracker) {
this.workExecutor = workExecutor;
this.context = context;
this.keyCoder = keyCoder;
this.executionStateTracker = executionStateTracker;
}
public DataflowWorkExecutor getWorkExecutor() {
return workExecutor;
}
public StreamingModeExecutionContext getContext() {
return context;
}
public ExecutionStateTracker getExecutionStateTracker() {
return executionStateTracker;
}
@Nullable
public Coder<?> getKeyCoder() {
return keyCoder;
}
}
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, \"&nbsp&nbsp\").replace(/\\n/g, \"<br>\"))</script>");
}
}
}
private class MetricsDataProvider implements StatusDataProvider {
@Override
public void appendSummaryHtml(PrintWriter writer) {
writer.println(
"Worker Threads: "
+ workUnitExecutor.getPoolSize()
+ "/"
+ workUnitExecutor.getMaximumPoolSize()
+ "<br>");
writer.println("Active Threads: " + workUnitExecutor.getActiveCount() + "<br>");
writer.println(
"Work Queue Size: "
+ workUnitExecutor.getQueue().size()
+ "/"
+ MAX_WORK_UNITS_QUEUED
+ "<br>");
writer.print("Commit Queue: ");
appendHumanizedBytes(commitQueue.weight(), writer);
writer.print(", ");
writer.print(commitQueue.size());
writer.println(" elements<br>");
writer.print("Active commit: ");
appendHumanizedBytes(activeCommitBytes.get(), 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");
}
}
}
}