From c20877fba59bf0728397ca3766d748fc14288c07 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Tue, 18 Aug 2026 16:59:32 +0800 Subject: [PATCH] fix(processing): finalize SuperSorter progress before completion --- .../druid/frame/processor/SuperSorter.java | 27 ++++++++++-- .../frame/processor/SuperSorterTest.java | 41 ++++++++++++++++++- 2 files changed, 62 insertions(+), 6 deletions(-) diff --git a/processing/src/main/java/org/apache/druid/frame/processor/SuperSorter.java b/processing/src/main/java/org/apache/druid/frame/processor/SuperSorter.java index 7253feb06aad..b669ecb35607 100644 --- a/processing/src/main/java/org/apache/druid/frame/processor/SuperSorter.java +++ b/processing/src/main/java/org/apache/druid/frame/processor/SuperSorter.java @@ -183,6 +183,9 @@ public class SuperSorter @GuardedBy("runWorkersLock") private SettableFuture allDone = null; + @GuardedBy("runWorkersLock") + private boolean totalMergersForUltimateLevelSet = false; + @GuardedBy("runWorkersLock") SuperSorterProgressTracker superSorterProgressTracker; @@ -299,7 +302,7 @@ public ListenableFuture run() () -> { synchronized (runWorkersLock) { if (outputPartitionsFuture.isDone()) { // Update the progress tracker - superSorterProgressTracker.setTotalMergersForUltimateLevel(getOutputPartitions().size()); + setTotalMergersForUltimateLevel(); } runWorkersIfPossible(); setAllDoneIfPossible(); @@ -415,7 +418,7 @@ private void setAllDoneIfPossible() } // OK to use wrap, not wrapReadOnly, because nil channels are already read-only. - allDone.set(OutputChannels.wrap(channels)); + setAllDone(OutputChannels.wrap(channels)); } else if (rowLimit == 0 && activeProcessors == 0) { // We had a row limit, and got it all the way down to zero. // Generate empty output channels for any partitions that we haven't written yet. @@ -427,14 +430,14 @@ private void setAllDoneIfPossible() } // OK to use wrap, not wrapReadOnly, because all channels in this list are already read-only. - allDone.set(OutputChannels.wrap(outputChannels)); + setAllDone(OutputChannels.wrap(outputChannels)); } else if (totalMergingLevels != UNKNOWN_LEVEL && outputsReadyByLevel.containsKey(totalMergingLevels - 1) && (outputsReadyByLevel.get(totalMergingLevels - 1).size() == getTotalMergersInLevel(totalMergingLevels - 1))) { // We're done!! // OK to use wrap, not wrapReadOnly, because all channels in this list are already read-only. - allDone.set(OutputChannels.wrap(outputChannels)); + setAllDone(OutputChannels.wrap(outputChannels)); } } catch (Throwable e) { @@ -442,6 +445,22 @@ private void setAllDoneIfPossible() } } + @GuardedBy("runWorkersLock") + private void setAllDone(final OutputChannels channels) + { + setTotalMergersForUltimateLevel(); + allDone.set(channels); + } + + @GuardedBy("runWorkersLock") + private void setTotalMergersForUltimateLevel() + { + if (!totalMergersForUltimateLevelSet) { + superSorterProgressTracker.setTotalMergersForUltimateLevel(getOutputPartitions().size()); + totalMergersForUltimateLevelSet = true; + } + } + @GuardedBy("runWorkersLock") private boolean runNextBatcher() { diff --git a/processing/src/test/java/org/apache/druid/frame/processor/SuperSorterTest.java b/processing/src/test/java/org/apache/druid/frame/processor/SuperSorterTest.java index 309baf36a0c0..7beeb5849710 100644 --- a/processing/src/test/java/org/apache/druid/frame/processor/SuperSorterTest.java +++ b/processing/src/test/java/org/apache/druid/frame/processor/SuperSorterTest.java @@ -74,6 +74,7 @@ import org.junit.jupiter.params.ParameterizedClass; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.MethodSource; +import org.mockito.Mockito; import java.io.File; import java.io.IOException; @@ -85,6 +86,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.concurrent.Executor; import java.util.concurrent.TimeUnit; import java.util.function.Consumer; @@ -119,6 +121,32 @@ public void tearDown() exec.getExecutorService().shutdownNow(); } + private static class ListenerDelayingFrameProcessorExecutor extends FrameProcessorExecutor + { + private int asExecutorCalls; + private Runnable pendingListener; + + private ListenerDelayingFrameProcessorExecutor() + { + super(MoreExecutors.listeningDecorator(Execs.multiThreaded(NUM_THREADS, "super-sorter-test-%d"))); + } + + @Override + public Executor asExecutor(final String cancellationId) + { + // The worker callback is registered first; delay the output-partitions listener registered by run(). + if (++asExecutorCalls == 2) { + return command -> pendingListener = command; + } + return super.asExecutor(cancellationId); + } + + private void runListener() + { + pendingListener.run(); + } + } + @Test public void testSingleEmptyInputChannel_fileStorage() throws Exception { @@ -161,10 +189,16 @@ public void testSingleEmptyInputChannel_fileStorage() throws Exception @Test public void testSingleEmptyInputChannel_immediately_fileStorage() throws Exception { + exec.getExecutorService().shutdownNow(); + final ListenerDelayingFrameProcessorExecutor listenerDelayingExec = + new ListenerDelayingFrameProcessorExecutor(); + exec = listenerDelayingExec; + final BlockingQueueFrameChannel inputChannel = BlockingQueueFrameChannel.minimal(); inputChannel.writable().close(); - final SuperSorterProgressTracker superSorterProgressTracker = new SuperSorterProgressTracker(); + final SuperSorterProgressTracker superSorterProgressTracker = + Mockito.spy(new SuperSorterProgressTracker()); final File tempFolder = temporaryFolder.newFolder(); final SuperSorter superSorter = new SuperSorter( @@ -188,10 +222,13 @@ public void testSingleEmptyInputChannel_immediately_fileStorage() throws Excepti final OutputChannels channels = superSorter.run().get(); Assertions.assertEquals(1, channels.getAllChannels().size()); + Mockito.verify(superSorterProgressTracker).setTotalMergersForUltimateLevel(1L); + Assertions.assertEquals(1.0, superSorterProgressTracker.snapshot().getProgressDigest(), 0.0f); + + listenerDelayingExec.runListener(); final ReadableFrameChannel channel = Iterables.getOnlyElement(channels.getAllChannels()).getReadableChannel(); Assertions.assertTrue(channel.isFinished()); - Assertions.assertEquals(1.0, superSorterProgressTracker.snapshot().getProgressDigest(), 0.0f); channel.close(); }