From 3cd2f0b1e29695ed08a7eb37f370e124fc14d39b Mon Sep 17 00:00:00 2001 From: Danny McCormick Date: Mon, 21 Sep 2026 16:51:39 +0000 Subject: [PATCH 1/3] Return last evaluated progress on lock timeout in RestrictionTrackers.getProgress --- .../splittabledofn/RestrictionTrackers.java | 46 ++++++------------ .../RestrictionTrackersTest.java | 48 ++++++++++++++++++- 2 files changed, 62 insertions(+), 32 deletions(-) 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..73f92d065a43 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 @@ -49,7 +49,6 @@ private static class RestrictionTrackerObserver extends RestrictionTracker { protected final RestrictionTracker delegate; protected ReentrantLock lock = new ReentrantLock(); - protected volatile boolean hasInitialProgress = false; private final ClaimObserver claimObserver; protected RestrictionTrackerObserver( @@ -110,16 +109,6 @@ public void checkDone() throws IllegalStateException { public IsBounded isBounded() { return delegate.isBounded(); } - - /** Evaluate progress if requested. */ - protected Progress getProgressBlocking() { - lock.lock(); - try { - return ((HasProgress) delegate).getProgress(); - } finally { - lock.unlock(); - } - } } /** @@ -129,7 +118,8 @@ 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; + private volatile Progress lastProgress = Progress.NONE; protected RestrictionTrackerObserverWithProgress( RestrictionTracker delegate, @@ -139,32 +129,26 @@ 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 { + lastProgress = ((HasProgress) delegate).getProgress(); + } finally { + lock.unlock(); } - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); } - return progress; - } else { - return getProgressBlocking(); + } catch (InterruptedException e) { + 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..3acce32703d4 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 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 @@ -212,7 +242,23 @@ public void testClaimObserversProgressNonBlockingOnTrySplit() throws Interrupted assertEquals(RestrictionTracker.Progress.NONE, progress); withProgress.releaseLock(); withProgress.waitUntilBlocking(false); + blocking.join(); progress = ((HasProgress) tracker).getProgress(); assertEquals(RestrictionTrackerWithProgress.REPORT_PROGRESS, progress); + + // Subsequent blocking trySplit should return lastProgress rather than blocking or Progress.NONE + withProgress.setProgress(RestrictionTrackerWithProgress.UPDATED_PROGRESS); + withProgress.setBlockTrySplit(true); + Thread secondBlocking = new Thread(() -> tracker.trySplit(0.5)); + 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); } } From 44dffcd0fa960ef5e685572a3e24b0064a6d7dac Mon Sep 17 00:00:00 2001 From: Danny McCormick Date: Mon, 21 Sep 2026 17:19:01 +0000 Subject: [PATCH 2/3] Update progress before unlocking when needsProgressUpdate is set --- .../splittabledofn/RestrictionTrackers.java | 37 ++++++++++++++++--- .../RestrictionTrackersTest.java | 16 +++++--- 2 files changed, 41 insertions(+), 12 deletions(-) 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 73f92d065a43..9093eca89d87 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 @@ -49,6 +49,8 @@ private static class RestrictionTrackerObserver extends RestrictionTracker { protected final RestrictionTracker delegate; protected ReentrantLock lock = new ReentrantLock(); + protected volatile Progress lastProgress = Progress.NONE; + protected volatile boolean needsProgressUpdate = false; private final ClaimObserver claimObserver; protected RestrictionTrackerObserver( @@ -58,6 +60,16 @@ protected RestrictionTrackerObserver( this.claimObserver = claimObserver; } + protected void unlock() { + try { + if (needsProgressUpdate) { + updateProgressBlocking(); + } + } finally { + lock.unlock(); + } + } + @Override public boolean tryClaim(PositionT position) { lock.lock(); @@ -70,7 +82,7 @@ public boolean tryClaim(PositionT position) { return false; } } finally { - lock.unlock(); + unlock(); } } @@ -80,7 +92,7 @@ public RestrictionT currentRestriction() { try { return delegate.currentRestriction(); } finally { - lock.unlock(); + unlock(); } } @@ -91,7 +103,7 @@ public SplitResult trySplit(double fractionOfRemainder) { SplitResult result = delegate.trySplit(fractionOfRemainder); return result; } finally { - lock.unlock(); + unlock(); } } @@ -101,7 +113,7 @@ public void checkDone() throws IllegalStateException { try { delegate.checkDone(); } finally { - lock.unlock(); + unlock(); } } @@ -109,6 +121,17 @@ public void checkDone() throws IllegalStateException { public IsBounded isBounded() { return delegate.isBounded(); } + + /** Evaluate progress if requested. */ + protected void updateProgressBlocking() { + lock.lock(); + try { + needsProgressUpdate = false; + lastProgress = ((HasProgress) delegate).getProgress(); + } finally { + lock.unlock(); + } + } } /** @@ -119,7 +142,6 @@ public IsBounded isBounded() { static class RestrictionTrackerObserverWithProgress extends RestrictionTrackerObserver implements HasProgress { private static final int PROGRESS_TIMEOUT_SEC = 60; - private volatile Progress lastProgress = Progress.NONE; protected RestrictionTrackerObserverWithProgress( RestrictionTracker delegate, @@ -140,12 +162,15 @@ Progress getProgress(int timeOutSec) { // when lock timeout occurs. if (lock.tryLock(timeOutSec, TimeUnit.SECONDS)) { try { - lastProgress = ((HasProgress) delegate).getProgress(); + updateProgressBlocking(); } finally { lock.unlock(); } + } else { + needsProgressUpdate = true; } } 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 3acce32703d4..66c28d89ac14 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 @@ -204,16 +204,18 @@ 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, unlock() sees needsProgressUpdate == true and evaluates + // REPORT_PROGRESS before releasing the lock. withProgress.releaseLock(); withProgress.waitUntilBlocking(false); blocking.join(); - progress = ((HasProgress) tracker).getProgress(); - assertEquals(RestrictionTrackerWithProgress.REPORT_PROGRESS, progress); - // Subsequent blocking tryClaim should return lastProgress rather than blocking or Progress.NONE + // 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 unlock). withProgress.setProgress(RestrictionTrackerWithProgress.UPDATED_PROGRESS); withProgress.setBlockTryClaim(true); Thread secondBlocking = new Thread(() -> tracker.tryClaim(new Object())); @@ -237,16 +239,18 @@ public void testClaimObserversProgressNonBlockingOnTrySplit() throws Interrupted Thread blocking = new Thread(() -> tracker.trySplit(0.5)); blocking.start(); withProgress.waitUntilBlocking(true); + // Times out while first trySplit holds lock; returns NONE and sets needsProgressUpdate = true RestrictionTracker.Progress progress = ((RestrictionTrackers.RestrictionTrackerObserverWithProgress) tracker).getProgress(1); assertEquals(RestrictionTracker.Progress.NONE, progress); + // When first trySplit finishes, unlock() sees needsProgressUpdate == true and evaluates + // REPORT_PROGRESS before releasing the lock. withProgress.releaseLock(); withProgress.waitUntilBlocking(false); blocking.join(); - progress = ((HasProgress) tracker).getProgress(); - assertEquals(RestrictionTrackerWithProgress.REPORT_PROGRESS, progress); - // Subsequent blocking trySplit should return lastProgress rather than blocking or Progress.NONE + // Even if a second blocking trySplit immediately grabs the lock before getProgress is called + // again, getProgress(1) returns REPORT_PROGRESS (updated during first trySplit's unlock). withProgress.setProgress(RestrictionTrackerWithProgress.UPDATED_PROGRESS); withProgress.setBlockTrySplit(true); Thread secondBlocking = new Thread(() -> tracker.trySplit(0.5)); From 409d7a289e8bb04b43bf7bab4b1708b71ada951f Mon Sep 17 00:00:00 2001 From: Danny McCormick Date: Tue, 22 Sep 2026 20:31:20 +0000 Subject: [PATCH 3/3] Unconditionally update progress after trySplit and rename unlock to updateProgressAndUnlock --- .../splittabledofn/RestrictionTrackers.java | 15 ++++++----- .../RestrictionTrackersTest.java | 25 ++++++++----------- 2 files changed, 20 insertions(+), 20 deletions(-) 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 9093eca89d87..370871fd9cae 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 @@ -60,7 +60,7 @@ protected RestrictionTrackerObserver( this.claimObserver = claimObserver; } - protected void unlock() { + protected void updateProgressAndUnlock() { try { if (needsProgressUpdate) { updateProgressBlocking(); @@ -82,7 +82,7 @@ public boolean tryClaim(PositionT position) { return false; } } finally { - unlock(); + updateProgressAndUnlock(); } } @@ -92,7 +92,7 @@ public RestrictionT currentRestriction() { try { return delegate.currentRestriction(); } finally { - unlock(); + updateProgressAndUnlock(); } } @@ -101,9 +101,10 @@ public SplitResult trySplit(double fractionOfRemainder) { lock.lock(); try { SplitResult result = delegate.trySplit(fractionOfRemainder); + needsProgressUpdate = true; return result; } finally { - unlock(); + updateProgressAndUnlock(); } } @@ -113,7 +114,7 @@ public void checkDone() throws IllegalStateException { try { delegate.checkDone(); } finally { - unlock(); + updateProgressAndUnlock(); } } @@ -127,7 +128,9 @@ protected void updateProgressBlocking() { lock.lock(); try { needsProgressUpdate = false; - lastProgress = ((HasProgress) delegate).getProgress(); + if (delegate instanceof HasProgress) { + lastProgress = ((HasProgress) delegate).getProgress(); + } } finally { lock.unlock(); } 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 66c28d89ac14..e24033224782 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 @@ -208,14 +208,15 @@ public void testClaimObserversProgressNonBlockingOnTryClaim() throws Interrupted RestrictionTracker.Progress progress = ((RestrictionTrackers.RestrictionTrackerObserverWithProgress) tracker).getProgress(1); assertEquals(RestrictionTracker.Progress.NONE, progress); - // When first tryClaim finishes, unlock() sees needsProgressUpdate == true and evaluates - // REPORT_PROGRESS before releasing the lock. + // When first tryClaim finishes, updateProgressAndUnlock() sees needsProgressUpdate == true and + // evaluates REPORT_PROGRESS before releasing the lock. withProgress.releaseLock(); withProgress.waitUntilBlocking(false); 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 unlock). + // 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())); @@ -236,27 +237,23 @@ 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); - // Times out while first trySplit holds lock; returns NONE and sets needsProgressUpdate = true - RestrictionTracker.Progress progress = - ((RestrictionTrackers.RestrictionTrackerObserverWithProgress) tracker).getProgress(1); - assertEquals(RestrictionTracker.Progress.NONE, progress); - // When first trySplit finishes, unlock() sees needsProgressUpdate == true and evaluates - // REPORT_PROGRESS before releasing the lock. withProgress.releaseLock(); withProgress.waitUntilBlocking(false); blocking.join(); - // Even if a second blocking trySplit immediately grabs the lock before getProgress is called - // again, getProgress(1) returns REPORT_PROGRESS (updated during first trySplit's unlock). + // 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.setBlockTrySplit(true); - Thread secondBlocking = new Thread(() -> tracker.trySplit(0.5)); + withProgress.setBlockTryClaim(true); + Thread secondBlocking = new Thread(() -> tracker.tryClaim(new Object())); secondBlocking.start(); withProgress.waitUntilBlocking(true); - progress = + RestrictionTracker.Progress progress = ((RestrictionTrackers.RestrictionTrackerObserverWithProgress) tracker).getProgress(1); assertEquals(RestrictionTrackerWithProgress.REPORT_PROGRESS, progress); withProgress.releaseLock();