[release-2.77] Cherry-pick RestrictionTrackers lock timeout fixes (#40205, #40251)

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

@@ -45,11 +45,13 @@ public interface ClaimObserver<PositionT> {
    * RestrictionTracker}.
    */
   @ThreadSafe
-  private static class RestrictionTrackerObserver<RestrictionT, PositionT>
+  static class RestrictionTrackerObserver<RestrictionT, PositionT>
       extends RestrictionTracker<RestrictionT, PositionT> {
+    private static final int SPLIT_TIMEOUT_SEC = 300;
     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 +61,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 +83,7 @@ public boolean tryClaim(PositionT position) {
           return false;
         }
       } finally {
-        lock.unlock();
+        updateProgressAndUnlock();
       }
     }
 
@@ -81,19 +93,48 @@ public RestrictionT currentRestriction() {
       try {
         return delegate.currentRestriction();
       } finally {
-        lock.unlock();
+        updateProgressAndUnlock();
       }
     }
 
     @Override
     public SplitResult<RestrictionT> trySplit(double fractionOfRemainder) {
-      lock.lock();
+      return trySplit(fractionOfRemainder, SPLIT_TIMEOUT_SEC);
+    }
+
+    @VisibleForTesting
+    SplitResult<RestrictionT> trySplit(double fractionOfRemainder, int timeOutSec) {
+      // When fractionOfRemainder == 0 (a checkpoint), returning null has a special meaning in the
+      // RestrictionTracker contract: it MUST imply that the restriction tracker is done and there
+      // is no more work left to do. Therefore, we cannot time out and return null when
+      // fractionOfRemainder == 0, and must block until the lock is acquired.
+      if (fractionOfRemainder == 0) {
+        lock.lock();
+        try {
+          SplitResult<RestrictionT> result = delegate.trySplit(fractionOfRemainder);
+          needsProgressUpdate = true;
+          return result;
+        } finally {
+          updateProgressAndUnlock();
+        }
+      }
       try {
-        SplitResult<RestrictionT> result = delegate.trySplit(fractionOfRemainder);
-        return result;
-      } finally {
-        lock.unlock();
+        // For dynamic splits (fractionOfRemainder > 0), lock can be held long by a long-running
+        // tryClaim. We tolerate this scenario by returning null (declining to split) when lock
+        // timeout occurs.
+        if (lock.tryLock(timeOutSec, TimeUnit.SECONDS)) {
+          try {
+            SplitResult<RestrictionT> result = delegate.trySplit(fractionOfRemainder);
+            needsProgressUpdate = true;
+            return result;
+          } finally {
+            updateProgressAndUnlock();
+          }
+        }
+      } catch (InterruptedException e) {
+        Thread.currentThread().interrupt();
       }
+      return null;
     }
 
     @Override
@@ -102,7 +143,7 @@ public void checkDone() throws IllegalStateException {
       try {
         delegate.checkDone();
       } finally {
-        lock.unlock();
+        updateProgressAndUnlock();
       }
     }
 
@@ -112,10 +153,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 +173,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 +183,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
@@ -150,7 +156,7 @@ public synchronized SplitResult<Object> trySplit(double fractionOfRemainder) {
         }
       }
       isBlocked = false;
-      return null;
+      return SplitResult.of("primary", "residual");
     }
 
     @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,117 @@ 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);
+  }
+
+  @Test
+  public void testClaimObserversTrySplitNonBlockingOnTryClaim() throws InterruptedException {
+    RestrictionTrackerWithProgress withProgress = new RestrictionTrackerWithProgress(true, false);
+    RestrictionTracker<Object, Object> tracker =
+        RestrictionTrackers.observe(withProgress, new RestrictionTrackers.NoopClaimObserver<>());
+    Thread blocking = new Thread(() -> tracker.tryClaim(new Object()));
+    blocking.start();
+    withProgress.waitUntilBlocking(true);
+
+    // While tryClaim holds the lock, trySplit times out and returns null instead of blocking
+    // indefinitely.
+    SplitResult<Object> splitResult =
+        ((RestrictionTrackers.RestrictionTrackerObserver<Object, Object>) tracker).trySplit(0.5, 1);
+    assertEquals(null, splitResult);
+
+    withProgress.releaseLock();
+    withProgress.waitUntilBlocking(false);
+    blocking.join();
+
+    // Once tryClaim releases the lock, trySplit succeeds.
+    splitResult =
+        ((RestrictionTrackers.RestrictionTrackerObserver<Object, Object>) tracker).trySplit(0.5, 1);
+    assertEquals(SplitResult.of("primary", "residual"), splitResult);
+  }
+
+  @Test
+  public void testClaimObserversTrySplitZeroFractionBlocksOnTryClaim() throws InterruptedException {
+    RestrictionTrackerWithProgress withProgress = new RestrictionTrackerWithProgress(true, false);
+    RestrictionTracker<Object, Object> tracker =
+        RestrictionTrackers.observe(withProgress, new RestrictionTrackers.NoopClaimObserver<>());
+    Thread blocking = new Thread(() -> tracker.tryClaim(new Object()));
+    blocking.start();
+    withProgress.waitUntilBlocking(true);
+
+    // When fractionOfRemainder == 0, trySplit must block until the lock is released rather than
+    // timing out and returning null (since null implies the restriction tracker is done).
+    Thread releaseLater =
+        new Thread(
+            () -> {
+              try {
+                Thread.sleep(1500);
+              } catch (InterruptedException e) {
+                Thread.currentThread().interrupt();
+              }
+              withProgress.releaseLock();
+            });
+    releaseLater.start();
+
+    SplitResult<Object> splitResult =
+        ((RestrictionTrackers.RestrictionTrackerObserver<Object, Object>) tracker).trySplit(0.0, 1);
+    assertEquals(SplitResult.of("primary", "residual"), splitResult);
+    blocking.join();
+    releaseLater.join();
   }
 }