Adds bytes read reporting for SDF.

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<Node, Edge> 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 <V> TupleTag<V> tupleTag(MultiOutputInfo multiOutputInfo) {

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<String, ElementCounter> 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<String, ElementCounter> 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<String, ElementCounter> 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 {

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<ParDoOperation> getReadReportingParDoOperations() {
+    List<ParDoOperation> 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()) {

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;
+  }
 }

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 =

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<String, String> {
+    @ProcessElement
+    public void processElement(ProcessContext c, RestrictionTracker<OffsetRange, Long> 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<String> 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.<Edge>of(
+                MultiOutputInfoEdge.create(
+                    instructionNode
+                        .getParallelInstruction()
+                        .getParDo()
+                        .getMultiOutputInfos()
+                        .get(0))));
+
+    ArrayList<Operation> 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(