From 973e0ca0341b6d893e720237ff5e1f8ad7502804 Mon Sep 17 00:00:00 2001 From: Danny McCormick Date: Wed, 23 Sep 2026 19:46:50 +0000 Subject: [PATCH 1/2] Time out RestrictionTrackers.trySplit after 5 minutes on lock contention --- .../splittabledofn/RestrictionTrackers.java | 28 ++++++++++++++----- .../RestrictionTrackersTest.java | 27 +++++++++++++++++- 2 files changed, 47 insertions(+), 8 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 370871fd9cae..c91e36b46d02 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,27 @@ public RestrictionT currentRestriction() { @Override public SplitResult trySplit(double fractionOfRemainder) { - lock.lock(); + return trySplit(fractionOfRemainder, SPLIT_TIMEOUT_SEC); + } + + @VisibleForTesting + SplitResult trySplit(double fractionOfRemainder, int timeOutSec) { try { - SplitResult result = delegate.trySplit(fractionOfRemainder); - needsProgressUpdate = true; - return result; - } finally { - updateProgressAndUnlock(); + // lock can be held long by 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..0b750badab3b 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,29 @@ 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); + } } From a8ca2f67dd94b4031579aa1bdf0d98577a67c11c Mon Sep 17 00:00:00 2001 From: Danny McCormick Date: Thu, 24 Sep 2026 13:50:46 +0000 Subject: [PATCH 2/2] Block on lock.lock() when fractionOfRemainder == 0 in trySplit --- .../splittabledofn/RestrictionTrackers.java | 19 ++++++++++-- .../RestrictionTrackersTest.java | 30 +++++++++++++++++++ 2 files changed, 47 insertions(+), 2 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 c91e36b46d02..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 @@ -104,9 +104,24 @@ public SplitResult trySplit(double fractionOfRemainder) { @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 { - // lock can be held long by long-running tryClaim. We tolerate this scenario by returning - // null (declining to split) when lock timeout occurs. + // 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); 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 0b750badab3b..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 @@ -287,4 +287,34 @@ public void testClaimObserversTrySplitNonBlockingOnTryClaim() throws Interrupted ((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(); + } }