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 6fefc6b184a5..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,11 +45,13 @@ 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 boolean hasInitialProgress = false; + protected volatile Progress lastProgress = Progress.NONE; + protected volatile boolean needsProgressUpdate = false; private final ClaimObserver 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 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); - 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 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 extends RestrictionTrackerObserver implements HasProgress { - private static final int FIRST_PROGRESS_TIMEOUT_SEC = 60; + private static final int PROGRESS_TIMEOUT_SEC = 60; protected RestrictionTrackerObserverWithProgress( RestrictionTracker 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; } } 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 7b6f3d47c273..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 @@ -103,7 +103,9 @@ private static class RestrictionTrackerWithProgress extends RestrictionTracker 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,13 +204,32 @@ 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 @@ -204,15 +237,84 @@ public void testClaimObserversProgressNonBlockingOnTrySplit() throws Interrupted RestrictionTrackerWithProgress withProgress = new RestrictionTrackerWithProgress(false, true); RestrictionTracker 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 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(); } }