Fix OrderByOperator.finish to wait for finishMemoryRevoke

core/trino-main/src/main/java/io/trino/operator/Operator.java

@@ -62,9 +62,12 @@ default ListenableFuture<Void> isBlocked()
      * Since memory revoking signal is delivered asynchronously to the Operator, implementation
      * must gracefully handle the case when there no longer is any revocable memory allocated.
      * <p>
-     * After this method is called on Operator the Driver is disallowed to call any
-     * processing methods on it (isBlocked/needsInput/addInput/getOutput) until
-     * {@link #finishMemoryRevoke()} is called.
+     * After this method is called on Operator the Driver is disallowed to call most of
+     * processing methods on it
+     * ({@link #isBlocked()}/{@link #needsInput()}/{@link #addInput(Page)}/{@link #getOutput()})
+     * until {@link #finishMemoryRevoke()} is called. {@link #finish()} is the only processing
+     * method that can be called during that time and {@link #close()} remains callable
+     * at any time.
      */
     default ListenableFuture<Void> startMemoryRevoke()
     {

core/trino-main/src/main/java/io/trino/operator/OrderByOperator.java

@@ -157,7 +157,7 @@ private enum State
 
     private Optional<Spiller> spiller = Optional.empty();
     private ListenableFuture<Void> spillInProgress = immediateVoidFuture();
-    private Runnable finishMemoryRevoke = () -> {};
+    private Optional<Runnable> finishMemoryRevoke = Optional.empty();
 
     private Iterator<Optional<Page>> sortedPages;
 
@@ -205,6 +205,9 @@ public void finish()
             return;
         }
         checkSuccess(spillInProgress, "spilling failed");
+        if (finishMemoryRevoke.isPresent()) {
+            return;
+        }
 
         if (state == State.NEEDS_INPUT) {
             state = State.HAS_OUTPUT;
@@ -220,7 +223,8 @@ public void finish()
                     // spill since revocable memory could not be converted to user memory immediately
                     // TODO: this should be asynchronous
                     getFutureValue(spillToDisk());
-                    finishMemoryRevoke.run();
+                    finishMemoryRevoke.orElseThrow().run();
+                    finishMemoryRevoke = Optional.empty();
                 }
             }
 
@@ -302,7 +306,7 @@ private ListenableFuture<Void> spillToDisk()
 
         if (revocableMemoryContext.getBytes() == 0) {
             verify(pageIndex.getPositionCount() == 0 || state == State.HAS_OUTPUT);
-            finishMemoryRevoke = () -> {};
+            finishMemoryRevoke = Optional.of(() -> {});
             return immediateVoidFuture();
         }
 
@@ -317,19 +321,19 @@ private ListenableFuture<Void> spillToDisk()
 
         pageIndex.sort(sortChannels, sortOrder);
         spillInProgress = spiller.get().spill(pageIndex.getSortedPages());
-        finishMemoryRevoke = () -> {
+        finishMemoryRevoke = Optional.of(() -> {
             pageIndex.clear();
             updateMemoryUsage();
-        };
+        });
 
         return spillInProgress;
     }
 
     @Override
     public void finishMemoryRevoke()
     {
-        finishMemoryRevoke.run();
-        finishMemoryRevoke = () -> {};
+        finishMemoryRevoke.orElseThrow().run();
+        finishMemoryRevoke = Optional.empty();
     }
 
     private List<WorkProcessor<Page>> getSpilledPages()