Adds backlog reporting support for non-fnapi based SDF's.

.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Dataflow_Streaming_Engine.json

@@ -1,4 +1,4 @@
 {
   "comment": "Modify this file in a trivial way to cause this test suite to run!",
-  "modification":  2,
+  "modification":  1,
 }

runners/core-java/src/main/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvoker.java

@@ -278,8 +278,109 @@ public FinishBundleContext finishBundleContext(DoFn<InputT, OutputT> doFn) {
     if (residual == null) {
       return new Result(null, cont, null, null);
     }
+    final KV<RestrictionT, KV<Instant, WatermarkEstimatorStateT>> residualForGetSize = residual;
+    // For a list of all DoFnInvoker arguments, see DoFn.java.
+    double backlogBytes =
+        invoker.invokeGetSize(
+            new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+              @Override
+              public String getErrorContext() {
+                return OutputAndTimeBoundedSplittableProcessElementInvoker.class.getSimpleName()
+                    + "/GetSize";
+              }
+
+              @Override
+              public InputT element(DoFn<InputT, OutputT> doFn) {
+                return element.getValue();
+              }
+
+              @Override
+              public Object restriction() {
+                return residualForGetSize.getKey();
+              }
+
+              @Override
+              public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+                return element.getTimestamp();
+              }
+
+              @Override
+              public RestrictionTracker<?, ?> restrictionTracker() {
+                return invoker.invokeNewTracker(
+                    new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
+                      @Override
+                      public String getErrorContext() {
+                        return OutputAndTimeBoundedSplittableProcessElementInvoker.class
+                                .getSimpleName()
+                            + "/NewTracker";
+                      }
+
+                      @Override
+                      public InputT element(DoFn<InputT, OutputT> doFn) {
+                        return element.getValue();
+                      }
+
+                      @Override
+                      public Object restriction() {
+                        return residualForGetSize.getKey();
+                      }
+
+                      @Override
+                      public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+                        return element.getTimestamp();
+                      }
+
+                      @Override
+                      public BoundedWindow window() {
+                        throw new IllegalStateException(
+                            "Attempting to access window outside of a windowed context");
+                      }
+
+                      @Override
+                      public PaneInfo paneInfo(DoFn<InputT, OutputT> doFn) {
+                        throw new IllegalStateException(
+                            "Attempting to access PaneInfo outside of a windowed context");
+                      }
+
+                      @Override
+                      public PipelineOptions pipelineOptions() {
+                        return pipelineOptions;
+                      }
+                    });
+              }
+
+              @Override
+              public BoundedWindow window() {
+                throw new IllegalStateException(
+                    "Attempting to access window outside of a windowed context");
+              }
+
+              @Override
+              public PaneInfo paneInfo(DoFn<InputT, OutputT> doFn) {
+                throw new IllegalStateException(
+                    "Attempting to access PaneInfo outside of a windowed context");
+              }
+
+              @Override
+              public PipelineOptions pipelineOptions() {
+                return pipelineOptions;
+              }
+
+              @Override
+              public Object sideInput(String tagId) {
+                PCollectionView<?> view = sideInputMapping.get(tagId);
+                if (view == null) {
+                  throw new IllegalArgumentException("calling getSideInput() with unknown view");
+                }
+                return processContext.sideInput(view);
+              }
+            });
     return new Result(
-        residual.getKey(), cont, residual.getValue().getKey(), residual.getValue().getValue());
+        residual.getKey(),
+        cont,
+        residual.getValue().getKey(),
+        residual.getValue().getValue(),
+        backlogBytes);
   }
 
   private class ProcessContext extends DoFn<InputT, OutputT>.ProcessContext

runners/core-java/src/main/java/org/apache/beam/runners/core/SplittableParDoViaKeyedWorkItems.java

@@ -22,6 +22,7 @@
 import com.google.auto.service.AutoService;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Consumer;
 import org.apache.beam.model.pipeline.v1.RunnerApi;
 import org.apache.beam.sdk.coders.ByteArrayCoder;
 import org.apache.beam.sdk.coders.Coder;
@@ -281,6 +282,7 @@ public static class ProcessFn<InputT, OutputT, RestrictionT, PositionT, Watermar
         processElementInvoker;
 
     private transient @Nullable DoFnInvoker<InputT, OutputT> invoker;
+    private transient @Nullable Consumer<Double> backlogBytesCallback;
 
     public ProcessFn(
         DoFn<InputT, OutputT> fn,
@@ -323,6 +325,10 @@ public void setProcessElementInvoker(
       this.processElementInvoker = invoker;
     }
 
+    public void setBacklogBytesCallback(Consumer<Double> backlogBytesCallback) {
+      this.backlogBytesCallback = backlogBytesCallback;
+    }
+
     public DoFn<InputT, OutputT> getFn() {
       return fn;
     }
@@ -622,6 +628,9 @@ public String getErrorContext() {
       } else {
         holdState.clear();
       }
+      if (backlogBytesCallback != null && result.getBacklogBytes() >= 0) {
+        backlogBytesCallback.accept(result.getBacklogBytes());
+      }
     }
 
     private DoFnInvoker.ArgumentProvider<InputT, OutputT> wrapOptionsAsSetup(

runners/core-java/src/main/java/org/apache/beam/runners/core/SplittableProcessElementInvoker.java

@@ -42,6 +42,10 @@ public class Result {
     private final DoFn.ProcessContinuation continuation;
     private final @Nullable Instant futureOutputWatermark;
     private final @Nullable WatermarkEstimatorStateT futureWatermarkEstimatorState;
+    private final double backlogBytes;
+
+    /* Constant representing an unknown amount of backlog. */
+    public static final double BACKLOG_UNKNOWN = -1.0;
 
     @SuppressFBWarnings(
         value = "NP_PARAMETER_MUST_BE_NONNULL_BUT_MARKED_AS_NULLABLE",
@@ -50,12 +54,27 @@ public Result(
         @Nullable RestrictionT residualRestriction,
         DoFn.ProcessContinuation continuation,
         @Nullable Instant futureOutputWatermark,
-        @Nullable WatermarkEstimatorStateT futureWatermarkEstimatorState) {
+        @Nullable WatermarkEstimatorStateT futureWatermarkEstimatorState,
+        double backlogBytes) {
       checkArgument(continuation != null, "continuation must not be null");
       this.continuation = continuation;
       this.residualRestriction = residualRestriction;
       this.futureOutputWatermark = futureOutputWatermark;
       this.futureWatermarkEstimatorState = futureWatermarkEstimatorState;
+      this.backlogBytes = backlogBytes;
+    }
+
+    public Result(
+        @Nullable RestrictionT residualRestriction,
+        DoFn.ProcessContinuation continuation,
+        @Nullable Instant futureOutputWatermark,
+        @Nullable WatermarkEstimatorStateT futureWatermarkEstimatorState) {
+      this(
+          residualRestriction,
+          continuation,
+          futureOutputWatermark,
+          futureWatermarkEstimatorState,
+          BACKLOG_UNKNOWN);
     }
 
     /**
@@ -76,6 +95,10 @@ public DoFn.ProcessContinuation getContinuation() {
     public @Nullable WatermarkEstimatorStateT getFutureWatermarkEstimatorState() {
       return futureWatermarkEstimatorState;
     }
+
+    public double getBacklogBytes() {
+      return backlogBytes;
+    }
   }
 
   /**

runners/core-java/src/main/java/org/apache/beam/runners/core/StepContext.java

@@ -36,4 +36,10 @@ public interface StepContext {
   default BundleFinalizer bundleFinalizer() {
     throw new UnsupportedOperationException("BundleFinalizer is unsupported.");
   }
+
+  /**
+   * Set the current backlog bytes for this step. This is mainly used by splittable DoFn to report
+   * the size of the residual restriction.
+   */
+  default void setBacklogBytes(double backlogBytes) {}
 }

runners/core-java/src/test/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvokerTest.java

@@ -95,6 +95,30 @@ public OffsetRange getInitialRestriction(@SuppressWarnings("unused") @Element Vo
     }
   }
 
+  private static class GetSizeFn extends DoFn<Void, String> {
+    @ProcessElement
+    public ProcessContinuation process(
+        ProcessContext c, RestrictionTracker<OffsetRange, Long> tracker) {
+      for (long i = tracker.currentRestriction().getFrom(); tracker.tryClaim(i); ++i) {
+        c.output(String.valueOf(i));
+        if (i == 2) {
+          return resume();
+        }
+      }
+      return stop();
+    }
+
+    @GetInitialRestriction
+    public OffsetRange getInitialRestriction() {
+      return new OffsetRange(0, 10);
+    }
+
+    @GetSize
+    public double getSize(@Restriction OffsetRange range) {
+      return range.getTo() - range.getFrom();
+    }
+  }
+
   private SplittableProcessElementInvoker<Void, String, OffsetRange, Long, Void>.Result runTest(
       int totalNumOutputs,
       Duration sleepBeforeFirstClaim,
@@ -103,11 +127,12 @@ private SplittableProcessElementInvoker<Void, String, OffsetRange, Long, Void>.R
       throws Exception {
     SomeFn fn = new SomeFn(sleepBeforeFirstClaim, numOutputsPerProcessCall, sleepBeforeEachOutput);
     OffsetRange initialRestriction = new OffsetRange(0, totalNumOutputs);
-    return runTest(fn, initialRestriction);
+    return runTest(fn, initialRestriction, Duration.standardSeconds(3));
   }
 
   private SplittableProcessElementInvoker<Void, String, OffsetRange, Long, Void>.Result runTest(
-      DoFn<Void, String> fn, OffsetRange initialRestriction) throws Exception {
+      DoFn<Void, String> fn, OffsetRange initialRestriction, Duration checkpointDuration)
+      throws Exception {
     SplittableProcessElementInvoker<Void, String, OffsetRange, Long, Void> invoker =
         new OutputAndTimeBoundedSplittableProcessElementInvoker<>(
             fn,
@@ -122,7 +147,7 @@ public <OutputT> void output(TupleTag<OutputT> tag, WindowedValue<OutputT> outpu
             NullSideInputReader.empty(),
             Executors.newSingleThreadScheduledExecutor(),
             1000,
-            Duration.standardSeconds(3),
+            checkpointDuration,
             () -> {
               throw new UnsupportedOperationException("BundleFinalizer not configured for test.");
             });
@@ -215,7 +240,7 @@ public OffsetRange getInitialRestriction(
           }
         };
     e.expectMessage("Output is not allowed before tryClaim()");
-    runTest(brokenFn, new OffsetRange(0, 5));
+    runTest(brokenFn, new OffsetRange(0, 5), Duration.standardSeconds(3));
   }
 
   @Test
@@ -235,6 +260,18 @@ public OffsetRange getInitialRestriction(
           }
         };
     e.expectMessage("Output is not allowed after a failed tryClaim()");
-    runTest(brokenFn, new OffsetRange(0, 5));
+    runTest(brokenFn, new OffsetRange(0, 5), Duration.standardSeconds(3));
+  }
+
+  @Test
+  public void testBacklogBytes() throws Exception {
+    GetSizeFn fn = new GetSizeFn();
+    OffsetRange initialRestriction = new OffsetRange(0, 10);
+    // Set a high checkpoint duration to prevent flakiness caused by early checkpointing.
+    SplittableProcessElementInvoker<Void, String, OffsetRange, Long, Void>.Result res =
+        runTest(fn, initialRestriction, Duration.standardMinutes(3));
+    // GetSizeFn claims 3 elements and then takes a checkpoint.
+    assertEquals(7.0, res.getBacklogBytes(), 0.001);
+    assertEquals(new OffsetRange(3, 10), res.getResidualRestriction());
   }
 }