diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/splittabledofn/RestrictionTrackers.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/splittabledofn/RestrictionTrackers.java index 370871fd9cae..ded09b10649b 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/splittabledofn/RestrictionTrackers.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/splittabledofn/RestrictionTrackers.java @@ -45,8 +45,9 @@ public interface ClaimObserver { * RestrictionTracker}. */ @ThreadSafe - private static class RestrictionTrackerObserver + static class RestrictionTrackerObserver extends RestrictionTracker { + private static final int SPLIT_TIMEOUT_SEC = 300; protected final RestrictionTracker delegate; protected ReentrantLock lock = new ReentrantLock(); protected volatile Progress lastProgress = Progress.NONE; @@ -98,14 +99,42 @@ public RestrictionT currentRestriction() { @Override public SplitResult trySplit(double fractionOfRemainder) { - lock.lock(); + return trySplit(fractionOfRemainder, SPLIT_TIMEOUT_SEC); + } + + @VisibleForTesting + SplitResult 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 result = delegate.trySplit(fractionOfRemainder); + needsProgressUpdate = true; + return result; + } finally { + updateProgressAndUnlock(); + } + } try { - SplitResult 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 result = delegate.trySplit(fractionOfRemainder); + needsProgressUpdate = true; + return result; + } finally { + updateProgressAndUnlock(); + } + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); } + return null; } @Override diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/splittabledofn/RestrictionTrackersTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/splittabledofn/RestrictionTrackersTest.java index e24033224782..3708f71587df 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/splittabledofn/RestrictionTrackersTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/splittabledofn/RestrictionTrackersTest.java @@ -156,7 +156,7 @@ public synchronized SplitResult 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 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 splitResult = + ((RestrictionTrackers.RestrictionTrackerObserver) 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) tracker).trySplit(0.5, 1); + assertEquals(SplitResult.of("primary", "residual"), splitResult); + } + + @Test + public void testClaimObserversTrySplitZeroFractionBlocksOnTryClaim() throws InterruptedException { + RestrictionTrackerWithProgress withProgress = new RestrictionTrackerWithProgress(true, false); + RestrictionTracker 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 splitResult = + ((RestrictionTrackers.RestrictionTrackerObserver) tracker).trySplit(0.0, 1); + assertEquals(SplitResult.of("primary", "residual"), splitResult); + blocking.join(); + releaseLater.join(); + } }