Adds bytes read reporting for SDF.

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

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

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

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/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.
    *

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,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<OutputReceiver, Coder<?>> 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 <V> TupleTag<V> tupleTag(MultiOutputInfo multiOutputInfo) {

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

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<K, InputT, AccumT, W extends BoundedWindow>
@@ -391,5 +397,10 @@ public void finishBundle() throws Exception {
 
     @Override
     public void abort() throws Exception {}
+
+    @Override
+    public @Nullable DoFnInfo<?, ?> getDoFnInfo() {
+      return null;
+    }
   }
 }

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

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<InputT, W>(
-                  helpers.fnInfo.getSideInputViews(),
-                  helpers.fnInfo.getInputCoder(),
-                  (WindowingStrategy<?, W>) helpers.fnInfo.getWindowingStrategy(),
+                  helpers.activeFnInfo.getSideInputViews(),
+                  helpers.activeFnInfo.getInputCoder(),
+                  (WindowingStrategy<?, W>) 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<?, W>) getDoFnInfo().getWindowingStrategy(),
+                    (WindowingStrategy<?, W>) helpers.activeFnInfo.getWindowingStrategy(),
                     (Collection<W>) unblockedElement.getWindows());
               }
             }
@@ -163,21 +163,23 @@ protected void onProcessWindowedValue(WindowedValue<InputT> elem) {
     }
     if (hasState) {
       helpers.registerStateCleanup(
-          (WindowingStrategy<?, W>) getDoFnInfo().getWindowingStrategy(), windowsProcessed);
+          (WindowingStrategy<?, W>) helpers.activeFnInfo.getWindowingStrategy(), windowsProcessed);
     }
   }
 
   @Override
   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<BoundedWindow> windowCoder =
         (Coder<BoundedWindow>)
-            (helpers.fnInfo != null ? helpers.fnInfo : helpers.doFnInstanceManager.peek())
+            (helpers.activeFnInfo != null
+                    ? helpers.activeFnInfo
+                    : helpers.doFnInstanceManager.peek())
                 .getWindowingStrategy()
                 .getWindowFn()
                 .windowCoder();
@@ -209,14 +211,20 @@ 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}.
    *
    * <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;
+  @Nullable DoFnInfo<?, ?> getActiveDoFnInfo() {
+    return helpers.activeFnInfo;
   }
 }