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..7589a5036b58 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 @@ -141,11 +141,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 +246,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 +270,29 @@ 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. + Coder bytesReadCoder = null; + if ("SplittableProcessFn".equals(cloudUserFn.getClassName())) { + Node firstOutput = Iterables.getFirst(network.successors(node), null); + if (firstOutput instanceof OutputReceiverNode) { + bytesReadCoder = ((OutputReceiverNode) firstOutput).getCoder(); + } + } OutputReceiver[] receivers = getOutputReceivers(network, node); - return OperationNode.create(new ParDoOperation(fn, receivers, operationContext)); + return OperationNode.create( + new ParDoOperation(fn, receivers, operationContext, bytesReadCoder)); } private static TupleTag tupleTag(MultiOutputInfo multiOutputInfo) { 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..d549d9163c7d 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 @@ -68,6 +68,8 @@ 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.OutputObjectAndByteCounter; +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.Windmill.GlobalDataId; @@ -796,19 +798,42 @@ 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(); - - return Optional.ofNullable(counters.get(sourceBytesCounterName)) - .map(counter -> ((OutputObjectAndByteCounter) counter).getByteCount().getAndReset()) - .orElse(0L); + // 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 = 0L; + ReadOperation readOperation = mapTaskExecutor.getReadOperation(); + if (readOperation != null + && readOperation.receivers != null + && readOperation.receivers.length > 0 + && readOperation.receivers[0] != null) { + HashMap readOpCounters = + readOperation.receivers[0].getOutputCounters(); + totalBytes += + Optional.ofNullable(readOpCounters.get(sourceBytesCounterName)) + .map(counter -> ((OutputObjectAndByteCounter) counter).getByteCount().getAndReset()) + .orElse(0L); + } + for (ParDoOperation parDoOp : mapTaskExecutor.getReadReportingParDoOperations()) { + if (parDoOp.receivers != null + && parDoOp.receivers.length > 0 + && parDoOp.receivers[0] != null) { + HashMap splittableProcessOpCounters = + parDoOp.receivers[0].getOutputCounters(); + totalBytes += + Optional.ofNullable(splittableProcessOpCounters.get(sourceBytesCounterName)) + .map(counter -> ((OutputObjectAndByteCounter) counter).getByteCount().getAndReset()) + .orElse(0L); + } + } + return totalBytes; } public boolean advance() throws CoderException { 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..68c94dc89f3c 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,21 @@ public ReadOperation getReadOperation() { return (ReadOperation) readOperation; } + public List getReadReportingParDoOperations() { + List reportingParDoOperations = Lists.newArrayList(); + if (operations != null) { + for (Operation op : operations) { + if (op instanceof ParDoOperation) { + ParDoOperation parDoOp = (ParDoOperation) op; + if (parDoOp.shouldReportBytesRead()) { + reportingParDoOperations.add(parDoOp); + } + } + } + } + return reportingParDoOperations; + } + 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/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..224a3b39577a 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,29 @@ package org.apache.beam.runners.dataflow.worker.util.common.worker; import java.io.Closeable; +import org.apache.beam.sdk.coders.Coder; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; import org.checkerframework.checker.nullness.qual.Nullable; /** A ParDo mapping function. */ public class ParDoOperation extends ReceivingOperation { private final ParDoFn fn; + // If present, this ParDo should report the amount of bytes read using this coder to serialize + // elements. + private final @Nullable Coder bytesReadCoder; public ParDoOperation(ParDoFn fn, OutputReceiver[] outputReceivers, OperationContext context) { + this(fn, outputReceivers, context, null); + } + + public ParDoOperation( + ParDoFn fn, + OutputReceiver[] outputReceivers, + OperationContext context, + @Nullable Coder bytesReadCoder) { super(outputReceivers, context); this.fn = fn; + this.bytesReadCoder = bytesReadCoder; } @Override @@ -82,4 +95,12 @@ public boolean supportsRestart() { public ParDoFn getFn() throws Exception { return fn; } + + public boolean shouldReportBytesRead() { + return bytesReadCoder != null; + } + + public @Nullable Coder getBytesReadCoder() { + return bytesReadCoder; + } } 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..9356f36f02cc 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 @@ -53,7 +53,9 @@ import org.apache.beam.runners.dataflow.worker.streaming.harness.StreamingCounters; import org.apache.beam.runners.dataflow.worker.streaming.sideinput.SideInputStateFetcherFactory; import org.apache.beam.runners.dataflow.worker.util.common.worker.MapTaskExecutor; +import org.apache.beam.runners.dataflow.worker.util.common.worker.Operation; import org.apache.beam.runners.dataflow.worker.util.common.worker.OutputObjectAndByteCounter; +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,23 +168,27 @@ private static void trackAutoscalingBytesRead( MapTask mapTask, Nodes.ParallelInstructionNode readNode, Coder readCoder, - ReadOperation readOperation, + Operation operation, MapTaskExecutor mapTaskExecutor, String counterName) { - NameContext nameContext = - NameContext.create( - mapTask.getStageName(), - readNode.getParallelInstruction().getOriginalName(), - readNode.getParallelInstruction().getSystemName(), - readNode.getParallelInstruction().getName()); - readOperation.receivers[0].addOutputCounter( - counterName, - new OutputObjectAndByteCounter( - new IntrinsicMapTaskExecutorFactory.ElementByteSizeObservableCoder<>(readCoder), - mapTaskExecutor.getOutputCounters(), - nameContext) - .setSamplingPeriod(100) - .countBytes(counterName)); + if (operation.receivers != null + && operation.receivers.length > 0 + && operation.receivers[0] != null) { + NameContext nameContext = + NameContext.create( + mapTask.getStageName(), + readNode.getParallelInstruction().getOriginalName(), + readNode.getParallelInstruction().getSystemName(), + readNode.getParallelInstruction().getName()); + operation.receivers[0].addOutputCounter( + counterName, + new OutputObjectAndByteCounter( + new IntrinsicMapTaskExecutorFactory.ElementByteSizeObservableCoder<>(readCoder), + mapTaskExecutor.getOutputCounters(), + nameContext) + .setSamplingPeriod(100) + .countBytes(counterName)); + } } private static ReadOperation getValidatedReadOperation(MapTaskExecutor mapTaskExecutor) { @@ -234,6 +240,19 @@ ComputationWorkExecutor createComputationWorkExecutor( readOperation, mapTaskExecutor, computationState.sourceBytesProcessCounterName()); + } else { + for (ParDoOperation parDoOp : mapTaskExecutor.getReadReportingParDoOperations()) { + Coder bytesReadCoder = parDoOp.getBytesReadCoder(); + if (bytesReadCoder != null) { + trackAutoscalingBytesRead( + mapTask, + readNode, + bytesReadCoder, + parDoOp, + 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..13c1a46704f0 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; @@ -567,6 +572,93 @@ 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()); + assertNull(parDoOperation.getBytesReadCoder()); + } + + 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( + WindowedValues.getValueOnlyCoder(StringUtf8Coder.of()), parDoOperation.getBytesReadCoder()); } static ParallelInstruction createPartialGroupByKeyInstruction( 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..99c598e56088 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; @@ -943,4 +950,150 @@ 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, 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_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, 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, 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/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..0cfda97fd149 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 @@ -62,6 +62,8 @@ 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; @@ -573,4 +575,36 @@ public void testExceptionAndThenCloseAbortsJustOnce() throws Exception { // Closing after already closed should not call abort again. executor.close(); } + + @Test + public void testGetReadReportingParDoOperations() throws Exception { + TestOutputReceiver receiver = + new TestOutputReceiver(counterSet, NameContextsForTests.nameContextForTest()); + ReadOperation readOp = new TestReadOperation(receiver, createContext("ReadOperation")); + ParDoOperation regularParDo = + new ParDoOperation( + new NoopParDoFn(), new OutputReceiver[0], createContext("RegularParDo"), null); + ParDoOperation reportingParDo1 = + new ParDoOperation( + new NoopParDoFn(), + new OutputReceiver[0], + createContext("ReportingParDo1"), + StringUtf8Coder.of()); + ParDoOperation reportingParDo2 = + new ParDoOperation( + new NoopParDoFn(), + new OutputReceiver[0], + createContext("ReportingParDo2"), + VarIntCoder.of()); + + List operations = + Arrays.asList(readOp, regularParDo, reportingParDo1, reportingParDo2); + ExecutionStateTracker stateTracker = ExecutionStateTracker.newForTest(); + try (MapTaskExecutor executor = new MapTaskExecutor(operations, counterSet, stateTracker)) { + List reportingOps = executor.getReadReportingParDoOperations(); + assertThat(reportingOps, contains(reportingParDo1, reportingParDo2)); + assertEquals(StringUtf8Coder.of(), reportingOps.get(0).getBytesReadCoder()); + assertEquals(VarIntCoder.of(), reportingOps.get(1).getBytesReadCoder()); + } + } } 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..69cbac36632e --- /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,337 @@ +/* + * 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.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.Iterables; +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 readReportingOps = mapTaskExecutor.getReadReportingParDoOperations(); + assertEquals(2, readReportingOps.size()); + + ParDoOperation stringSplittableOp = + readReportingOps.stream() + .filter(op -> windowedStringCoder.equals(op.getBytesReadCoder())) + .findFirst() + .orElseThrow( + () -> new AssertionError("Expected ParDoOperation with windowedStringCoder")); + ParDoOperation intSplittableOp = + readReportingOps.stream() + .filter(op -> windowedIntCoder.equals(op.getBytesReadCoder())) + .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. + List allParDoOps = + ImmutableList.copyOf(Iterables.filter(mapTaskExecutor.operations, ParDoOperation.class)); + assertEquals(3, allParDoOps.size()); + ParDoOperation regularParDoOp = + allParDoOps.stream() + .filter(op -> !op.shouldReportBytesRead()) + .findFirst() + .orElseThrow(() -> new AssertionError("Expected non-reporting ParDoOperation")); + assertNull(regularParDoOp.getBytesReadCoder()); + 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()); + } +}