Time out RestrictionTrackers.trySplit after 5 minutes on lock contention

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

@@ -45,8 +45,9 @@ 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 Progress lastProgress = Progress.NONE;
@@ -98,14 +99,42 @@ public RestrictionT currentRestriction() {
 
     @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);
-        needsProgressUpdate = true;
-        return result;
-      } finally {
-        updateProgressAndUnlock();
+        // 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

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

@@ -156,7 +156,7 @@ public synchronized SplitResult<Object> trySplit(double fractionOfRemainder) {
         }
       }
       isBlocked = false;
-      return null;
+      return SplitResult.of("primary", "residual");
     }
 
     @Override
@@ -262,4 +262,59 @@ public void testClaimObserversProgressNonBlockingOnTrySplit() throws Interrupted
     progress = ((HasProgress) tracker).getProgress();
     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();
+  }
 }