Return last evaluated progress on lock timeout in RestrictionTrackers.getProgress

sdks/java/core/src/main/java/org/apache/beam/sdk/fn/splittabledofn/RestrictionTrackers.java

@@ -49,7 +49,8 @@ private static class RestrictionTrackerObserver<RestrictionT, PositionT>
       extends RestrictionTracker<RestrictionT, PositionT> {
     protected final RestrictionTracker<RestrictionT, PositionT> delegate;
     protected ReentrantLock lock = new ReentrantLock();
-    protected volatile boolean hasInitialProgress = false;
+    protected volatile Progress lastProgress = Progress.NONE;
+    protected volatile boolean needsProgressUpdate = false;
     private final ClaimObserver<PositionT> claimObserver;
 
     protected RestrictionTrackerObserver(
@@ -59,6 +60,16 @@ protected RestrictionTrackerObserver(
       this.claimObserver = claimObserver;
     }
 
+    protected void updateProgressAndUnlock() {
+      try {
+        if (needsProgressUpdate) {
+          updateProgressBlocking();
+        }
+      } finally {
+        lock.unlock();
+      }
+    }
+
     @Override
     public boolean tryClaim(PositionT position) {
       lock.lock();
@@ -71,7 +82,7 @@ public boolean tryClaim(PositionT position) {
           return false;
         }
       } finally {
-        lock.unlock();
+        updateProgressAndUnlock();
       }
     }
 
@@ -81,7 +92,7 @@ public RestrictionT currentRestriction() {
       try {
         return delegate.currentRestriction();
       } finally {
-        lock.unlock();
+        updateProgressAndUnlock();
       }
     }
 
@@ -90,9 +101,10 @@ public SplitResult<RestrictionT> trySplit(double fractionOfRemainder) {
       lock.lock();
       try {
         SplitResult<RestrictionT> result = delegate.trySplit(fractionOfRemainder);
+        needsProgressUpdate = true;
         return result;
       } finally {
-        lock.unlock();
+        updateProgressAndUnlock();
       }
     }
 
@@ -102,7 +114,7 @@ public void checkDone() throws IllegalStateException {
       try {
         delegate.checkDone();
       } finally {
-        lock.unlock();
+        updateProgressAndUnlock();
       }
     }
 
@@ -112,10 +124,13 @@ public IsBounded isBounded() {
     }
 
     /** Evaluate progress if requested. */
-    protected Progress getProgressBlocking() {
+    protected void updateProgressBlocking() {
       lock.lock();
       try {
-        return ((HasProgress) delegate).getProgress();
+        needsProgressUpdate = false;
+        if (delegate instanceof HasProgress) {
+          lastProgress = ((HasProgress) delegate).getProgress();
+        }
       } finally {
         lock.unlock();
       }
@@ -129,7 +144,7 @@ protected Progress getProgressBlocking() {
   @ThreadSafe
   static class RestrictionTrackerObserverWithProgress<RestrictionT, PositionT>
       extends RestrictionTrackerObserver<RestrictionT, PositionT> implements HasProgress {
-    private static final int FIRST_PROGRESS_TIMEOUT_SEC = 60;
+    private static final int PROGRESS_TIMEOUT_SEC = 60;
 
     protected RestrictionTrackerObserverWithProgress(
         RestrictionTracker<RestrictionT, PositionT> delegate,
@@ -139,32 +154,29 @@ protected RestrictionTrackerObserverWithProgress(
 
     @Override
     public Progress getProgress() {
-      return getProgress(FIRST_PROGRESS_TIMEOUT_SEC);
+      return getProgress(PROGRESS_TIMEOUT_SEC);
     }
 
     @VisibleForTesting
     Progress getProgress(int timeOutSec) {
-      if (!hasInitialProgress) {
-        Progress progress = Progress.NONE;
-        try {
-          // lock can be held long by long-running tryClaim/trySplit. We tolerate this scenario
-          // by returning zero progress when initial progress never evaluated before due to lock
-          // timeout.
-          if (lock.tryLock(timeOutSec, TimeUnit.SECONDS)) {
-            try {
-              progress = getProgressBlocking();
-              hasInitialProgress = true;
-            } finally {
-              lock.unlock();
-            }
+      try {
+        // lock can be held long by long-running tryClaim/trySplit. We tolerate this scenario
+        // by returning the last evaluated progress (or zero progress if never evaluated before)
+        // when lock timeout occurs.
+        if (lock.tryLock(timeOutSec, TimeUnit.SECONDS)) {
+          try {
+            updateProgressBlocking();
+          } finally {
+            lock.unlock();
           }
-        } catch (InterruptedException e) {
-          Thread.currentThread().interrupt();
+        } else {
+          needsProgressUpdate = true;
         }
-        return progress;
-      } else {
-        return getProgressBlocking();
+      } catch (InterruptedException e) {
+        needsProgressUpdate = true;
+        Thread.currentThread().interrupt();
       }
+      return lastProgress;
     }
   }
 

sdks/java/core/src/test/java/org/apache/beam/sdk/fn/splittabledofn/RestrictionTrackersTest.java

@@ -103,7 +103,9 @@ private static class RestrictionTrackerWithProgress extends RestrictionTracker<O
     private boolean blockTryClaim;
     private boolean blockTrySplit;
     private volatile boolean isBlocked;
+    private volatile Progress currentProgress = REPORT_PROGRESS;
     public static final Progress REPORT_PROGRESS = Progress.from(2.0, 3.0);
+    public static final Progress UPDATED_PROGRESS = Progress.from(4.0, 1.0);
 
     public RestrictionTrackerWithProgress() {
       this(false, false);
@@ -117,7 +119,11 @@ public RestrictionTrackerWithProgress(boolean blockTryClaim, boolean blockTrySpl
 
     @Override
     public Progress getProgress() {
-      return REPORT_PROGRESS;
+      return currentProgress;
+    }
+
+    public void setProgress(Progress progress) {
+      this.currentProgress = progress;
     }
 
     @Override
@@ -161,6 +167,14 @@ public IsBounded isBounded() {
       return IsBounded.BOUNDED;
     }
 
+    public synchronized void setBlockTryClaim(boolean blockTryClaim) {
+      this.blockTryClaim = blockTryClaim;
+    }
+
+    public synchronized void setBlockTrySplit(boolean blockTrySplit) {
+      this.blockTrySplit = blockTrySplit;
+    }
+
     public synchronized void releaseLock() {
       blockTrySplit = false;
       blockTryClaim = false;
@@ -190,29 +204,62 @@ public void testClaimObserversProgressNonBlockingOnTryClaim() throws Interrupted
     Thread blocking = new Thread(() -> tracker.tryClaim(new Object()));
     blocking.start();
     withProgress.waitUntilBlocking(true);
+    // Times out while first tryClaim holds lock; returns NONE and sets needsProgressUpdate = true
     RestrictionTracker.Progress progress =
         ((RestrictionTrackers.RestrictionTrackerObserverWithProgress) tracker).getProgress(1);
     assertEquals(RestrictionTracker.Progress.NONE, progress);
+    // When first tryClaim finishes, updateProgressAndUnlock() sees needsProgressUpdate == true and
+    // evaluates REPORT_PROGRESS before releasing the lock.
     withProgress.releaseLock();
     withProgress.waitUntilBlocking(false);
-    progress = ((HasProgress) tracker).getProgress();
+    blocking.join();
+
+    // Even if a second blocking tryClaim immediately grabs the lock before getProgress is called
+    // again, getProgress(1) returns REPORT_PROGRESS (updated during first tryClaim's
+    // updateProgressAndUnlock).
+    withProgress.setProgress(RestrictionTrackerWithProgress.UPDATED_PROGRESS);
+    withProgress.setBlockTryClaim(true);
+    Thread secondBlocking = new Thread(() -> tracker.tryClaim(new Object()));
+    secondBlocking.start();
+    withProgress.waitUntilBlocking(true);
+    progress =
+        ((RestrictionTrackers.RestrictionTrackerObserverWithProgress) tracker).getProgress(1);
     assertEquals(RestrictionTrackerWithProgress.REPORT_PROGRESS, progress);
+    withProgress.releaseLock();
+    withProgress.waitUntilBlocking(false);
+    secondBlocking.join();
+    progress = ((HasProgress) tracker).getProgress();
+    assertEquals(RestrictionTrackerWithProgress.UPDATED_PROGRESS, progress);
   }
 
   @Test
   public void testClaimObserversProgressNonBlockingOnTrySplit() throws InterruptedException {
     RestrictionTrackerWithProgress withProgress = new RestrictionTrackerWithProgress(false, true);
     RestrictionTracker<Object, Object> tracker =
         RestrictionTrackers.observe(withProgress, new RestrictionTrackers.NoopClaimObserver<>());
+    // trySplit unconditionally refreshes lastProgress via updateProgressAndUnlock()
     Thread blocking = new Thread(() -> tracker.trySplit(0.5));
     blocking.start();
     withProgress.waitUntilBlocking(true);
+    withProgress.releaseLock();
+    withProgress.waitUntilBlocking(false);
+    blocking.join();
+
+    // Even though getProgress was never called before, trySplit unconditionally updated
+    // lastProgress to REPORT_PROGRESS. If a subsequent tryClaim blocks, getProgress(1) returns
+    // REPORT_PROGRESS rather than NONE or a stale pre-split progress.
+    withProgress.setProgress(RestrictionTrackerWithProgress.UPDATED_PROGRESS);
+    withProgress.setBlockTryClaim(true);
+    Thread secondBlocking = new Thread(() -> tracker.tryClaim(new Object()));
+    secondBlocking.start();
+    withProgress.waitUntilBlocking(true);
     RestrictionTracker.Progress progress =
         ((RestrictionTrackers.RestrictionTrackerObserverWithProgress) tracker).getProgress(1);
-    assertEquals(RestrictionTracker.Progress.NONE, progress);
+    assertEquals(RestrictionTrackerWithProgress.REPORT_PROGRESS, progress);
     withProgress.releaseLock();
     withProgress.waitUntilBlocking(false);
+    secondBlocking.join();
     progress = ((HasProgress) tracker).getProgress();
-    assertEquals(RestrictionTrackerWithProgress.REPORT_PROGRESS, progress);
+    assertEquals(RestrictionTrackerWithProgress.UPDATED_PROGRESS, progress);
   }
 }