diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/AssignWindowsParDoFnFactory.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/AssignWindowsParDoFnFactory.java index 673130a6048a..897d707456e2 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/AssignWindowsParDoFnFactory.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/AssignWindowsParDoFnFactory.java @@ -31,6 +31,7 @@ import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; import org.apache.beam.sdk.transforms.windowing.WindowFn; +import org.apache.beam.sdk.util.DoFnInfo; import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.sdk.values.WindowedValue; import org.apache.beam.sdk.values.WindowedValues; @@ -131,5 +132,10 @@ public void finishBundle() throws Exception { public void abort() throws Exception { receiver = null; } + + @Override + public @Nullable DoFnInfo getDoFnInfo() { + return null; + } } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/BatchModeUngroupingParDoFn.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/BatchModeUngroupingParDoFn.java index c3cf6d9e67e2..df1438b36e00 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/BatchModeUngroupingParDoFn.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/BatchModeUngroupingParDoFn.java @@ -19,8 +19,10 @@ import org.apache.beam.runners.dataflow.worker.util.common.worker.ParDoFn; import org.apache.beam.runners.dataflow.worker.util.common.worker.Receiver; +import org.apache.beam.sdk.util.DoFnInfo; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.WindowedValue; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; /** @@ -85,4 +87,9 @@ public void finishBundle() throws Exception { public void abort() throws Exception { underlyingParDoFn.abort(); } + + @Override + public @Nullable DoFnInfo getDoFnInfo() throws Exception { + return underlyingParDoFn.getDoFnInfo(); + } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/CreateIsmShardKeyAndSortKeyDoFnFactory.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/CreateIsmShardKeyAndSortKeyDoFnFactory.java index afe2e28b2a67..9d28b4720aa0 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/CreateIsmShardKeyAndSortKeyDoFnFactory.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/CreateIsmShardKeyAndSortKeyDoFnFactory.java @@ -32,10 +32,12 @@ import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.util.DoFnInfo; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.sdk.values.WindowedValue; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.checkerframework.checker.nullness.qual.Nullable; /** * A {@link ParDoFnFactory} that creates a system {@link ParDoFn} responsible for limiting the users @@ -122,5 +124,10 @@ public void finishBundle() {} @Override public void abort() {} + + @Override + public @Nullable DoFnInfo getDoFnInfo() { + return null; + } } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ForwardingParDoFn.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ForwardingParDoFn.java index 3a864d6caf2f..0c16325e56c4 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ForwardingParDoFn.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ForwardingParDoFn.java @@ -19,6 +19,7 @@ import org.apache.beam.runners.dataflow.worker.util.common.worker.ParDoFn; import org.apache.beam.runners.dataflow.worker.util.common.worker.Receiver; +import org.apache.beam.sdk.util.DoFnInfo; import org.checkerframework.checker.nullness.qual.Nullable; /** @@ -62,4 +63,9 @@ public void finishBundle() throws Exception { public void abort() throws Exception { delegate.abort(); } + + @Override + public @Nullable DoFnInfo getDoFnInfo() throws Exception { + return delegate.getDoFnInfo(); + } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/GroupAlsoByWindowsParDoFn.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/GroupAlsoByWindowsParDoFn.java index 6c92b7f10ea1..f4364f006c51 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/GroupAlsoByWindowsParDoFn.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/GroupAlsoByWindowsParDoFn.java @@ -34,6 +34,7 @@ import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.options.StreamingOptions; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.util.DoFnInfo; import org.apache.beam.sdk.util.WindowedValueReceiver; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.PCollectionView; @@ -160,6 +161,11 @@ public void abort() throws Exception { fnRunner = null; } + @Override + public @Nullable DoFnInfo getDoFnInfo() { + return null; + } + /** * Composes and returns a {@link DoFnRunner} based on the parameters. * diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/IntrinsicMapTaskExecutorFactory.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/IntrinsicMapTaskExecutorFactory.java index 3ea29787eb3b..ee3e6ac691a3 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/IntrinsicMapTaskExecutorFactory.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/IntrinsicMapTaskExecutorFactory.java @@ -32,6 +32,7 @@ import java.util.List; import java.util.function.Function; import org.apache.beam.runners.core.ElementByteSizeObservable; +import org.apache.beam.runners.core.SplittableParDoViaKeyedWorkItems.ProcessFn; import org.apache.beam.runners.dataflow.util.CloudObject; import org.apache.beam.runners.dataflow.util.CloudObjects; import org.apache.beam.runners.dataflow.worker.counters.CounterFactory; @@ -64,6 +65,9 @@ import org.apache.beam.sdk.fn.IdGenerator; import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.options.StreamingOptions; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.reflect.DoFnSignatures; +import org.apache.beam.sdk.util.DoFnInfo; import org.apache.beam.sdk.util.common.ElementByteSizeObserver; import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.sdk.values.WindowedValues.WindowedValueCoder; @@ -141,11 +145,9 @@ public DataflowMapTaskExecutor create( Iterables.filter(Networks.topologicalOrder(network), OperationNode.class)) { topoSortedOperations.add(node.getOperation()); - for (Node predecessor : + for (OutputReceiverNode predecessor : Iterables.filter(network.predecessors(node), OutputReceiverNode.class)) { - ((OutputReceiverNode) predecessor) - .getOutputReceiver() - .addOutput((Receiver) node.getOperation()); + predecessor.getOutputReceiver().addOutput((Receiver) node.getOperation()); } } @@ -248,7 +250,7 @@ OperationNode createWriteOperation( return OperationNode.create(WriteOperation.create(sink, EMPTY_OUTPUT_RECEIVER_ARRAY, context)); } - private ParDoFnFactory parDoFnFactory = new DefaultParDoFnFactory(); + private final ParDoFnFactory parDoFnFactory = new DefaultParDoFnFactory(); private OperationNode createParDoOperation( Network network, @@ -272,18 +274,44 @@ private OperationNode createParDoOperation( } successorOffset += 1; } + CloudObject cloudUserFn = CloudObject.fromSpec(parDo.getUserFn()); ParDoFn fn = parDoFnFactory.create( options, - CloudObject.fromSpec(parDo.getUserFn()), + cloudUserFn, parDo.getSideInputs(), mainOutputTag, outputTagsToReceiverIndicesBuilder.build(), executionContext, operationContext); + // An SDF's SplittableProcessFn behaves like a source, so we need to report its bytes read even + // though it's not a read operator for autoscaling to work properly. + ImmutableMap.Builder> bytesReadCodersBuilder = ImmutableMap.builder(); + if (isSplittable(fn)) { + for (Node successor : network.successors(node)) { + if (successor instanceof OutputReceiverNode) { + OutputReceiverNode outputReceiverNode = (OutputReceiverNode) successor; + bytesReadCodersBuilder.put( + outputReceiverNode.getOutputReceiver(), outputReceiverNode.getCoder()); + } + } + } OutputReceiver[] receivers = getOutputReceivers(network, node); - return OperationNode.create(new ParDoOperation(fn, receivers, operationContext)); + return OperationNode.create( + new ParDoOperation(fn, receivers, operationContext, bytesReadCodersBuilder.build())); + } + + private static boolean isSplittable(ParDoFn fn) throws Exception { + DoFnInfo doFnInfo = fn.getDoFnInfo(); + if (doFnInfo == null) { + return false; + } + DoFn doFn = doFnInfo.getDoFn(); + if (doFn instanceof ProcessFn) { + doFn = ((ProcessFn) doFn).getFn(); + } + return DoFnSignatures.isSplittable(doFn); } private static TupleTag tupleTag(MultiOutputInfo multiOutputInfo) { diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PairWithConstantKeyDoFnFactory.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PairWithConstantKeyDoFnFactory.java index 425184a4a126..5a0406d9cdd7 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PairWithConstantKeyDoFnFactory.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PairWithConstantKeyDoFnFactory.java @@ -32,9 +32,11 @@ import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.util.CoderUtils; +import org.apache.beam.sdk.util.DoFnInfo; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.sdk.values.WindowedValue; +import org.checkerframework.checker.nullness.qual.Nullable; /** * A {@link ParDoFnFactory} which returns a {@link ParDoFn} that transforms all {@code @@ -106,5 +108,10 @@ public void finishBundle() {} @Override public void abort() {} + + @Override + public @Nullable DoFnInfo getDoFnInfo() { + return null; + } } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PartialGroupByKeyParDoFns.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PartialGroupByKeyParDoFns.java index a6d7810412a1..c81ad13ef38b 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PartialGroupByKeyParDoFns.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PartialGroupByKeyParDoFns.java @@ -43,6 +43,7 @@ import org.apache.beam.sdk.state.BagState; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; import org.apache.beam.sdk.util.AppliedCombineFn; +import org.apache.beam.sdk.util.DoFnInfo; import org.apache.beam.sdk.util.SerializableUtils; import org.apache.beam.sdk.util.common.ElementByteSizeObserver; import org.apache.beam.sdk.values.KV; @@ -327,6 +328,11 @@ public void finishBundle() throws Exception { @Override public void abort() throws Exception {} + + @Override + public @Nullable DoFnInfo getDoFnInfo() { + return null; + } } static class StreamingSideInputPGBKParDoFn @@ -391,5 +397,10 @@ public void finishBundle() throws Exception { @Override public void abort() throws Exception {} + + @Override + public @Nullable DoFnInfo getDoFnInfo() { + return null; + } } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ReifyTimestampAndWindowsParDoFnFactory.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ReifyTimestampAndWindowsParDoFnFactory.java index 0438b525b6b9..eb3d6fcdfeab 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ReifyTimestampAndWindowsParDoFnFactory.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ReifyTimestampAndWindowsParDoFnFactory.java @@ -26,6 +26,7 @@ import org.apache.beam.runners.dataflow.worker.util.common.worker.ParDoFn; import org.apache.beam.runners.dataflow.worker.util.common.worker.Receiver; import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.util.DoFnInfo; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.sdk.values.WindowedValue; @@ -98,5 +99,10 @@ public void finishBundle() throws Exception { public void abort() throws Exception { this.receiver = null; } + + @Override + public @Nullable DoFnInfo getDoFnInfo() { + return null; + } } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFn.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFn.java index badc1533b2ad..4b43a84174f5 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFn.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFn.java @@ -97,9 +97,9 @@ protected void onStartKey() { sideInputProcessor = new StreamingSideInputProcessor<>( new StreamingSideInputFetcher( - helpers.fnInfo.getSideInputViews(), - helpers.fnInfo.getInputCoder(), - (WindowingStrategy) helpers.fnInfo.getWindowingStrategy(), + helpers.activeFnInfo.getSideInputViews(), + helpers.activeFnInfo.getInputCoder(), + (WindowingStrategy) helpers.activeFnInfo.getWindowingStrategy(), (StreamingModeExecutionContext.StreamingModeStepContext) helpers.userStepContext)); @@ -112,7 +112,7 @@ protected void onStartKey() { // These elements are now processed. Register cleanup timers for all the unblocked // windows. helpers.registerStateCleanup( - (WindowingStrategy) getDoFnInfo().getWindowingStrategy(), + (WindowingStrategy) helpers.activeFnInfo.getWindowingStrategy(), (Collection) unblockedElement.getWindows()); } } @@ -163,7 +163,7 @@ protected void onProcessWindowedValue(WindowedValue elem) { } if (hasState) { helpers.registerStateCleanup( - (WindowingStrategy) getDoFnInfo().getWindowingStrategy(), windowsProcessed); + (WindowingStrategy) helpers.activeFnInfo.getWindowingStrategy(), windowsProcessed); } } @@ -171,13 +171,15 @@ protected void onProcessWindowedValue(WindowedValue elem) { public void processTimers() throws Exception { // Note: We need to get windowCoder to decode the timers. If we haven't already deserialized - // the fnInfo, we peek at a new instance to retrieve that. If this extra deserialization becomes - // excessively costly, we could either (1) have the DoFnInstanceManager remember the associated - // windowCoder (allowing us to get it without a DoFnInfo instance) or (2) check whether timers - // exist without actually decoding them. + // the activeFnInfo, we peek at a new instance to retrieve that. If this extra deserialization + // becomes excessively costly, we could either (1) have the DoFnInstanceManager remember the + // associated windowCoder (allowing us to get it without a DoFnInfo instance) or (2) check + // whether timers exist without actually decoding them. Coder windowCoder = (Coder) - (helpers.fnInfo != null ? helpers.fnInfo : helpers.doFnInstanceManager.peek()) + (helpers.activeFnInfo != null + ? helpers.activeFnInfo + : helpers.doFnInstanceManager.peek()) .getWindowingStrategy() .getWindowFn() .windowCoder(); @@ -209,6 +211,12 @@ public void abort() throws Exception { helpers.abort(); } + /** Returns the {@link DoFnInfo} used by this {@link SimpleParDoFn}. */ + @Override + public DoFnInfo getDoFnInfo() throws Exception { + return helpers.doFnInstanceManager.peek(); + } + /** * Returns the {@link DoFnInfo} currently being used by this {@link SimpleParDoFn}. * @@ -216,7 +224,7 @@ public void abort() throws Exception { * finished. */ @VisibleForTesting - @Nullable DoFnInfo getDoFnInfo() { - return helpers.fnInfo; + @Nullable DoFnInfo getActiveDoFnInfo() { + return helpers.activeFnInfo; } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFnHelpers.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFnHelpers.java index 15bfba9bbc49..a69b3af7e256 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFnHelpers.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFnHelpers.java @@ -97,7 +97,7 @@ class SimpleParDoFnHelpers { // Various DoFn helpers, null between bundles @Nullable DoFnRunner fnRunner; - @Nullable DoFnInfo fnInfo; + @Nullable DoFnInfo activeFnInfo; private Receiver @Nullable [] receivers; // This may additionally be null if it is not a real DoFn but an OldDoFn or @@ -212,20 +212,20 @@ public void output(TupleTag tag, WindowedValue output) { } } }; - fnInfo = (DoFnInfo) doFnInstanceManager.get(); - fnSignature = DoFnSignatures.getSignature(fnInfo.getDoFn().getClass()); + activeFnInfo = (DoFnInfo) doFnInstanceManager.get(); + fnSignature = DoFnSignatures.getSignature(activeFnInfo.getDoFn().getClass()); fnRunner = runnerFactory.createRunner( - fnInfo.getDoFn(), + activeFnInfo.getDoFn(), options, mainOutputTag, sideOutputTags, - fnInfo.getSideInputViews(), + activeFnInfo.getSideInputViews(), sideInputReader, - fnInfo.getInputCoder(), - fnInfo.getOutputCoders(), - fnInfo.getWindowingStrategy(), + activeFnInfo.getInputCoder(), + activeFnInfo.getOutputCoders(), + activeFnInfo.getWindowingStrategy(), stepContext, userStepContext, outputManager, @@ -240,18 +240,18 @@ void finishBundle(StreamingSideInputProcessor sideInputProcessor) throws E if (sideInputProcessor != null) { sideInputProcessor.handleFinishBundle(); } - doFnInstanceManager.complete(fnInfo); + doFnInstanceManager.complete(activeFnInfo); fnRunner = null; - fnInfo = null; + activeFnInfo = null; fnSignature = null; sideInputProcessor = null; } } void abort() throws Exception { - doFnInstanceManager.abort(fnInfo); + doFnInstanceManager.abort(activeFnInfo); fnRunner = null; - fnInfo = null; + activeFnInfo = null; } @VisibleForTesting static final String CLEANUP_TIMER_ID = "cleanup-timer"; @@ -366,7 +366,8 @@ private void processSystemTimer( } BoundedWindow window = ((StateNamespaces.WindowNamespace) timer.getNamespace()).getWindow(); - Instant targetTime = earliestAllowableCleanupTime(window, fnInfo.getWindowingStrategy()); + Instant targetTime = + earliestAllowableCleanupTime(window, activeFnInfo.getWindowingStrategy()); checkState( !targetTime.isAfter(timer.getTimestamp()), @@ -385,12 +386,12 @@ private void processSystemTimer( try { tag = StateTags.tagForSpec( - stateDecl.id(), (StateSpec) stateDecl.field().get(fnInfo.getDoFn())); + stateDecl.id(), (StateSpec) stateDecl.field().get(activeFnInfo.getDoFn())); } catch (IllegalAccessException e) { throw new RuntimeException( String.format( "Error accessing %s for %s", - StateSpec.class.getName(), fnInfo.getDoFn().getClass().getName()), + StateSpec.class.getName(), activeFnInfo.getDoFn().getClass().getName()), e); } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputParDoFn.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputParDoFn.java index 5e61ee3a57f3..c0081e35e4bf 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputParDoFn.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputParDoFn.java @@ -109,9 +109,9 @@ protected void onStartKey() { sideInputProcessor = new StreamingSideInputProcessor<>( new StreamingSideInputFetcher( - helpers.fnInfo.getSideInputViews(), + helpers.activeFnInfo.getSideInputViews(), inputCoder, - (WindowingStrategy) helpers.fnInfo.getWindowingStrategy(), + (WindowingStrategy) helpers.activeFnInfo.getWindowingStrategy(), (StreamingModeExecutionContext.StreamingModeStepContext) helpers.userStepContext)); } @@ -138,7 +138,7 @@ protected void onStartKey() { // These elements are now processed. Register cleanup timers for all the unblocked // windows. helpers.registerStateCleanup( - (WindowingStrategy) getDoFnInfo().getWindowingStrategy(), windows); + (WindowingStrategy) helpers.activeFnInfo.getWindowingStrategy(), windows); } }); } @@ -169,13 +169,15 @@ public void processElement(Object untypedElem) throws Exception { public void processTimers() throws Exception { // Note: We need to get windowCoder to decode the timers. If we haven't already deserialized - // the fnInfo, we peek at a new instance to retrieve that. If this extra deserialization becomes - // excessively costly, we could either (1) have the DoFnInstanceManager remember the associated - // windowCoder (allowing us to get it without a DoFnInfo instance) or (2) check whether timers - // exist without actually decoding them. + // the activeFnInfo, we peek at a new instance to retrieve that. If this extra deserialization + // becomes excessively costly, we could either (1) have the DoFnInstanceManager remember the + // associated windowCoder (allowing us to get it without a DoFnInfo instance) or (2) check + // whether timers exist without actually decoding them. Coder windowCoder = (Coder) - (helpers.fnInfo != null ? helpers.fnInfo : helpers.doFnInstanceManager.peek()) + (helpers.activeFnInfo != null + ? helpers.activeFnInfo + : helpers.doFnInstanceManager.peek()) .getWindowingStrategy() .getWindowFn() .windowCoder(); @@ -231,18 +233,25 @@ protected void onProcessWindowedValue(WindowedValue> el } if (hasState) { helpers.registerStateCleanup( - (WindowingStrategy) getDoFnInfo().getWindowingStrategy(), windowsProcessed); + (WindowingStrategy) helpers.activeFnInfo.getWindowingStrategy(), windowsProcessed); } } + /** Returns the {@link DoFnInfo} used by this {@link StreamingKeyedWorkItemSideInputParDoFn}. */ + @Override + public DoFnInfo getDoFnInfo() throws Exception { + return helpers.doFnInstanceManager.peek(); + } + /** - * Returns the {@link DoFnInfo} currently being used by this {@link SimpleParDoFn}. + * Returns the {@link DoFnInfo} currently being used by this {@link + * StreamingKeyedWorkItemSideInputParDoFn}. * - *

May be null if no element has been processed yet, or if the {@link SimpleParDoFn} has - * finished. + *

May be null if no element has been processed yet, or if the {@link + * StreamingKeyedWorkItemSideInputParDoFn} has finished. */ @VisibleForTesting - @Nullable DoFnInfo getDoFnInfo() { - return helpers.fnInfo; + @Nullable DoFnInfo getActiveDoFnInfo() { + return helpers.activeFnInfo; } } 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 c7036aed0bae..b400ad61e44e 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 @@ -48,6 +48,7 @@ import org.apache.beam.runners.core.metrics.ExecutionStateTracker; import org.apache.beam.runners.core.metrics.ExecutionStateTracker.ExecutionState; import org.apache.beam.runners.dataflow.worker.DataflowOperationContext.DataflowExecutionState; +import org.apache.beam.runners.dataflow.worker.counters.Counter; import org.apache.beam.runners.dataflow.worker.counters.CounterFactory; import org.apache.beam.runners.dataflow.worker.counters.NameContext; import org.apache.beam.runners.dataflow.worker.profiler.ScopedProfiler.ProfileScope; @@ -67,7 +68,10 @@ import org.apache.beam.runners.dataflow.worker.streaming.sideinput.SideInputStateFetcherFactory; import org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor; import org.apache.beam.runners.dataflow.worker.util.common.worker.ElementCounter; +import org.apache.beam.runners.dataflow.worker.util.common.worker.Operation; import org.apache.beam.runners.dataflow.worker.util.common.worker.OutputObjectAndByteCounter; +import org.apache.beam.runners.dataflow.worker.util.common.worker.OutputReceiver; +import org.apache.beam.runners.dataflow.worker.util.common.worker.ParDoOperation; import org.apache.beam.runners.dataflow.worker.util.common.worker.WorkExecutor; import org.apache.beam.runners.dataflow.worker.windmill.Windmill; import org.apache.beam.runners.dataflow.worker.windmill.Windmill.GlobalDataId; @@ -796,19 +800,46 @@ private Windmill.WorkItemCommitRequest.Builder buildWorkItemTruncationRequestBui return outputBuilder; } - private final long computeSourceBytesProcessed(String sourceBytesCounterName) { + @VisibleForTesting + long computeSourceBytesProcessed(String sourceBytesCounterName) { if (!(workExecutor instanceof DataflowMapTaskExecutor)) { return 0L; } - HashMap counters = - ((DataflowMapTaskExecutor) workExecutor) - .getReadOperation() - .receivers[0] - .getOutputCounters(); + // Note this is a problem for both custom source read operators and SDF ParDos: extracting + // the byte count counter values here races with StreamingWorkerStatusReporter, which also + // calls getAndReset on all modified counters to report to Dataflow. We should find a more + // dedicated counter that is only extracted here. + DataflowMapTaskExecutor mapTaskExecutor = (DataflowMapTaskExecutor) workExecutor; + long totalBytes = + getAndResetOperationBytesProcessed( + mapTaskExecutor.getReadOperation(), sourceBytesCounterName); + for (ParDoOperation parDoOp : mapTaskExecutor.getParDoOperations()) { + if (parDoOp.shouldReportBytesRead()) { + totalBytes += getAndResetOperationBytesProcessed(parDoOp, sourceBytesCounterName); + } + } + return totalBytes; + } - return Optional.ofNullable(counters.get(sourceBytesCounterName)) - .map(counter -> ((OutputObjectAndByteCounter) counter).getByteCount().getAndReset()) - .orElse(0L); + private static long getAndResetOperationBytesProcessed( + @Nullable Operation operation, String sourceBytesCounterName) { + if (operation == null || operation.receivers == null) { + return 0L; + } + long bytes = 0L; + for (OutputReceiver receiver : operation.receivers) { + if (receiver == null) { + continue; + } + ElementCounter counter = receiver.getOutputCounters().get(sourceBytesCounterName); + if (counter instanceof OutputObjectAndByteCounter) { + Counter byteCount = ((OutputObjectAndByteCounter) counter).getByteCount(); + if (byteCount != null) { + bytes += byteCount.getAndReset(); + } + } + } + return bytes; } public boolean advance() throws CoderException { diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingPCollectionViewWriterParDoFn.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingPCollectionViewWriterParDoFn.java index 6b51427bb930..21ce2574bc81 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingPCollectionViewWriterParDoFn.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingPCollectionViewWriterParDoFn.java @@ -24,9 +24,11 @@ import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.coders.IterableCoder; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.util.DoFnInfo; import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.sdk.values.WindowedValue; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; +import org.checkerframework.checker.nullness.qual.Nullable; /** * A {@link ParDoFn} that writes side input data using {@link @@ -83,4 +85,9 @@ public void finishBundle() throws Exception {} public void abort() throws Exception { // no-op as long as it remains stateless } + + @Override + public @Nullable DoFnInfo getDoFnInfo() { + return null; + } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ToIsmRecordForMultimapDoFnFactory.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ToIsmRecordForMultimapDoFnFactory.java index 261ed69fb5d7..c275fb8e3d09 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ToIsmRecordForMultimapDoFnFactory.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ToIsmRecordForMultimapDoFnFactory.java @@ -35,10 +35,12 @@ import org.apache.beam.sdk.coders.KvCoder; import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.util.DoFnInfo; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.sdk.values.WindowedValue; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.checkerframework.checker.nullness.qual.Nullable; /** * A {@link ParDoFnFactory} that creates a system {@link ParDoFn} responsible for transforming @@ -157,5 +159,10 @@ public void finishBundle() {} @Override public void abort() {} + + @Override + public @Nullable DoFnInfo getDoFnInfo() { + return null; + } } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ValuesDoFnFactory.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ValuesDoFnFactory.java index acf143a5467b..ecf1d731a80d 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ValuesDoFnFactory.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ValuesDoFnFactory.java @@ -27,9 +27,11 @@ import org.apache.beam.runners.dataflow.worker.util.common.worker.Receiver; import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.transforms.Values; +import org.apache.beam.sdk.util.DoFnInfo; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.sdk.values.WindowedValue; +import org.checkerframework.checker.nullness.qual.Nullable; /** * A {@link ParDoFnFactory} which returns a {@link ParDoFn} with similar behavior to {@link @@ -88,5 +90,10 @@ public void finishBundle() {} @Override public void abort() {} + + @Override + public @Nullable DoFnInfo getDoFnInfo() { + return null; + } } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/MapTaskExecutor.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/MapTaskExecutor.java index 90fc76276940..919b0acf4de3 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/MapTaskExecutor.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/MapTaskExecutor.java @@ -145,6 +145,18 @@ public ReadOperation getReadOperation() { return (ReadOperation) readOperation; } + public List getParDoOperations() { + List parDoOperations = Lists.newArrayList(); + if (operations != null) { + for (Operation op : operations) { + if (op instanceof ParDoOperation) { + parDoOperations.add((ParDoOperation) op); + } + } + } + return parDoOperations; + } + public boolean supportsRestart() { for (Operation op : operations) { if (!op.supportsRestart()) { diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoFn.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoFn.java index 75ecdb67dd42..d06de2be088e 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoFn.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoFn.java @@ -17,6 +17,7 @@ */ package org.apache.beam.runners.dataflow.worker.util.common.worker; +import org.apache.beam.sdk.util.DoFnInfo; import org.checkerframework.checker.nullness.qual.Nullable; /** @@ -37,4 +38,10 @@ public interface ParDoFn { void finishBundle() throws Exception; void abort() throws Exception; + + /** + * For ParDoFn's that wrap an underlying DoFn, returns a potentially cached copy of its DoFnInfo. + * Otherwise, returns null. + */ + @Nullable DoFnInfo getDoFnInfo() throws Exception; } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoOperation.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoOperation.java index a814b90cae21..5376beb1829a 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoOperation.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoOperation.java @@ -18,16 +18,32 @@ package org.apache.beam.runners.dataflow.worker.util.common.worker; import java.io.Closeable; +import java.util.Map; +import org.apache.beam.sdk.coders.Coder; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; import org.checkerframework.checker.nullness.qual.Nullable; /** A ParDo mapping function. */ public class ParDoOperation extends ReceivingOperation { private final ParDoFn fn; + // If non-empty, this ParDo should report the amount of bytes read using these per-output-receiver + // coders to serialize elements. + private final Map> bytesReadCoders; public ParDoOperation(ParDoFn fn, OutputReceiver[] outputReceivers, OperationContext context) { + this(fn, outputReceivers, context, ImmutableMap.of()); + } + + public ParDoOperation( + ParDoFn fn, + OutputReceiver[] outputReceivers, + OperationContext context, + @Nullable Map> bytesReadCoders) { super(outputReceivers, context); this.fn = fn; + this.bytesReadCoders = + bytesReadCoders != null ? ImmutableMap.copyOf(bytesReadCoders) : ImmutableMap.of(); } @Override @@ -82,4 +98,16 @@ public boolean supportsRestart() { public ParDoFn getFn() throws Exception { return fn; } + + public boolean shouldReportBytesRead() { + return !bytesReadCoders.isEmpty(); + } + + public Map> getBytesReadCoders() { + return bytesReadCoders; + } + + public @Nullable Coder getBytesReadCoder(OutputReceiver receiver) { + return bytesReadCoders.get(receiver); + } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/SimplePartialGroupByKeyParDoFn.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/SimplePartialGroupByKeyParDoFn.java index cce77f3aceb6..eaf1f835a662 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/SimplePartialGroupByKeyParDoFn.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/SimplePartialGroupByKeyParDoFn.java @@ -17,6 +17,7 @@ */ package org.apache.beam.runners.dataflow.worker.util.common.worker; +import org.apache.beam.sdk.util.DoFnInfo; import org.checkerframework.checker.nullness.qual.Nullable; /** A partial group-by-key {@link ParDoFn} implementation. */ @@ -51,4 +52,9 @@ public void finishBundle() throws Exception { @Override public void abort() {} + + @Override + public @Nullable DoFnInfo getDoFnInfo() { + return null; + } } 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 b51512252e37..81ece67b74c8 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 @@ -54,6 +54,8 @@ 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.OutputReceiver; +import org.apache.beam.runners.dataflow.worker.util.common.worker.ParDoOperation; 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; @@ -166,7 +168,7 @@ private static void trackAutoscalingBytesRead( MapTask mapTask, Nodes.ParallelInstructionNode readNode, Coder readCoder, - ReadOperation readOperation, + OutputReceiver receiver, MapTaskExecutor mapTaskExecutor, String counterName) { NameContext nameContext = @@ -175,7 +177,7 @@ private static void trackAutoscalingBytesRead( readNode.getParallelInstruction().getOriginalName(), readNode.getParallelInstruction().getSystemName(), readNode.getParallelInstruction().getName()); - readOperation.receivers[0].addOutputCounter( + receiver.addOutputCounter( counterName, new OutputObjectAndByteCounter( new IntrinsicMapTaskExecutorFactory.ElementByteSizeObservableCoder<>(readCoder), @@ -227,13 +229,38 @@ ComputationWorkExecutor createComputationWorkExecutor( // If using a custom source, count bytes read for autoscaling. if (isCustomSource(readNode)) { - trackAutoscalingBytesRead( - mapTask, - readNode, - readCoder, - readOperation, - mapTaskExecutor, - computationState.sourceBytesProcessCounterName()); + if (readOperation.receivers != null) { + for (OutputReceiver receiver : readOperation.receivers) { + if (receiver != null) { + trackAutoscalingBytesRead( + mapTask, + readNode, + readCoder, + receiver, + mapTaskExecutor, + computationState.sourceBytesProcessCounterName()); + } + } + } + } else { + for (ParDoOperation parDoOp : mapTaskExecutor.getParDoOperations()) { + if (parDoOp.shouldReportBytesRead() && parDoOp.receivers != null) { + for (OutputReceiver receiver : parDoOp.receivers) { + if (receiver != null) { + Coder bytesReadCoder = parDoOp.getBytesReadCoder(receiver); + if (bytesReadCoder != null) { + trackAutoscalingBytesRead( + mapTask, + readNode, + bytesReadCoder, + receiver, + mapTaskExecutor, + computationState.sourceBytesProcessCounterName()); + } + } + } + } + } } ComputationWorkExecutor.Builder executionStateBuilder = diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/IntrinsicMapTaskExecutorFactoryTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/IntrinsicMapTaskExecutorFactoryTest.java index 85973372e419..83e2495a0839 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/IntrinsicMapTaskExecutorFactoryTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/IntrinsicMapTaskExecutorFactoryTest.java @@ -17,6 +17,7 @@ */ package org.apache.beam.runners.dataflow.worker; +import static org.apache.beam.runners.dataflow.util.Structs.addObject; import static org.apache.beam.runners.dataflow.util.Structs.addString; import static org.apache.beam.runners.dataflow.worker.DataflowOutputCounter.getElementCounterName; import static org.apache.beam.runners.dataflow.worker.DataflowOutputCounter.getMeanByteCounterName; @@ -75,6 +76,7 @@ import org.apache.beam.runners.dataflow.worker.graph.Nodes.Node; import org.apache.beam.runners.dataflow.worker.graph.Nodes.OperationNode; import org.apache.beam.runners.dataflow.worker.graph.Nodes.ParallelInstructionNode; +import org.apache.beam.runners.dataflow.worker.util.WorkerPropertyNames; import org.apache.beam.runners.dataflow.worker.util.common.worker.FlattenOperation; import org.apache.beam.runners.dataflow.worker.util.common.worker.Operation; import org.apache.beam.runners.dataflow.worker.util.common.worker.ParDoOperation; @@ -87,14 +89,17 @@ import org.apache.beam.sdk.coders.IterableCoder; import org.apache.beam.sdk.coders.KvCoder; import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.coders.VoidCoder; 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.io.range.OffsetRange; import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.options.PipelineOptionsFactory; import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.DoFnSchemaInformation; import org.apache.beam.sdk.transforms.Sum; +import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker; import org.apache.beam.sdk.transforms.windowing.IntervalWindow.IntervalWindowCoder; import org.apache.beam.sdk.util.AppliedCombineFn; import org.apache.beam.sdk.util.DoFnInfo; @@ -105,6 +110,7 @@ import org.apache.beam.sdk.values.WindowedValues.FullWindowedValueCoder; import org.apache.beam.sdk.values.WindowingStrategy; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.graph.MutableNetwork; @@ -567,6 +573,193 @@ public void testCreateParDoOperation() throws Exception { assertEquals(1, parDoOperation.receivers.length); assertEquals(0, parDoOperation.receivers[0].getReceiverCount()); assertEquals(Operation.InitializationState.UNSTARTED, parDoOperation.initializationState); + assertFalse(parDoOperation.shouldReportBytesRead()); + assertTrue(parDoOperation.getBytesReadCoders().isEmpty()); + assertNull(parDoOperation.getBytesReadCoder(parDoOperation.receivers[0])); + } + + static class TestSplittableDoFn extends DoFn { + @ProcessElement + public void processElement(ProcessContext c, RestrictionTracker tracker) { + if (tracker.tryClaim(tracker.currentRestriction().getFrom())) { + c.output(c.element()); + } + } + + @GetInitialRestriction + public OffsetRange getInitialRestriction(@SuppressWarnings("unused") @Element String element) { + return new OffsetRange(0, 1); + } + } + + @Test + public void testCreateSplittableProcessFnParDoOperation() throws Exception { + int producerIndex = 1; + int producerOutputNum = 2; + + BatchModeExecutionContext context = + BatchModeExecutionContext.forTesting(options, counterSet, "testStage"); + TupleTag mainOutput = new TupleTag<>(PropertyNames.OUTPUT); + String serializedFn = + StringUtils.byteArrayToJsonString( + SerializableUtils.serializeToByteArray( + DoFnInfo.forFn( + new TestSplittableDoFn(), + WindowingStrategy.globalDefault(), + null /* side input views */, + StringUtf8Coder.of() /* input coder */, + Collections.singletonMap(mainOutput, StringUtf8Coder.of()), + mainOutput /* main output id */, + DoFnSchemaInformation.create(), + Collections.emptyMap()))); + + CloudObject cloudUserFn = CloudObject.forClassName("SplittableProcessFn"); + addString(cloudUserFn, PropertyNames.SERIALIZED_FN, serializedFn); + addObject( + cloudUserFn, + WorkerPropertyNames.RESTRICTION_CODER, + CloudObjects.asCloudObject( + KvCoder.of(OffsetRange.Coder.of(), VoidCoder.of()), /* sdkComponents= */ null)); + + ParallelInstruction instruction = + createParDoInstruction( + producerIndex, producerOutputNum, "SplittableParDo", "", new TestSplittableDoFn()); + instruction.getParDo().setUserFn(cloudUserFn); + + ParallelInstructionNode instructionNode = + ParallelInstructionNode.create(instruction, ExecutionLocation.UNKNOWN); + + Node outputReceiverNode = + IntrinsicMapTaskExecutorFactory.createOutputReceiversTransform(STAGE, counterSet, false) + .apply( + InstructionOutputNode.create( + instructionNode.getParallelInstruction().getOutputs().get(0), PCOLLECTION_ID)); + + when(network.successors(instructionNode)).thenReturn(ImmutableSet.of(outputReceiverNode)); + when(network.outDegree(instructionNode)).thenReturn(1); + when(network.edgesConnecting(instructionNode, outputReceiverNode)) + .thenReturn( + ImmutableSet.of( + MultiOutputInfoEdge.create( + instructionNode + .getParallelInstruction() + .getParDo() + .getMultiOutputInfos() + .get(0)))); + + ArrayList createdOperations = new ArrayList<>(); + Node operationNode = + mapTaskExecutorFactory + .createOperationTransformForParallelInstructionNodes( + STAGE, network, options, readerRegistry, sinkRegistry, context, createdOperations) + .apply(instructionNode); + assertThat(operationNode, instanceOf(OperationNode.class)); + assertThat(((OperationNode) operationNode).getOperation(), instanceOf(ParDoOperation.class)); + ParDoOperation parDoOperation = (ParDoOperation) ((OperationNode) operationNode).getOperation(); + assertThat(createdOperations, contains(parDoOperation)); + + assertTrue(parDoOperation.shouldReportBytesRead()); + assertEquals( + ImmutableMap.of( + parDoOperation.receivers[0], WindowedValues.getValueOnlyCoder(StringUtf8Coder.of())), + parDoOperation.getBytesReadCoders()); + assertEquals( + WindowedValues.getValueOnlyCoder(StringUtf8Coder.of()), + parDoOperation.getBytesReadCoder(parDoOperation.receivers[0])); + } + + @Test + public void testCreateMultiOutputSplittableProcessFnParDoOperation() throws Exception { + int producerIndex = 1; + int producerOutputNum = 2; + + BatchModeExecutionContext context = + BatchModeExecutionContext.forTesting(options, counterSet, "testStage"); + TupleTag mainOutput = new TupleTag<>("main"); + TupleTag sideOutput = new TupleTag<>("side"); + String serializedFn = + StringUtils.byteArrayToJsonString( + SerializableUtils.serializeToByteArray( + DoFnInfo.forFn( + new TestSplittableDoFn(), + WindowingStrategy.globalDefault(), + null /* side input views */, + StringUtf8Coder.of() /* input coder */, + ImmutableMap.of( + mainOutput, StringUtf8Coder.of(), sideOutput, BigEndianIntegerCoder.of()), + mainOutput /* main output id */, + DoFnSchemaInformation.create(), + Collections.emptyMap()))); + + CloudObject cloudUserFn = CloudObject.forClassName("SplittableProcessFn"); + addString(cloudUserFn, PropertyNames.SERIALIZED_FN, serializedFn); + addObject( + cloudUserFn, + WorkerPropertyNames.RESTRICTION_CODER, + CloudObjects.asCloudObject( + KvCoder.of(OffsetRange.Coder.of(), VoidCoder.of()), /* sdkComponents= */ null)); + + MultiOutputInfo mainOutputInfo = new MultiOutputInfo(); + mainOutputInfo.setTag("main"); + MultiOutputInfo sideOutputInfo = new MultiOutputInfo(); + sideOutputInfo.setTag("side"); + + InstructionOutput mainCloudOutput = new InstructionOutput(); + mainCloudOutput.setName("SplittableParDo_main"); + mainCloudOutput.setCodec(windowedStringCoder); + mainCloudOutput.setOriginalName("originalName"); + mainCloudOutput.setSystemName("systemName"); + + Coder windowedIntCoder = WindowedValues.getValueOnlyCoder(BigEndianIntegerCoder.of()); + InstructionOutput sideCloudOutput = new InstructionOutput(); + sideCloudOutput.setName("SplittableParDo_side"); + sideCloudOutput.setCodec( + CloudObjects.asCloudObject(windowedIntCoder, /* sdkComponents= */ null)); + sideCloudOutput.setOriginalName("originalName"); + sideCloudOutput.setSystemName("systemName"); + + ParallelInstruction instruction = + createParDoInstruction( + producerIndex, producerOutputNum, "SplittableParDo", "", new TestSplittableDoFn()); + instruction.getParDo().setUserFn(cloudUserFn); + instruction.getParDo().setNumOutputs(2); + instruction.getParDo().setMultiOutputInfos(ImmutableList.of(mainOutputInfo, sideOutputInfo)); + instruction.setOutputs(ImmutableList.of(mainCloudOutput, sideCloudOutput)); + + ParallelInstructionNode instructionNode = + ParallelInstructionNode.create(instruction, ExecutionLocation.UNKNOWN); + + Function outputReceiversTransform = + IntrinsicMapTaskExecutorFactory.createOutputReceiversTransform(STAGE, counterSet, false); + Node mainOutputReceiverNode = + outputReceiversTransform.apply( + InstructionOutputNode.create(mainCloudOutput, PCOLLECTION_ID + "_main")); + Node sideOutputReceiverNode = + outputReceiversTransform.apply( + InstructionOutputNode.create(sideCloudOutput, PCOLLECTION_ID + "_side")); + + when(network.successors(instructionNode)) + .thenReturn(ImmutableSet.of(mainOutputReceiverNode, sideOutputReceiverNode)); + when(network.outDegree(instructionNode)).thenReturn(2); + when(network.edgesConnecting(instructionNode, mainOutputReceiverNode)) + .thenReturn(ImmutableSet.of(MultiOutputInfoEdge.create(mainOutputInfo))); + when(network.edgesConnecting(instructionNode, sideOutputReceiverNode)) + .thenReturn(ImmutableSet.of(MultiOutputInfoEdge.create(sideOutputInfo))); + + ArrayList createdOperations = new ArrayList<>(); + Node operationNode = + mapTaskExecutorFactory + .createOperationTransformForParallelInstructionNodes( + STAGE, network, options, readerRegistry, sinkRegistry, context, createdOperations) + .apply(instructionNode); + ParDoOperation parDoOperation = (ParDoOperation) ((OperationNode) operationNode).getOperation(); + + assertEquals(2, parDoOperation.receivers.length); + assertTrue(parDoOperation.shouldReportBytesRead()); + assertEquals( + WindowedValues.getValueOnlyCoder(StringUtf8Coder.of()), + parDoOperation.getBytesReadCoder(parDoOperation.receivers[0])); + assertEquals(windowedIntCoder, parDoOperation.getBytesReadCoder(parDoOperation.receivers[1])); } static ParallelInstruction createPartialGroupByKeyInstruction( diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/IntrinsicMapTaskExecutorTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/IntrinsicMapTaskExecutorTest.java index 73d69bd0f617..4d0d24a5886f 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/IntrinsicMapTaskExecutorTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/IntrinsicMapTaskExecutorTest.java @@ -69,7 +69,9 @@ import org.apache.beam.sdk.metrics.Metrics; import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.util.DoFnInfo; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; +import org.checkerframework.checker.nullness.qual.Nullable; import org.junit.Assert; import org.junit.Rule; import org.junit.Test; @@ -223,6 +225,11 @@ public void abort() {} @Override public void finishKey(Object key) throws Exception {} + + @Override + public @Nullable DoFnInfo getDoFnInfo() { + return null; + } } /** Verify counts for the per-element-output-time counter are correct. */ 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 0251b075becf..bc808865cc68 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 @@ -43,6 +43,7 @@ import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; +import org.apache.beam.runners.core.ElementByteSizeObservable; import org.apache.beam.runners.core.SideInputReader; import org.apache.beam.runners.core.StateInternals; import org.apache.beam.runners.core.StateNamespaceForTest; @@ -72,6 +73,12 @@ 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.BoundedQueueExecutor; +import org.apache.beam.runners.dataflow.worker.util.common.worker.NativeReader; +import org.apache.beam.runners.dataflow.worker.util.common.worker.OutputObjectAndByteCounter; +import org.apache.beam.runners.dataflow.worker.util.common.worker.OutputReceiver; +import org.apache.beam.runners.dataflow.worker.util.common.worker.ParDoFn; +import org.apache.beam.runners.dataflow.worker.util.common.worker.ParDoOperation; +import org.apache.beam.runners.dataflow.worker.util.common.worker.ReadOperation; import org.apache.beam.runners.dataflow.worker.util.common.worker.WorkExecutor; import org.apache.beam.runners.dataflow.worker.windmill.Windmill; import org.apache.beam.runners.dataflow.worker.windmill.client.getdata.FakeGetDataClient; @@ -95,6 +102,7 @@ import org.apache.beam.sdk.values.PCollectionView; import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; import org.hamcrest.Matchers; import org.joda.time.Duration; @@ -943,4 +951,237 @@ public void testAdvance_stopsWhenCurrentWorkBatchingDisabled() throws Exception assertFalse(transitionListenerCalled.get()); verifyNoInteractions(mockExecutor); } + + @Test + public void testComputeSourceBytesProcessed_customSource() throws Exception { + CounterSet counterSet = new CounterSet(); + NameContext nameContext = NameContext.create("stage", "original", "system", "out"); + DataflowOperationContext opContext = executionContext.createOperationContext(nameContext); + + OutputReceiver readReceiver = new OutputReceiver(); + OutputObjectAndByteCounter customSourceCounter = + new OutputObjectAndByteCounter( + mock(ElementByteSizeObservable.class), counterSet, nameContext) + .countBytes("sourceBytesProcessCounterName"); + customSourceCounter.getByteCount().addValue(100L); + readReceiver.addOutputCounter("sourceBytesProcessCounterName", customSourceCounter); + + ReadOperation readOp = ReadOperation.forTest(mock(NativeReader.class), readReceiver, opContext); + + workExecutor = + IntrinsicMapTaskExecutor.forOperations( + Arrays.asList(readOp), mock(ExecutionStateTracker.class)); + + Work work = + createMockWork( + Windmill.WorkItem.newBuilder() + .setKey(ByteString.copyFromUtf8("key1")) + .setWorkToken(1L) + .build(), + Watermarks.builder().setInputDataWatermark(Instant.EPOCH).build()); + + start(work); + + long processed = executionContext.computeSourceBytesProcessed("sourceBytesProcessCounterName"); + assertEquals(100L, processed); + // Verified getAndReset resets to 0. + assertEquals(0L, executionContext.computeSourceBytesProcessed("sourceBytesProcessCounterName")); + } + + @Test + public void testComputeSourceBytesProcessed_splittableDoFn() throws Exception { + CounterSet counterSet = new CounterSet(); + NameContext nameContext = NameContext.create("stage", "original", "system", "out"); + DataflowOperationContext opContext = executionContext.createOperationContext(nameContext); + + OutputReceiver readReceiver = new OutputReceiver(); + ReadOperation readOp = ReadOperation.forTest(mock(NativeReader.class), readReceiver, opContext); + + ParDoFn mockParDoFn = mock(ParDoFn.class); + OutputReceiver parDoReceiver = new OutputReceiver(); + OutputObjectAndByteCounter sdfCounter = + new OutputObjectAndByteCounter( + mock(ElementByteSizeObservable.class), counterSet, nameContext) + .countBytes("sourceBytesProcessCounterName"); + sdfCounter.getByteCount().addValue(42L); + parDoReceiver.addOutputCounter("sourceBytesProcessCounterName", sdfCounter); + + ParDoOperation parDoOp = + new ParDoOperation( + mockParDoFn, + new OutputReceiver[] {parDoReceiver}, + opContext, + ImmutableMap.of(parDoReceiver, mock(Coder.class))); + + workExecutor = + IntrinsicMapTaskExecutor.forOperations( + Arrays.asList(readOp, parDoOp), mock(ExecutionStateTracker.class)); + + Work work = + createMockWork( + Windmill.WorkItem.newBuilder() + .setKey(ByteString.copyFromUtf8("key1")) + .setWorkToken(1L) + .build(), + Watermarks.builder().setInputDataWatermark(Instant.EPOCH).build()); + + start(work); + + long processed = executionContext.computeSourceBytesProcessed("sourceBytesProcessCounterName"); + assertEquals(42L, processed); + // Verified getAndReset resets to 0. + assertEquals(0L, executionContext.computeSourceBytesProcessed("sourceBytesProcessCounterName")); + } + + @Test + public void testComputeSourceBytesProcessed_multipleOutputReceivers() throws Exception { + NameContext nameContext = NameContext.create("stage", "original", "system", "out"); + DataflowOperationContext opContext = executionContext.createOperationContext(nameContext); + + // Custom source ReadOperation with a null receiver[0] and an active receiver[1]. + OutputReceiver readSideReceiver = new OutputReceiver(); + OutputObjectAndByteCounter readSideCounter = + new OutputObjectAndByteCounter( + mock(ElementByteSizeObservable.class), new CounterSet(), nameContext) + .countBytes("sourceBytesProcessCounterName"); + readSideCounter.getByteCount().addValue(15L); + readSideReceiver.addOutputCounter("sourceBytesProcessCounterName", readSideCounter); + ReadOperation readOp = + ReadOperation.create( + mock(NativeReader.class), new OutputReceiver[] {null, readSideReceiver}, opContext); + + // Multi-output SDF ParDoOperation with two receivers sharing a CounterSet (like production) + // plus an additional receiver with its own CounterSet. + ParDoFn mockParDoFn = mock(ParDoFn.class); + CounterSet sharedCounterSet = new CounterSet(); + + OutputReceiver sdfMainReceiver = new OutputReceiver(); + OutputObjectAndByteCounter sdfMainCounter = + new OutputObjectAndByteCounter( + mock(ElementByteSizeObservable.class), sharedCounterSet, nameContext) + .countBytes("sourceBytesProcessCounterName"); + sdfMainCounter.getByteCount().addValue(25L); + sdfMainReceiver.addOutputCounter("sourceBytesProcessCounterName", sdfMainCounter); + + OutputReceiver sdfSideReceiver1 = new OutputReceiver(); + OutputObjectAndByteCounter sdfSideCounter1 = + new OutputObjectAndByteCounter( + mock(ElementByteSizeObservable.class), sharedCounterSet, nameContext) + .countBytes("sourceBytesProcessCounterName"); + sdfSideCounter1.getByteCount().addValue(30L); + sdfSideReceiver1.addOutputCounter("sourceBytesProcessCounterName", sdfSideCounter1); + + OutputReceiver sdfSideReceiver2 = new OutputReceiver(); + OutputObjectAndByteCounter sdfSideCounter2 = + new OutputObjectAndByteCounter( + mock(ElementByteSizeObservable.class), new CounterSet(), nameContext) + .countBytes("sourceBytesProcessCounterName"); + sdfSideCounter2.getByteCount().addValue(35L); + sdfSideReceiver2.addOutputCounter("sourceBytesProcessCounterName", sdfSideCounter2); + + ParDoOperation sdfParDoOp = + new ParDoOperation( + mockParDoFn, + new OutputReceiver[] {sdfMainReceiver, null, sdfSideReceiver1, sdfSideReceiver2}, + opContext, + ImmutableMap.of( + sdfMainReceiver, + mock(Coder.class), + sdfSideReceiver1, + mock(Coder.class), + sdfSideReceiver2, + mock(Coder.class))); + + workExecutor = + IntrinsicMapTaskExecutor.forOperations( + Arrays.asList(readOp, sdfParDoOp), mock(ExecutionStateTracker.class)); + + Work work = + createMockWork( + Windmill.WorkItem.newBuilder() + .setKey(ByteString.copyFromUtf8("key1")) + .setWorkToken(1L) + .build(), + Watermarks.builder().setInputDataWatermark(Instant.EPOCH).build()); + + start(work); + + long processed = executionContext.computeSourceBytesProcessed("sourceBytesProcessCounterName"); + assertEquals(15L + 25L + 30L + 35L, processed); + assertEquals(0L, executionContext.computeSourceBytesProcessed("sourceBytesProcessCounterName")); + } + + @Test + public void testComputeSourceBytesProcessed_multipleParDoOperationsAndFlushState() + throws Exception { + NameContext nameContext = NameContext.create("stage", "original", "system", "out"); + DataflowOperationContext opContext = executionContext.createOperationContext(nameContext); + + OutputReceiver readReceiver = new OutputReceiver(); + ReadOperation readOp = ReadOperation.forTest(mock(NativeReader.class), readReceiver, opContext); + + ParDoFn mockParDoFn = mock(ParDoFn.class); + + OutputReceiver sdfReceiver1 = new OutputReceiver(); + OutputObjectAndByteCounter sdfCounter1 = + new OutputObjectAndByteCounter( + mock(ElementByteSizeObservable.class), new CounterSet(), nameContext) + .countBytes("sourceBytesProcessCounterName"); + sdfCounter1.getByteCount().addValue(40L); + sdfReceiver1.addOutputCounter("sourceBytesProcessCounterName", sdfCounter1); + ParDoOperation sdfParDoOp1 = + new ParDoOperation( + mockParDoFn, + new OutputReceiver[] {sdfReceiver1}, + opContext, + ImmutableMap.of(sdfReceiver1, mock(Coder.class))); + + OutputReceiver sdfReceiver2 = new OutputReceiver(); + OutputObjectAndByteCounter sdfCounter2 = + new OutputObjectAndByteCounter( + mock(ElementByteSizeObservable.class), new CounterSet(), nameContext) + .countBytes("sourceBytesProcessCounterName"); + sdfCounter2.getByteCount().addValue(60L); + sdfReceiver2.addOutputCounter("sourceBytesProcessCounterName", sdfCounter2); + ParDoOperation sdfParDoOp2 = + new ParDoOperation( + mockParDoFn, + new OutputReceiver[] {sdfReceiver2}, + opContext, + ImmutableMap.of(sdfReceiver2, mock(Coder.class))); + + OutputReceiver regularReceiver = new OutputReceiver(); + OutputObjectAndByteCounter regularCounter = + new OutputObjectAndByteCounter( + mock(ElementByteSizeObservable.class), new CounterSet(), nameContext) + .countBytes("sourceBytesProcessCounterName"); + regularCounter.getByteCount().addValue(999L); + regularReceiver.addOutputCounter("sourceBytesProcessCounterName", regularCounter); + ParDoOperation regularParDoOp = + new ParDoOperation(mockParDoFn, new OutputReceiver[] {regularReceiver}, opContext, null); + + workExecutor = + IntrinsicMapTaskExecutor.forOperations( + Arrays.asList(readOp, sdfParDoOp1, regularParDoOp, sdfParDoOp2), + mock(ExecutionStateTracker.class)); + + Work work = + createMockWork( + Windmill.WorkItem.newBuilder() + .setKey(ByteString.copyFromUtf8("key1")) + .setWorkToken(1L) + .build(), + Watermarks.builder().setInputDataWatermark(Instant.EPOCH).build()); + + sdfParDoOp1.start(); + sdfParDoOp2.start(); + regularParDoOp.start(); + + start(work); + executionContext.finishKey(); + ExecuteWorkResult result = executionContext.flushStateAndReset(); + + assertEquals(1, result.workItemCommits().size()); + assertEquals(100L, result.workItemCommits().get(0).getSourceBytesProcessed()); + } } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/UserParDoFnFactoryTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/UserParDoFnFactoryTest.java index 43331d11a7ee..0168c8aa0ada 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/UserParDoFnFactoryTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/UserParDoFnFactoryTest.java @@ -196,7 +196,7 @@ public void testFactoryReuseInStep() throws Exception { parDoFn.startBundle(rcvr); parDoFn.processElement(WindowedValues.valueInGlobalWindow("foo")); - TestDoFn fn = (TestDoFn) ((SimpleParDoFn) parDoFn).getDoFnInfo().getDoFn(); + TestDoFn fn = (TestDoFn) ((SimpleParDoFn) parDoFn).getActiveDoFnInfo().getDoFn(); assertThat(fn, not(theInstance(initialFn))); parDoFn.finishBundle(); @@ -217,7 +217,8 @@ public void testFactoryReuseInStep() throws Exception { secondParDoFn.startBundle(rcvr); secondParDoFn.processElement(WindowedValues.valueInGlobalWindow("spam")); - TestDoFn reobtainedFn = (TestDoFn) ((SimpleParDoFn) secondParDoFn).getDoFnInfo().getDoFn(); + TestDoFn reobtainedFn = + (TestDoFn) ((SimpleParDoFn) secondParDoFn).getActiveDoFnInfo().getDoFn(); secondParDoFn.finishBundle(); assertThat(reobtainedFn.state, equalTo(TestDoFn.State.FINISHED)); @@ -256,13 +257,13 @@ public void testFactorySimultaneousUse() throws Exception { parDoFn.processElement(WindowedValues.valueInGlobalWindow("foo")); // Must be after the first call to process element for reallyStartBundle to have been called - TestDoFn firstDoFn = (TestDoFn) ((SimpleParDoFn) parDoFn).getDoFnInfo().getDoFn(); + TestDoFn firstDoFn = (TestDoFn) ((SimpleParDoFn) parDoFn).getActiveDoFnInfo().getDoFn(); secondParDoFn.startBundle(rcvr); secondParDoFn.processElement(WindowedValues.valueInGlobalWindow("spam")); // Must be after the first call to process element for reallyStartBundle to have been called - TestDoFn secondDoFn = (TestDoFn) ((SimpleParDoFn) secondParDoFn).getDoFnInfo().getDoFn(); + TestDoFn secondDoFn = (TestDoFn) ((SimpleParDoFn) secondParDoFn).getActiveDoFnInfo().getDoFn(); parDoFn.finishBundle(); secondParDoFn.finishBundle(); @@ -292,7 +293,7 @@ public void testFactoryDoesNotReuseAfterAborted() throws Exception { parDoFn.startBundle(rcvr); parDoFn.processElement(WindowedValues.valueInGlobalWindow("foo")); - TestDoFn fn = (TestDoFn) ((SimpleParDoFn) parDoFn).getDoFnInfo().getDoFn(); + TestDoFn fn = (TestDoFn) ((SimpleParDoFn) parDoFn).getActiveDoFnInfo().getDoFn(); parDoFn.abort(); assertThat(fn.state, equalTo(TestDoFn.State.TORN_DOWN)); @@ -310,7 +311,7 @@ public void testFactoryDoesNotReuseAfterAborted() throws Exception { secondParDoFn.startBundle(rcvr); secondParDoFn.processElement(WindowedValues.valueInGlobalWindow("foo")); - TestDoFn secondFn = (TestDoFn) ((SimpleParDoFn) secondParDoFn).getDoFnInfo().getDoFn(); + TestDoFn secondFn = (TestDoFn) ((SimpleParDoFn) secondParDoFn).getActiveDoFnInfo().getDoFn(); assertThat(secondFn, not(theInstance(fn))); assertThat(fn.state, equalTo(TestDoFn.State.TORN_DOWN)); diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/MapTaskExecutorTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/MapTaskExecutorTest.java index 40f4c7b0e715..48f0a7d1bfae 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/MapTaskExecutorTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/MapTaskExecutorTest.java @@ -32,6 +32,7 @@ import static org.hamcrest.Matchers.contains; import static org.hamcrest.Matchers.equalTo; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; import static org.junit.Assert.fail; import static org.mockito.ArgumentMatchers.anyBoolean; import static org.mockito.ArgumentMatchers.eq; @@ -62,10 +63,15 @@ import org.apache.beam.runners.dataflow.worker.counters.NameContext; import org.apache.beam.runners.dataflow.worker.profiler.ScopedProfiler.NoopProfileScope; import org.apache.beam.runners.dataflow.worker.util.common.worker.ExecutorTestUtils.TestReader; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.coders.VarIntCoder; import org.apache.beam.sdk.metrics.Metrics; import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.util.DoFnInfo; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; +import org.checkerframework.checker.nullness.qual.Nullable; import org.junit.Assert; import org.junit.Rule; import org.junit.Test; @@ -219,6 +225,11 @@ public void abort() {} @Override public void finishKey(Object key) throws Exception {} + + @Override + public @Nullable DoFnInfo getDoFnInfo() { + return null; + } } /** Verify counts for the per-element-output-time counter are correct. */ @@ -573,4 +584,40 @@ public void testExceptionAndThenCloseAbortsJustOnce() throws Exception { // Closing after already closed should not call abort again. executor.close(); } + + @Test + public void testGetParDoOperations() throws Exception { + TestOutputReceiver receiver = + new TestOutputReceiver(counterSet, NameContextsForTests.nameContextForTest()); + ReadOperation readOp = new TestReadOperation(receiver, createContext("ReadOperation")); + OutputReceiver reportingReceiver1 = new OutputReceiver(); + OutputReceiver reportingReceiver2 = new OutputReceiver(); + ParDoOperation regularParDo = + new ParDoOperation( + new NoopParDoFn(), new OutputReceiver[0], createContext("RegularParDo"), null); + ParDoOperation reportingParDo1 = + new ParDoOperation( + new NoopParDoFn(), + new OutputReceiver[] {reportingReceiver1}, + createContext("ReportingParDo1"), + ImmutableMap.of(reportingReceiver1, StringUtf8Coder.of())); + ParDoOperation reportingParDo2 = + new ParDoOperation( + new NoopParDoFn(), + new OutputReceiver[] {reportingReceiver2}, + createContext("ReportingParDo2"), + ImmutableMap.of(reportingReceiver2, VarIntCoder.of())); + + List operations = + Arrays.asList(readOp, regularParDo, reportingParDo1, reportingParDo2); + ExecutionStateTracker stateTracker = ExecutionStateTracker.newForTest(); + try (MapTaskExecutor executor = new MapTaskExecutor(operations, counterSet, stateTracker)) { + List parDoOps = executor.getParDoOperations(); + assertThat(parDoOps, contains(regularParDo, reportingParDo1, reportingParDo2)); + assertTrue(parDoOps.get(0).getBytesReadCoders().isEmpty()); + Assert.assertNull(parDoOps.get(0).getBytesReadCoder(reportingReceiver1)); + assertEquals(StringUtf8Coder.of(), parDoOps.get(1).getBytesReadCoder(reportingReceiver1)); + assertEquals(VarIntCoder.of(), parDoOps.get(2).getBytesReadCoder(reportingReceiver2)); + } + } } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoOperationTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoOperationTest.java index 0c9667acf38b..10203834d24f 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoOperationTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ParDoOperationTest.java @@ -32,6 +32,8 @@ import org.apache.beam.runners.dataflow.worker.counters.CounterFactory.LongCounterMean; import org.apache.beam.runners.dataflow.worker.counters.CounterSet; import org.apache.beam.runners.dataflow.worker.counters.NameContext; +import org.apache.beam.sdk.util.DoFnInfo; +import org.checkerframework.checker.nullness.qual.Nullable; import org.hamcrest.CoreMatchers; import org.junit.Test; import org.junit.runner.RunWith; @@ -88,6 +90,11 @@ public void abort() throws Exception { @Override public void finishKey(Object key) throws Exception {} + + @Override + public @Nullable DoFnInfo getDoFnInfo() { + return null; + } } @Test diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactoryTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactoryTest.java new file mode 100644 index 000000000000..4022ea2c62da --- /dev/null +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactoryTest.java @@ -0,0 +1,475 @@ +/* + * 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.windmill.work.processing; + +import static org.apache.beam.runners.dataflow.util.Structs.addObject; +import static org.apache.beam.runners.dataflow.util.Structs.addString; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.mock; + +import com.google.api.services.dataflow.model.InstructionInput; +import com.google.api.services.dataflow.model.InstructionOutput; +import com.google.api.services.dataflow.model.MapTask; +import com.google.api.services.dataflow.model.MultiOutputInfo; +import com.google.api.services.dataflow.model.ParDoInstruction; +import com.google.api.services.dataflow.model.ParallelInstruction; +import com.google.api.services.dataflow.model.ReadInstruction; +import com.google.api.services.dataflow.model.Source; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +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.util.PropertyNames; +import org.apache.beam.runners.dataflow.worker.DataflowExecutionStateSampler; +import org.apache.beam.runners.dataflow.worker.DataflowMapTaskExecutor; +import org.apache.beam.runners.dataflow.worker.HotKeyLogger; +import org.apache.beam.runners.dataflow.worker.IntrinsicMapTaskExecutorFactory; +import org.apache.beam.runners.dataflow.worker.MultiKeyBundleOptions; +import org.apache.beam.runners.dataflow.worker.ReaderCache; +import org.apache.beam.runners.dataflow.worker.counters.Counter; +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.StageInfo; +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.harness.StreamingCounters; +import org.apache.beam.runners.dataflow.worker.streaming.sideinput.SideInputStateFetcherFactory; +import org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor; +import org.apache.beam.runners.dataflow.worker.util.WorkerPropertyNames; +import org.apache.beam.runners.dataflow.worker.util.common.worker.ParDoOperation; +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.coders.StringUtf8Coder; +import org.apache.beam.sdk.coders.VarIntCoder; +import org.apache.beam.sdk.coders.VoidCoder; +import org.apache.beam.sdk.extensions.gcp.util.Transport; +import org.apache.beam.sdk.fn.IdGenerators; +import org.apache.beam.sdk.io.range.OffsetRange; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.DoFnSchemaInformation; +import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker; +import org.apache.beam.sdk.transforms.windowing.GlobalWindow; +import org.apache.beam.sdk.util.CoderUtils; +import org.apache.beam.sdk.util.DoFnInfo; +import org.apache.beam.sdk.util.SerializableUtils; +import org.apache.beam.sdk.util.StringUtils; +import org.apache.beam.sdk.values.TupleTag; +import org.apache.beam.sdk.values.WindowedValue; +import org.apache.beam.sdk.values.WindowedValues; +import org.apache.beam.sdk.values.WindowingStrategy; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Tests for {@link ComputationWorkExecutorFactory}. */ +@RunWith(JUnit4.class) +public class ComputationWorkExecutorFactoryTest { + + private static class StringSplittableDoFn extends DoFn { + @ProcessElement + public void processElement(ProcessContext c, RestrictionTracker tracker) { + if (tracker.tryClaim(tracker.currentRestriction().getFrom())) { + c.output(c.element()); + } + } + + @GetInitialRestriction + public OffsetRange getInitialRestriction(@SuppressWarnings("unused") @Element String element) { + return new OffsetRange(0, 1); + } + } + + private static class IntegerSplittableDoFn extends DoFn { + @ProcessElement + public void processElement(ProcessContext c, RestrictionTracker tracker) { + if (tracker.tryClaim(tracker.currentRestriction().getFrom())) { + c.output(c.element().length()); + } + } + + @GetInitialRestriction + public OffsetRange getInitialRestriction(@SuppressWarnings("unused") @Element String element) { + return new OffsetRange(0, 1); + } + } + + private static class IdentityDoFn extends DoFn { + @ProcessElement + public void processElement(ProcessContext c) { + c.output(c.element()); + } + } + + private static ParallelInstruction createWindmillReadInstruction( + String name, Coder outputCoder) { + CloudObject encodedCoder = CloudObjects.asCloudObject(outputCoder, /* sdkComponents= */ null); + + Source cloudSource = new Source(); + cloudSource.setSpec(CloudObject.forClassName("UngroupedWindmillReader")); + cloudSource.setCodec(encodedCoder); + + ReadInstruction readInstruction = new ReadInstruction(); + readInstruction.setSource(cloudSource); + + InstructionOutput output = new InstructionOutput(); + output.setName(name + "_output"); + output.setCodec(encodedCoder); + output.setOriginalName(name + "Original"); + output.setSystemName(name + "System"); + + ParallelInstruction instruction = new ParallelInstruction(); + instruction.setName(name); + instruction.setSystemName(name + "System"); + instruction.setOriginalName(name + "Original"); + instruction.setRead(readInstruction); + instruction.setOutputs(Collections.singletonList(output)); + return instruction; + } + + private static ParallelInstruction createParDoInstruction( + int producerIndex, + int producerOutputNum, + String systemName, + DoFn fn, + Coder elementOutputCoder, + Coder> windowedOutputCoder, + boolean isSplittable) { + InstructionInput cloudInput = new InstructionInput(); + cloudInput.setProducerInstructionIndex(producerIndex); + cloudInput.setOutputNum(producerOutputNum); + + TupleTag mainOutputTag = new TupleTag<>(PropertyNames.OUTPUT); + String serializedFn = + StringUtils.byteArrayToJsonString( + SerializableUtils.serializeToByteArray( + DoFnInfo.forFn( + fn, + WindowingStrategy.globalDefault(), + null /* side input views */, + StringUtf8Coder.of() /* input coder */, + Collections.singletonMap(mainOutputTag, elementOutputCoder), + mainOutputTag, + DoFnSchemaInformation.create(), + Collections.emptyMap()))); + + CloudObject cloudUserFn = + CloudObject.forClassName(isSplittable ? "SplittableProcessFn" : "DoFn"); + addString(cloudUserFn, PropertyNames.SERIALIZED_FN, serializedFn); + if (isSplittable) { + addObject( + cloudUserFn, + WorkerPropertyNames.RESTRICTION_CODER, + CloudObjects.asCloudObject( + KvCoder.of(OffsetRange.Coder.of(), VoidCoder.of()), /* sdkComponents= */ null)); + } + + MultiOutputInfo multiOutputInfo = new MultiOutputInfo(); + multiOutputInfo.setTag(PropertyNames.OUTPUT); + + ParDoInstruction parDoInstruction = new ParDoInstruction(); + parDoInstruction.setInput(cloudInput); + parDoInstruction.setNumOutputs(1); + parDoInstruction.setMultiOutputInfos(ImmutableList.of(multiOutputInfo)); + parDoInstruction.setUserFn(cloudUserFn); + + InstructionOutput output = new InstructionOutput(); + output.setName(systemName + "_output"); + output.setCodec(CloudObjects.asCloudObject(windowedOutputCoder, /* sdkComponents= */ null)); + output.setOriginalName(systemName + "Original"); + output.setSystemName(systemName); + + ParallelInstruction instruction = new ParallelInstruction(); + instruction.setParDo(parDoInstruction); + instruction.setOutputs(Collections.singletonList(output)); + instruction.setName(systemName); + instruction.setSystemName(systemName); + instruction.setOriginalName(systemName + "Original"); + return instruction; + } + + @Test + public void + testCreateComputationWorkExecutor_multipleReadReportingParDoOperationsWithDifferentCoders() + throws Exception { + DataflowWorkerHarnessOptions options = + PipelineOptionsFactory.as(DataflowWorkerHarnessOptions.class); + ComputationWorkExecutorFactory factory = + new ComputationWorkExecutorFactory( + options, + IntrinsicMapTaskExecutorFactory.defaultFactory(), + mock(ReaderCache.class), + computationId -> mock(WindmillStateCache.ForComputation.class), + DataflowExecutionStateSampler.instance(), + StreamingCounters.create(), + mock(FailureTracker.class), + IdGenerators.decrementingLongs(), + new FakeGlobalConfigHandle(StreamingGlobalConfig.builder().build()), + mock(HotKeyLogger.class), + SideInputStateFetcherFactory.fromOptions(options), + MultiKeyBundleOptions.fromOptions(options)); + + Coder> windowedStringCoder = + WindowedValues.getFullCoder(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE); + Coder> windowedIntCoder = + WindowedValues.getFullCoder(VarIntCoder.of(), GlobalWindow.Coder.INSTANCE); + + List instructions = + Arrays.asList( + createWindmillReadInstruction("Read", windowedStringCoder), + createParDoInstruction( + 0, + 0, + "SplittableParDoString", + new StringSplittableDoFn(), + StringUtf8Coder.of(), + windowedStringCoder, + true), + createParDoInstruction( + 0, + 0, + "SplittableParDoInt", + new IntegerSplittableDoFn(), + VarIntCoder.of(), + windowedIntCoder, + true), + createParDoInstruction( + 0, + 0, + "RegularParDo", + new IdentityDoFn(), + StringUtf8Coder.of(), + windowedStringCoder, + false)); + + MapTask mapTask = new MapTask(); + mapTask.setStageName("testStage"); + mapTask.setSystemName("testSystem"); + mapTask.setInstructions(instructions); + mapTask.setFactory(Transport.getJsonFactory()); + + StageInfo stageInfo = StageInfo.create("testStage", "testSystem"); + ComputationState computationState = + new ComputationState( + "computationId", + mapTask, + mock(BoundedQueueExecutor.class), + Collections.emptyMap(), + mock(WindmillStateCache.ForComputation.class)); + + ComputationWorkExecutor workExecutor = + factory.createComputationWorkExecutor(stageInfo, computationState, "workLatencyTrackingId"); + + DataflowMapTaskExecutor mapTaskExecutor = (DataflowMapTaskExecutor) workExecutor.workExecutor(); + List allParDoOps = mapTaskExecutor.getParDoOperations(); + assertEquals(3, allParDoOps.size()); + + ParDoOperation stringSplittableOp = + allParDoOps.stream() + .filter(op -> windowedStringCoder.equals(op.getBytesReadCoder(op.receivers[0]))) + .findFirst() + .orElseThrow( + () -> new AssertionError("Expected ParDoOperation with windowedStringCoder")); + ParDoOperation intSplittableOp = + allParDoOps.stream() + .filter(op -> windowedIntCoder.equals(op.getBytesReadCoder(op.receivers[0]))) + .findFirst() + .orElseThrow(() -> new AssertionError("Expected ParDoOperation with windowedIntCoder")); + + String counterName = computationState.sourceBytesProcessCounterName(); + assertNotNull(stringSplittableOp.receivers[0].getOutputCounters().get(counterName)); + assertNotNull(intSplittableOp.receivers[0].getOutputCounters().get(counterName)); + + // Verify that the non-reporting ParDoOperation does not have the source bytes counter attached. + ParDoOperation regularParDoOp = + allParDoOps.stream() + .filter(op -> !op.shouldReportBytesRead()) + .findFirst() + .orElseThrow(() -> new AssertionError("Expected non-reporting ParDoOperation")); + assertTrue(regularParDoOp.getBytesReadCoders().isEmpty()); + assertNull(regularParDoOp.getBytesReadCoder(regularParDoOp.receivers[0])); + assertNull(regularParDoOp.receivers[0].getOutputCounters().get(counterName)); + + // Process elements through each operation's receiver and verify each uses its own coder. + WindowedValue stringElem = WindowedValues.valueInGlobalWindow("hello"); + WindowedValue intElem = WindowedValues.valueInGlobalWindow(12345); + stringSplittableOp.receivers[0].process(stringElem); + intSplittableOp.receivers[0].process(intElem); + regularParDoOp.receivers[0].process(stringElem); + + long expectedStringBytes = + CoderUtils.encodeToByteArray(windowedStringCoder, stringElem, Coder.Context.NESTED).length; + long expectedIntBytes = + CoderUtils.encodeToByteArray(windowedIntCoder, intElem, Coder.Context.NESTED).length; + + @SuppressWarnings("unchecked") + Counter bytesReadCounter = + (Counter) mapTaskExecutor.getOutputCounters().getExistingCounter(counterName); + assertNotNull(bytesReadCounter); + assertEquals( + Long.valueOf(expectedStringBytes + expectedIntBytes), bytesReadCounter.getAggregate()); + } + + @Test + public void testCreateComputationWorkExecutor_multiOutputSplittableParDoWithDifferentCoders() + throws Exception { + DataflowWorkerHarnessOptions options = + PipelineOptionsFactory.as(DataflowWorkerHarnessOptions.class); + ComputationWorkExecutorFactory factory = + new ComputationWorkExecutorFactory( + options, + IntrinsicMapTaskExecutorFactory.defaultFactory(), + mock(ReaderCache.class), + computationId -> mock(WindmillStateCache.ForComputation.class), + DataflowExecutionStateSampler.instance(), + StreamingCounters.create(), + mock(FailureTracker.class), + IdGenerators.decrementingLongs(), + new FakeGlobalConfigHandle(StreamingGlobalConfig.builder().build()), + mock(HotKeyLogger.class), + SideInputStateFetcherFactory.fromOptions(options), + MultiKeyBundleOptions.fromOptions(options)); + + Coder> windowedStringCoder = + WindowedValues.getFullCoder(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE); + Coder> windowedIntCoder = + WindowedValues.getFullCoder(VarIntCoder.of(), GlobalWindow.Coder.INSTANCE); + + TupleTag mainOutputTag = new TupleTag<>("main"); + TupleTag sideOutputTag = new TupleTag<>("side"); + String serializedFn = + StringUtils.byteArrayToJsonString( + SerializableUtils.serializeToByteArray( + DoFnInfo.forFn( + new StringSplittableDoFn(), + WindowingStrategy.globalDefault(), + null /* side input views */, + StringUtf8Coder.of() /* input coder */, + ImmutableMap.of( + mainOutputTag, StringUtf8Coder.of(), sideOutputTag, VarIntCoder.of()), + mainOutputTag, + DoFnSchemaInformation.create(), + Collections.emptyMap()))); + + CloudObject cloudUserFn = CloudObject.forClassName("SplittableProcessFn"); + addString(cloudUserFn, PropertyNames.SERIALIZED_FN, serializedFn); + addObject( + cloudUserFn, + WorkerPropertyNames.RESTRICTION_CODER, + CloudObjects.asCloudObject( + KvCoder.of(OffsetRange.Coder.of(), VoidCoder.of()), /* sdkComponents= */ null)); + + InstructionInput cloudInput = new InstructionInput(); + cloudInput.setProducerInstructionIndex(0); + cloudInput.setOutputNum(0); + + MultiOutputInfo mainOutputInfo = new MultiOutputInfo(); + mainOutputInfo.setTag("main"); + MultiOutputInfo sideOutputInfo = new MultiOutputInfo(); + sideOutputInfo.setTag("side"); + + ParDoInstruction parDoInstruction = new ParDoInstruction(); + parDoInstruction.setInput(cloudInput); + parDoInstruction.setNumOutputs(2); + parDoInstruction.setMultiOutputInfos(ImmutableList.of(mainOutputInfo, sideOutputInfo)); + parDoInstruction.setUserFn(cloudUserFn); + + InstructionOutput mainOutput = new InstructionOutput(); + mainOutput.setName("MultiOutputSplittableParDo_main"); + mainOutput.setCodec(CloudObjects.asCloudObject(windowedStringCoder, /* sdkComponents= */ null)); + mainOutput.setOriginalName("MultiOutputSplittableParDoOriginal"); + mainOutput.setSystemName("MultiOutputSplittableParDo"); + + InstructionOutput sideOutput = new InstructionOutput(); + sideOutput.setName("MultiOutputSplittableParDo_side"); + sideOutput.setCodec(CloudObjects.asCloudObject(windowedIntCoder, /* sdkComponents= */ null)); + sideOutput.setOriginalName("MultiOutputSplittableParDoOriginal"); + sideOutput.setSystemName("MultiOutputSplittableParDo"); + + ParallelInstruction multiOutputSdfInstruction = new ParallelInstruction(); + multiOutputSdfInstruction.setParDo(parDoInstruction); + multiOutputSdfInstruction.setOutputs(ImmutableList.of(mainOutput, sideOutput)); + multiOutputSdfInstruction.setName("MultiOutputSplittableParDo"); + multiOutputSdfInstruction.setSystemName("MultiOutputSplittableParDo"); + multiOutputSdfInstruction.setOriginalName("MultiOutputSplittableParDoOriginal"); + + MapTask mapTask = new MapTask(); + mapTask.setStageName("testStage"); + mapTask.setSystemName("testSystem"); + mapTask.setInstructions( + Arrays.asList( + createWindmillReadInstruction("Read", windowedStringCoder), multiOutputSdfInstruction)); + mapTask.setFactory(Transport.getJsonFactory()); + + StageInfo stageInfo = StageInfo.create("testStage", "testSystem"); + ComputationState computationState = + new ComputationState( + "computationId", + mapTask, + mock(BoundedQueueExecutor.class), + Collections.emptyMap(), + mock(WindmillStateCache.ForComputation.class)); + + ComputationWorkExecutor workExecutor = + factory.createComputationWorkExecutor(stageInfo, computationState, "workLatencyTrackingId"); + + DataflowMapTaskExecutor mapTaskExecutor = (DataflowMapTaskExecutor) workExecutor.workExecutor(); + List allParDoOps = mapTaskExecutor.getParDoOperations(); + assertEquals(1, allParDoOps.size()); + + ParDoOperation multiOutputSdfOp = allParDoOps.get(0); + assertTrue(multiOutputSdfOp.shouldReportBytesRead()); + assertEquals(2, multiOutputSdfOp.receivers.length); + assertEquals( + windowedStringCoder, multiOutputSdfOp.getBytesReadCoder(multiOutputSdfOp.receivers[0])); + assertEquals( + windowedIntCoder, multiOutputSdfOp.getBytesReadCoder(multiOutputSdfOp.receivers[1])); + + String counterName = computationState.sourceBytesProcessCounterName(); + assertNotNull(multiOutputSdfOp.receivers[0].getOutputCounters().get(counterName)); + assertNotNull(multiOutputSdfOp.receivers[1].getOutputCounters().get(counterName)); + + // Process differently typed elements on main and side output receivers without + // ClassCastException. + WindowedValue stringElem = WindowedValues.valueInGlobalWindow("hello"); + WindowedValue intElem = WindowedValues.valueInGlobalWindow(12345); + multiOutputSdfOp.receivers[0].process(stringElem); + multiOutputSdfOp.receivers[1].process(intElem); + + long expectedStringBytes = + CoderUtils.encodeToByteArray(windowedStringCoder, stringElem, Coder.Context.NESTED).length; + long expectedIntBytes = + CoderUtils.encodeToByteArray(windowedIntCoder, intElem, Coder.Context.NESTED).length; + + @SuppressWarnings("unchecked") + Counter bytesReadCounter = + (Counter) mapTaskExecutor.getOutputCounters().getExistingCounter(counterName); + assertNotNull(bytesReadCounter); + assertEquals( + Long.valueOf(expectedStringBytes + expectedIntBytes), bytesReadCounter.getAggregate()); + } +}