diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java index 9401cc5f8ed9..41119742b15b 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java @@ -52,10 +52,12 @@ import org.apache.beam.runners.dataflow.worker.counters.NameContext; import org.apache.beam.runners.dataflow.worker.profiler.ScopedProfiler.ProfileScope; import org.apache.beam.runners.dataflow.worker.streaming.BoundedQueueExecutorWorkHandle; +import org.apache.beam.runners.dataflow.worker.streaming.KeyCommitTooLargeException; import org.apache.beam.runners.dataflow.worker.streaming.Watermarks; import org.apache.beam.runners.dataflow.worker.streaming.Work; import org.apache.beam.runners.dataflow.worker.streaming.config.StreamingGlobalConfig; import org.apache.beam.runners.dataflow.worker.streaming.config.StreamingGlobalConfigHandle; +import org.apache.beam.runners.dataflow.worker.streaming.harness.StreamingCounters; import org.apache.beam.runners.dataflow.worker.streaming.sideinput.SideInput; import org.apache.beam.runners.dataflow.worker.streaming.sideinput.SideInputState; import org.apache.beam.runners.dataflow.worker.streaming.sideinput.SideInputStateFetcher; @@ -77,6 +79,7 @@ import org.apache.beam.runners.dataflow.worker.windmill.state.WindmillTagEncodingV1; import org.apache.beam.runners.dataflow.worker.windmill.state.WindmillTagEncodingV2; import org.apache.beam.runners.dataflow.worker.windmill.state.WindmillTimerData; +import org.apache.beam.runners.dataflow.worker.windmill.work.processing.failures.FailureTracker; import org.apache.beam.sdk.annotations.Internal; import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.coders.CoderException; @@ -177,6 +180,9 @@ public class StreamingModeExecutionContext private final HotKeyLogger hotKeyLogger; private final boolean hotKeyLoggingEnabled; private final String stepName; + private final String systemName; + private final StreamingCounters streamingCounters; + private final FailureTracker failureTracker; private @Nullable Coder keyCoder; // Key switch listener to delegate MDC logging context and thread name updates @@ -212,6 +218,9 @@ public StreamingModeExecutionContext( HotKeyLogger hotKeyLogger, boolean hotKeyLoggingEnabled, String stepName, + String systemName, + StreamingCounters streamingCounters, + FailureTracker failureTracker, String sourceBytesProcessCounterName, SideInputStateFetcherFactory sideInputStateFetcherFactory) { super( @@ -231,6 +240,9 @@ public StreamingModeExecutionContext( this.hotKeyLogger = checkNotNull(hotKeyLogger); this.hotKeyLoggingEnabled = hotKeyLoggingEnabled; this.stepName = checkNotNull(stepName); + this.systemName = checkNotNull(systemName); + this.streamingCounters = checkNotNull(streamingCounters); + this.failureTracker = checkNotNull(failureTracker); this.sourceBytesProcessCounterName = checkNotNull(sourceBytesProcessCounterName); this.sideInputStateFetcherFactory = sideInputStateFetcherFactory; StreamingGlobalConfig config = globalConfigHandle.getConfig(); @@ -669,6 +681,52 @@ private void flushStateInternal() { getOutputBuilder() .setSourceBytesProcessed(computeSourceBytesProcessed(sourceBytesProcessCounterName)); + + validateCommitRequestSize(); + } + + private void validateCommitRequestSize() { + Windmill.WorkItemCommitRequest.Builder currentBuilder = getOutputBuilder(); + Work currentWork = getWork(); + long byteLimit = operationalLimits.getMaxWorkItemCommitBytes(); + Windmill.WorkItemCommitRequest commitRequest = currentBuilder.build(); + 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. + // Commit is too large if overflow has occurred or the commitSize has exceeded the allowed + // commit byte limit. + streamingCounters.windmillMaxObservedWorkItemCommitBytes().addValue(estimatedCommitSize); + if (commitSize >= 0 && commitSize < byteLimit) { + return; + } + + KeyCommitTooLargeException e = + KeyCommitTooLargeException.causedBy( + systemName, byteLimit, commitRequest, key, hotKeyLoggingEnabled); + failureTracker.trackFailure(systemName, currentWork.getWorkItem(), 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 + Windmill.WorkItemCommitRequest.Builder truncationBuilder = + buildWorkItemTruncationRequestBuilder(currentWork, estimatedCommitSize); + for (int i = 0; i < outputBuilders.size(); i++) { + if (outputBuilders.get(i) == currentBuilder) { + outputBuilders.set(i, truncationBuilder); + break; + } + } + this.outputBuilder = truncationBuilder; + } + + private Windmill.WorkItemCommitRequest.Builder buildWorkItemTruncationRequestBuilder( + Work work, int estimatedCommitSize) { + Windmill.WorkItemCommitRequest.Builder outputBuilder = createOutputBuilder(work); + outputBuilder.setExceedsMaxWorkItemCommitBytes(true); + outputBuilder.setEstimatedWorkItemCommitBytes(estimatedCommitSize); + return outputBuilder; } private final long computeSourceBytesProcessed(String sourceBytesCounterName) { diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/KeyCommitTooLargeException.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/KeyCommitTooLargeException.java index 4d6ae8a208c1..c57921abbf46 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/KeyCommitTooLargeException.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/KeyCommitTooLargeException.java @@ -18,16 +18,38 @@ package org.apache.beam.runners.dataflow.worker.streaming; import org.apache.beam.runners.dataflow.worker.windmill.Windmill; +import org.checkerframework.checker.nullness.qual.Nullable; public final class KeyCommitTooLargeException extends Exception { public static KeyCommitTooLargeException causedBy( - String computationId, long byteLimit, Windmill.WorkItemCommitRequest request) { + String stageName, long byteLimit, Windmill.WorkItemCommitRequest request) { + return causedBy(stageName, byteLimit, request, null, false); + } + + public static KeyCommitTooLargeException causedBy( + String stageName, + long byteLimit, + Windmill.WorkItemCommitRequest request, + boolean hotKeyLoggingEnabled) { + return causedBy(stageName, byteLimit, request, null, hotKeyLoggingEnabled); + } + + public static KeyCommitTooLargeException causedBy( + String stageName, + long byteLimit, + Windmill.WorkItemCommitRequest request, + @Nullable Object decodedKey, + boolean hotKeyLoggingEnabled) { StringBuilder message = new StringBuilder(); message.append("Commit request for stage "); - message.append(computationId); + message.append(stageName); message.append(" and sharding key "); - message.append(request.getShardingKey()); + message.append(Long.toUnsignedString(request.getShardingKey())); + if (decodedKey != null && hotKeyLoggingEnabled) { + message.append(" and key "); + message.append(decodedKey); + } if (request.getSerializedSize() > 0) { message.append( " has size " @@ -38,9 +60,8 @@ public static KeyCommitTooLargeException causedBy( 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." + ". 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." + " See https://cloud.google.com/dataflow/docs/guides/common-errors#key-commit-too-large-exception."); return new KeyCommitTooLargeException(message.toString()); } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactory.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactory.java index 4a52d9fde771..0c3591102f9c 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactory.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactory.java @@ -49,11 +49,13 @@ import org.apache.beam.runners.dataflow.worker.streaming.ComputationWorkExecutor; import org.apache.beam.runners.dataflow.worker.streaming.StageInfo; import org.apache.beam.runners.dataflow.worker.streaming.config.StreamingGlobalConfigHandle; +import org.apache.beam.runners.dataflow.worker.streaming.harness.StreamingCounters; import org.apache.beam.runners.dataflow.worker.streaming.sideinput.SideInputStateFetcherFactory; import org.apache.beam.runners.dataflow.worker.util.common.worker.MapTaskExecutor; 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.state.WindmillStateCache; +import org.apache.beam.runners.dataflow.worker.windmill.work.processing.failures.FailureTracker; import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.coders.KvCoder; import org.apache.beam.sdk.fn.IdGenerator; @@ -84,6 +86,8 @@ final class ComputationWorkExecutorFactory { private final SinkRegistry sinkRegistry; private final DataflowExecutionStateSampler sampler; private final CounterSet pendingDeltaCounters; + private final StreamingCounters streamingCounters; + private final FailureTracker failureTracker; /** * Function which converts map tasks to their network representation for execution. @@ -108,7 +112,8 @@ final class ComputationWorkExecutorFactory { ReaderCache readerCache, Function stateCacheFactory, DataflowExecutionStateSampler sampler, - CounterSet pendingDeltaCounters, + StreamingCounters streamingCounters, + FailureTracker failureTracker, IdGenerator idGenerator, StreamingGlobalConfigHandle globalConfigHandle, HotKeyLogger hotKeyLogger, @@ -122,7 +127,9 @@ final class ComputationWorkExecutorFactory { this.readerRegistry = ReaderRegistry.defaultRegistry(); this.sinkRegistry = SinkRegistry.defaultRegistry(); this.sampler = sampler; - this.pendingDeltaCounters = pendingDeltaCounters; + this.streamingCounters = streamingCounters; + this.failureTracker = failureTracker; + this.pendingDeltaCounters = streamingCounters.pendingDeltaCounters(); this.mapTaskToNetwork = new MapTaskToNetworkFunction(idGenerator); this.maxSinkBytes = hasExperiment(options, DISABLE_SINK_BYTE_LIMIT_EXPERIMENT) @@ -286,6 +293,9 @@ private StreamingModeExecutionContext createExecutionContext( hotKeyLogger, hotKeyLoggingEnabled, stepName, + stageInfo.systemName(), + streamingCounters, + failureTracker, computationState.sourceBytesProcessCounterName(), sideInputStateFetcherFactory); } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java index e9b85d720d2b..990393f51f8f 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java @@ -43,7 +43,6 @@ import org.apache.beam.runners.dataflow.worker.streaming.ComputationState; import org.apache.beam.runners.dataflow.worker.streaming.ComputationWorkExecutor; import org.apache.beam.runners.dataflow.worker.streaming.ExecutableWork; -import org.apache.beam.runners.dataflow.worker.streaming.KeyCommitTooLargeException; import org.apache.beam.runners.dataflow.worker.streaming.StageInfo; import org.apache.beam.runners.dataflow.worker.streaming.Watermarks; import org.apache.beam.runners.dataflow.worker.streaming.Work; @@ -81,36 +80,30 @@ public class StreamingWorkScheduler { private final Supplier clock; private final ComputationWorkExecutorFactory computationWorkExecutorFactory; - private final FailureTracker failureTracker; private final WorkFailureProcessor workFailureProcessor; private final StreamingCommitFinalizer commitFinalizer; private final StreamingCounters streamingCounters; private final ConcurrentMap stageInfoMap; private final DataflowExecutionStateSampler sampler; - private final StreamingGlobalConfigHandle globalConfigHandle; private final BoundedQueueExecutor workExecutor; public StreamingWorkScheduler( Supplier clock, BoundedQueueExecutor workExecutor, ComputationWorkExecutorFactory computationWorkExecutorFactory, - FailureTracker failureTracker, WorkFailureProcessor workFailureProcessor, StreamingCommitFinalizer commitFinalizer, StreamingCounters streamingCounters, ConcurrentMap stageInfoMap, - DataflowExecutionStateSampler sampler, - StreamingGlobalConfigHandle globalConfigHandle) { + DataflowExecutionStateSampler sampler) { this.clock = clock; this.workExecutor = workExecutor; this.computationWorkExecutorFactory = computationWorkExecutorFactory; - this.failureTracker = failureTracker; this.workFailureProcessor = workFailureProcessor; this.commitFinalizer = commitFinalizer; this.streamingCounters = streamingCounters; this.stageInfoMap = stageInfoMap; this.sampler = sampler; - this.globalConfigHandle = globalConfigHandle; } public static StreamingWorkScheduler create( @@ -139,7 +132,8 @@ public static StreamingWorkScheduler create( readerCache, stateCacheFactory, sampler, - streamingCounters.pendingDeltaCounters(), + streamingCounters, + failureTracker, idGenerator, globalConfigHandle, hotKeyLogger, @@ -149,13 +143,11 @@ public static StreamingWorkScheduler create( clock, workExecutor, computationWorkExecutorFactory, - failureTracker, workFailureProcessor, StreamingCommitFinalizer.create(workExecutor, commitFinalizerCleanupExecutor), streamingCounters, stageInfoMap, - sampler, - globalConfigHandle); + sampler); } private static long computeShuffleBytesRead(Windmill.WorkItem workItem) { @@ -175,14 +167,6 @@ private static Windmill.WorkItemCommitRequest.Builder initializeOutputBuilder( .setCacheToken(workItem.getCacheToken()); } - private static Windmill.WorkItemCommitRequest buildWorkItemTruncationRequest( - ByteString key, Windmill.WorkItem workItem, int estimatedCommitSize) { - Windmill.WorkItemCommitRequest.Builder outputBuilder = initializeOutputBuilder(key, workItem); - outputBuilder.setExceedsMaxWorkItemCommitBytes(true); - outputBuilder.setEstimatedWorkItemCommitBytes(estimatedCommitSize); - return outputBuilder.build(); - } - /** Sets the stage name and workId of the Thread executing the {@link Work} for logging. */ private static void setUpWorkLoggingContext(String workLatencyTrackingId, String computationId) { setLoggingContextWorkId(workLatencyTrackingId); @@ -305,33 +289,6 @@ private void processWork( } } - private Windmill.WorkItemCommitRequest validateCommitRequestSize( - Windmill.WorkItemCommitRequest commitRequest, - String computationId, - Windmill.WorkItem workItem) { - long byteLimit = globalConfigHandle.getConfig().operationalLimits().getMaxWorkItemCommitBytes(); - 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. - // Commit is too large if overflow has occurred or the commitSize has exceeded the allowed - // commit byte limit. - streamingCounters.windmillMaxObservedWorkItemCommitBytes().addValue(estimatedCommitSize); - if (commitSize >= 0 && commitSize < byteLimit) { - return commitRequest; - } - - KeyCommitTooLargeException e = - KeyCommitTooLargeException.causedBy(computationId, byteLimit, commitRequest); - failureTracker.trackFailure(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 - return buildWorkItemTruncationRequest(workItem.getKey(), workItem, estimatedCommitSize); - } - private void recordProcessingStats( List workBatch, List workItemCommits, @@ -448,17 +405,13 @@ private void commitWorkBatch( private void commitSingleKeyWork( ComputationState computationState, Work work, Windmill.WorkItemCommitRequest commitRequest) { - // Validate the commit request, possibly requesting truncation if the commitSize is too large. - Windmill.WorkItemCommitRequest validatedCommitRequest = - validateCommitRequestSize( - commitRequest, computationState.getComputationId(), work.getWorkItem()); work.setState(Work.State.COMMIT_QUEUED); - validatedCommitRequest = - validatedCommitRequest + Windmill.WorkItemCommitRequest commitRequestWithAttributions = + commitRequest .toBuilder() .addAllPerWorkItemLatencyAttributions(work.getLatencyAttributions(sampler)) .build(); - work.queueCommit(validatedCommitRequest, computationState); + work.queueCommit(commitRequestWithAttributions, computationState); } private void recordProcessingTime( 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 23730bc57705..fc8f12c45cb3 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 @@ -1290,8 +1290,9 @@ public void testKeyTokenInvalidException() throws Exception { worker.stop(); } - @Test - public void testKeyCommitTooLargeException() throws Exception { + private void runKeyCommitTooLargeExceptionTest( + StreamingDataflowWorkerTestParams.Builder workerParams, boolean expectKeyInErrorMessage) + throws Exception { KvCoder kvCoder = KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of()); List instructions = @@ -1304,7 +1305,7 @@ public void testKeyCommitTooLargeException() throws Exception { StreamingDataflowWorker worker = makeWorker( - defaultWorkerParams() + workerParams .setInstructions(instructions) .setStreamingGlobalConfig( StreamingGlobalConfig.builder() @@ -1337,18 +1338,14 @@ public void testKeyCommitTooLargeException() throws Exception { .build(), removeDynamicFields(largeCommit)); - // Check this explicitly since the estimated commit bytes weren't actually - // checked against an expected value in the previous step assertTrue(largeCommit.getEstimatedWorkItemCommitBytes() > 1000); - // Spam worker updates a few times. int maxTries = 10; while (--maxTries > 0) { worker.reportPeriodicWorkerUpdatesForTest(); Uninterruptibles.sleepUninterruptibly(100, TimeUnit.MILLISECONDS); } - // We should see an exception reported for the large commit but not the small one. ArgumentCaptor workItemStatusCaptor = ArgumentCaptor.forClass(WorkItemStatus.class); verify(mockWorkUnitClient, atLeast(2)).reportWorkItemStatus(workItemStatusCaptor.capture()); @@ -1360,12 +1357,35 @@ public void testKeyCommitTooLargeException() throws Exception { foundErrors = true; String errorMessage = status.getErrors().get(0).getMessage(); assertThat(errorMessage, Matchers.containsString("KeyCommitTooLargeException")); + assertThat(errorMessage, Matchers.containsString("Commit request for stage computation")); + if (expectKeyInErrorMessage) { + assertThat(errorMessage, Matchers.containsString("and key large_key")); + } else { + assertThat(errorMessage, Matchers.not(Matchers.containsString("and key large_key"))); + } } } assertTrue(foundErrors); worker.stop(); } + @Test + public void testKeyCommitTooLargeException() throws Exception { + runKeyCommitTooLargeExceptionTest(defaultWorkerParams(), /* expectKeyInErrorMessage= */ false); + } + + @Test + public void testKeyCommitTooLargeException_withHotKeyLoggingEnabled() throws Exception { + runKeyCommitTooLargeExceptionTest( + defaultWorkerParams("--hotKeyLoggingEnabled=true"), /* expectKeyInErrorMessage= */ true); + } + + @Test + public void testKeyCommitTooLargeException_withHotKeyLoggingDisabled() throws Exception { + runKeyCommitTooLargeExceptionTest( + defaultWorkerParams("--hotKeyLoggingEnabled=false"), /* expectKeyInErrorMessage= */ false); + } + @Test public void testOutputKeyTooLargeException() throws Exception { KvCoder kvCoder = KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of()); diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java index 561596f68d0f..3d11e5097463 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java @@ -63,6 +63,7 @@ import org.apache.beam.runners.dataflow.worker.streaming.config.FakeGlobalConfigHandle; import org.apache.beam.runners.dataflow.worker.streaming.config.StreamingGlobalConfig; import org.apache.beam.runners.dataflow.worker.streaming.config.StreamingGlobalConfigHandle; +import org.apache.beam.runners.dataflow.worker.streaming.harness.StreamingCounters; import org.apache.beam.runners.dataflow.worker.streaming.sideinput.SideInputStateFetcherFactory; import org.apache.beam.runners.dataflow.worker.util.common.worker.WorkExecutor; import org.apache.beam.runners.dataflow.worker.windmill.Windmill; @@ -71,6 +72,7 @@ import org.apache.beam.runners.dataflow.worker.windmill.state.WindmillStateReader; import org.apache.beam.runners.dataflow.worker.windmill.state.WindmillTagEncodingV1; import org.apache.beam.runners.dataflow.worker.windmill.state.WindmillTagEncodingV2; +import org.apache.beam.runners.dataflow.worker.windmill.work.processing.failures.FailureTracker; import org.apache.beam.runners.dataflow.worker.windmill.work.refresh.HeartbeatSender; import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.coders.Coder; @@ -142,6 +144,9 @@ private StreamingModeExecutionContext createExecutionContext( new HotKeyLogger(), /*hotKeyLoggingEnabled=*/ false, /*stepName=*/ "stepName", + /*systemName=*/ "systemName", + StreamingCounters.create(), + mock(FailureTracker.class), "sourceBytesProcessCounterName", SideInputStateFetcherFactory.fromOptions(options)); } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSourcesTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSourcesTest.java index 31ea1bab07af..2da52bea8eba 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSourcesTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSourcesTest.java @@ -94,6 +94,7 @@ import org.apache.beam.runners.dataflow.worker.streaming.config.FixedGlobalConfigHandle; import org.apache.beam.runners.dataflow.worker.streaming.config.StreamingGlobalConfig; import org.apache.beam.runners.dataflow.worker.streaming.config.StreamingGlobalConfigHandle; +import org.apache.beam.runners.dataflow.worker.streaming.harness.StreamingCounters; import org.apache.beam.runners.dataflow.worker.streaming.sideinput.SideInputStateFetcherFactory; import org.apache.beam.runners.dataflow.worker.testing.TestCountingSource; import org.apache.beam.runners.dataflow.worker.util.common.worker.NativeReader; @@ -103,6 +104,7 @@ import org.apache.beam.runners.dataflow.worker.windmill.client.getdata.FakeGetDataClient; 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; import org.apache.beam.runners.dataflow.worker.windmill.work.refresh.HeartbeatSender; import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.coders.BigEndianIntegerCoder; @@ -641,6 +643,9 @@ public void testReadUnboundedReader() throws Exception { new HotKeyLogger(), /*hotKeyLoggingEnabled=*/ false, /*stepName=*/ "stepName", + /*systemName=*/ "systemName", + StreamingCounters.create(), + mock(FailureTracker.class), "sourceBytesProcessCounterName", SideInputStateFetcherFactory.fromOptions(options)); @@ -1014,6 +1019,9 @@ public void testFailedWorkItemsAbort() throws Exception { new HotKeyLogger(), /*hotKeyLoggingEnabled=*/ false, /*stepName=*/ "stepName", + /*systemName=*/ "systemName", + StreamingCounters.create(), + mock(FailureTracker.class), "sourceBytesProcessCounterName", SideInputStateFetcherFactory.fromOptions(options));