Ensure 0 backlog is sent when finishing processing restrictions.

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

@@ -282,7 +282,7 @@ public FinishBundleContext finishBundleContext(DoFn<InputT, OutputT> doFn) {
       processContext.tracker.checkDone();
     }
     if (residual == null) {
-      return new Result(null, cont, null, null);
+      return new Result(null, cont, null, null, 0.0);
     }
     final KV<RestrictionT, KV<Instant, WatermarkEstimatorStateT>> residualForGetSize = residual;
     // For a list of all DoFnInvoker arguments, see DoFn.java.

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

@@ -606,6 +606,9 @@ public String getErrorContext() {
         restrictionState.clear();
         watermarkEstimatorState.clear();
         holdState.clear();
+        if (backlogBytesCallback != null) {
+          backlogBytesCallback.accept(0.0);
+        }
         return;
       }
 

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

@@ -214,6 +214,7 @@ public void testInvokeProcessElementVoluntaryReturnStop() throws Exception {
         runTest(5, Duration.ZERO, Integer.MAX_VALUE, Duration.millis(100));
     assertFalse(res.getContinuation().shouldResume());
     assertNull(res.getResidualRestriction());
+    assertEquals(0.0, res.getBacklogBytes(), 0.001);
   }
 
   @Test
@@ -274,4 +275,16 @@ public void testBacklogBytes() throws Exception {
     assertEquals(7.0, res.getBacklogBytes(), 0.001);
     assertEquals(new OffsetRange(3, 10), res.getResidualRestriction());
   }
+
+  @Test
+  public void testBacklogBytesWhenDone() throws Exception {
+    GetSizeFn fn = new GetSizeFn();
+    OffsetRange initialRestriction = new OffsetRange(0, 2);
+    // 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 2 elements and finishes.
+    assertEquals(0.0, res.getBacklogBytes(), 0.001);
+    assertNull(res.getResidualRestriction());
+  }
 }

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

@@ -763,6 +763,10 @@ public void testReportsBacklog() throws Exception {
       // The residual range should be [3, 10), so size is 7.
       assertEquals(1, backlogs.size());
       assertEquals(7.0, backlogs.get(0), 0.001);
+
+      assertTrue(tester.advanceProcessingTimeBy(Duration.standardSeconds(1)));
+      assertEquals(2, backlogs.size());
+      assertEquals(0.0, backlogs.get(1), 0.001);
     }
   }
 
@@ -788,6 +792,34 @@ public void testReportsBacklogWithoutGetSize() throws Exception {
       // The residual range should be [3, 10), so size is 7.
       assertEquals(1, backlogs.size());
       assertEquals(7.0, backlogs.get(0), 0.001);
+
+      assertTrue(tester.advanceProcessingTimeBy(Duration.standardSeconds(1)));
+      assertEquals(2, backlogs.size());
+      assertEquals(0.0, backlogs.get(1), 0.001);
+    }
+  }
+
+  @Test
+  public void testReportsZeroBacklogWhenDone() throws Exception {
+    DoFn<Integer, String> fn = new GetSizeFn();
+    Instant base = Instant.now();
+    final List<Double> backlogs = new ArrayList<>();
+
+    try (ProcessFnTester<Integer, String, OffsetRange, Long, Void> tester =
+        new ProcessFnTester<>(
+            base,
+            fn,
+            BigEndianIntegerCoder.of(),
+            SerializableCoder.of(OffsetRange.class),
+            VoidCoder.of(),
+            MAX_OUTPUTS_PER_BUNDLE,
+            MAX_BUNDLE_DURATION)) {
+      tester.processFn.setBacklogBytesCallback(backlogs::add);
+
+      // OffsetRange(0, 2) completes immediately without resume.
+      tester.startElement(42, new OffsetRange(0, 2));
+      assertEquals(1, backlogs.size());
+      assertEquals(0.0, backlogs.get(0), 0.001);
     }
   }
 }