fix(processing): finalize SuperSorter progress before completion

processing/src/main/java/org/apache/druid/frame/processor/SuperSorter.java

@@ -183,6 +183,9 @@ public class SuperSorter
   @GuardedBy("runWorkersLock")
   private SettableFuture<OutputChannels> allDone = null;
 
+  @GuardedBy("runWorkersLock")
+  private boolean totalMergersForUltimateLevelSet = false;
+
   @GuardedBy("runWorkersLock")
   SuperSorterProgressTracker superSorterProgressTracker;
 
@@ -299,7 +302,7 @@ public ListenableFuture<OutputChannels> 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,21 +430,37 @@ 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) {
       allDone.setException(e);
     }
   }
 
+  @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()
   {

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();
     }