Adds bytes read reporting for SDF.

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

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<Node, Edge> network,
@@ -272,18 +274,41 @@ 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 (isSplittable(fn)) {
+      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 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 <V> TupleTag<V> tupleTag(MultiOutputInfo multiOutputInfo) {

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFn.java

@@ -34,7 +34,6 @@
 import org.apache.beam.sdk.values.TupleTag;
 import org.apache.beam.sdk.values.WindowedValue;
 import org.apache.beam.sdk.values.WindowingStrategy;
-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.Lists;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
@@ -112,7 +111,7 @@ protected void onStartKey() {
                 // These elements are now processed. Register cleanup timers for all the unblocked
                 // windows.
                 helpers.registerStateCleanup(
-                    (WindowingStrategy<?, W>) getDoFnInfo().getWindowingStrategy(),
+                    (WindowingStrategy<?, W>) helpers.fnInfo.getWindowingStrategy(),
                     (Collection<W>) unblockedElement.getWindows());
               }
             }
@@ -163,7 +162,7 @@ protected void onProcessWindowedValue(WindowedValue<InputT> elem) {
     }
     if (hasState) {
       helpers.registerStateCleanup(
-          (WindowingStrategy<?, W>) getDoFnInfo().getWindowingStrategy(), windowsProcessed);
+          (WindowingStrategy<?, W>) helpers.fnInfo.getWindowingStrategy(), windowsProcessed);
     }
   }
 
@@ -176,11 +175,7 @@ public void processTimers() throws Exception {
     // windowCoder (allowing us to get it without a DoFnInfo instance) or (2) check whether timers
     // exist without actually decoding them.
     Coder<BoundedWindow> windowCoder =
-        (Coder<BoundedWindow>)
-            (helpers.fnInfo != null ? helpers.fnInfo : helpers.doFnInstanceManager.peek())
-                .getWindowingStrategy()
-                .getWindowFn()
-                .windowCoder();
+        (Coder<BoundedWindow>) getDoFnInfo().getWindowingStrategy().getWindowFn().windowCoder();
     helpers.processTimers(
         SimpleParDoFnHelpers.TimerType.USER,
         helpers.userStepContext,
@@ -209,14 +204,9 @@ public void abort() throws Exception {
     helpers.abort();
   }
 
-  /**
-   * Returns the {@link DoFnInfo} currently being used by this {@link SimpleParDoFn}.
-   *
-   * <p>May be null if no element has been processed yet, or if the {@link SimpleParDoFn} has
-   * finished.
-   */
-  @VisibleForTesting
-  @Nullable DoFnInfo<?, ?> getDoFnInfo() {
-    return helpers.fnInfo;
+  /** Returns the {@link DoFnInfo} used by this {@link SimpleParDoFn}. */
+  @Override
+  public DoFnInfo<?, ?> getDoFnInfo() throws Exception {
+    return helpers.fnInfo != null ? helpers.fnInfo : helpers.doFnInstanceManager.peek();
   }
 }

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputParDoFn.java

@@ -45,7 +45,6 @@
 import org.apache.beam.sdk.values.TupleTag;
 import org.apache.beam.sdk.values.WindowedValue;
 import org.apache.beam.sdk.values.WindowingStrategy;
-import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
 @SuppressWarnings({
@@ -138,7 +137,7 @@ protected void onStartKey() {
                 // These elements are now processed. Register cleanup timers for all the unblocked
                 // windows.
                 helpers.registerStateCleanup(
-                    (WindowingStrategy<?, W>) getDoFnInfo().getWindowingStrategy(), windows);
+                    (WindowingStrategy<?, W>) helpers.fnInfo.getWindowingStrategy(), windows);
               }
             });
       }
@@ -174,11 +173,7 @@ public void processTimers() throws Exception {
     // windowCoder (allowing us to get it without a DoFnInfo instance) or (2) check whether timers
     // exist without actually decoding them.
     Coder<BoundedWindow> windowCoder =
-        (Coder<BoundedWindow>)
-            (helpers.fnInfo != null ? helpers.fnInfo : helpers.doFnInstanceManager.peek())
-                .getWindowingStrategy()
-                .getWindowFn()
-                .windowCoder();
+        (Coder<BoundedWindow>) getDoFnInfo().getWindowingStrategy().getWindowFn().windowCoder();
     helpers.processTimers(
         SimpleParDoFnHelpers.TimerType.FAIL_USER,
         helpers.userStepContext,
@@ -231,18 +226,13 @@ protected void onProcessWindowedValue(WindowedValue<KeyedWorkItem<K, InputT>> el
     }
     if (hasState) {
       helpers.registerStateCleanup(
-          (WindowingStrategy<?, W>) getDoFnInfo().getWindowingStrategy(), windowsProcessed);
+          (WindowingStrategy<?, W>) helpers.fnInfo.getWindowingStrategy(), windowsProcessed);
     }
   }
 
-  /**
-   * Returns the {@link DoFnInfo} currently being used by this {@link SimpleParDoFn}.
-   *
-   * <p>May be null if no element has been processed yet, or if the {@link SimpleParDoFn} has
-   * finished.
-   */
-  @VisibleForTesting
-  @Nullable DoFnInfo<?, ?> getDoFnInfo() {
-    return helpers.fnInfo;
+  /** Returns the {@link DoFnInfo} used by this {@link StreamingKeyedWorkItemSideInputParDoFn}. */
+  @Override
+  public DoFnInfo<?, ?> getDoFnInfo() throws Exception {
+    return helpers.fnInfo != null ? helpers.fnInfo : helpers.doFnInstanceManager.peek();
   }
 }

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,43 @@ 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.getParDoOperations()) {
+      if (parDoOp.shouldReportBytesRead()
+          && 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,18 @@ public ReadOperation getReadOperation() {
     return (ReadOperation) readOperation;
   }
 
+  public List<ParDoOperation> getParDoOperations() {
+    List<ParDoOperation> 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()) {

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,8 @@ public interface ParDoFn {
   void finishBundle() throws Exception;
 
   void abort() throws Exception;
+
+  default @Nullable DoFnInfo<?, ?> getDoFnInfo() throws Exception {
+    return null;
+  }
 }

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