Allow close exchange early

Author: dnhatn

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/exchange/ExchangeService.java

@@ -14,6 +14,7 @@
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.ActionListenerResponseHandler;
 import org.elasticsearch.action.support.ChannelActionListener;
+import org.elasticsearch.action.support.SubscribableListener;
 import org.elasticsearch.common.component.AbstractLifecycleComponent;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
@@ -23,6 +24,7 @@
 import org.elasticsearch.common.util.concurrent.EsRejectedExecutionException;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.BlockStreamInput;
+import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.tasks.CancellableTask;
@@ -40,10 +42,11 @@
 
 import java.io.IOException;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.Executor;
-import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.atomic.AtomicReference;
 
 /**
  * {@link ExchangeService} is responsible for exchanging pages between exchange sinks and sources on the same or different nodes.
@@ -293,7 +296,7 @@ static final class TransportRemoteSink implements RemoteSink {
         final Executor responseExecutor;
 
         final AtomicLong estimatedPageSizeInBytes = new AtomicLong(0L);
-        final AtomicBoolean finished = new AtomicBoolean(false);
+        final AtomicReference<SubscribableListener<Void>> completionListenerRef = new AtomicReference<>(null);
 
         TransportRemoteSink(
             TransportService transportService,
@@ -318,13 +321,14 @@ public void fetchPageAsync(boolean allSourcesFinished, ActionListener<ExchangeRe
                 return;
             }
             // already finished
-            if (finished.get()) {
-                listener.onResponse(new ExchangeResponse(blockFactory, null, true));
+            SubscribableListener<Void> completionListener = completionListenerRef.get();
+            if (completionListener != null) {
+                completionListener.addListener(listener.map(unused -> new ExchangeResponse(blockFactory, null, true)));
                 return;
             }
             doFetchPageAsync(false, ActionListener.wrap(r -> {
                 if (r.finished()) {
-                    finished.set(true);
+                    completionListenerRef.compareAndSet(null, SubscribableListener.newSucceeded(null));
                 }
                 listener.onResponse(r);
             }, e -> close(ActionListener.running(() -> listener.onFailure(e)))));
@@ -356,10 +360,19 @@ private void doFetchPageAsync(boolean allSourcesFinished, ActionListener<Exchang
 
         @Override
         public void close(ActionListener<Void> listener) {
-            if (finished.compareAndSet(false, true)) {
-                doFetchPageAsync(true, listener.delegateFailure((l, unused) -> l.onResponse(null)));
-            } else {
-                listener.onResponse(null);
+            final SubscribableListener<Void> candidate = new SubscribableListener<>();
+            final SubscribableListener<Void> actual = completionListenerRef.updateAndGet(
+                curr -> Objects.requireNonNullElse(curr, candidate)
+            );
+            actual.addListener(listener);
+            if (candidate == actual) {
+                doFetchPageAsync(true, candidate.delegateFailure((l, r) -> {
+                    final Page page = r.takePage();
+                    if (page != null) {
+                        page.releaseBlocks();
+                    }
+                    l.onResponse(null);
+                }));
             }
         }
     }

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/exchange/ExchangeSourceHandler.java

@@ -13,12 +13,14 @@
 import org.elasticsearch.action.support.SubscribableListener;
 import org.elasticsearch.common.util.concurrent.AbstractRunnable;
 import org.elasticsearch.compute.EsqlRefCountingListener;
+import org.elasticsearch.common.util.concurrent.ConcurrentCollections;
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.FailureCollector;
 import org.elasticsearch.compute.operator.IsBlockedResult;
 import org.elasticsearch.core.Releasable;
 
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.Executor;
 import java.util.concurrent.atomic.AtomicInteger;
 
@@ -41,6 +43,9 @@ public final class ExchangeSourceHandler {
     // The final failure collected will be notified to callers via the {@code completionListener}.
     private final FailureCollector failure = new FailureCollector();
 
+    private final AtomicInteger nextSinkId = new AtomicInteger();
+    private final Map<Integer, RemoteSink> remoteSinks = ConcurrentCollections.newConcurrentMap();
+
     /**
      * Creates a new ExchangeSourceHandler.
      *
@@ -53,7 +58,12 @@ public ExchangeSourceHandler(int maxBufferSize, Executor fetchExecutor, ActionLi
         this.buffer = new ExchangeBuffer(maxBufferSize);
         this.fetchExecutor = fetchExecutor;
         this.outstandingSinks = new PendingInstances(() -> buffer.finish(false));
-        this.outstandingSources = new PendingInstances(() -> buffer.finish(true));
+        final PendingInstances closingSinks = new PendingInstances(() -> {});
+        closingSinks.trackNewInstance();
+        this.outstandingSources = new PendingInstances(() -> {
+            buffer.finish(true);
+            finishEarly(ActionListener.running(closingSinks::finishInstance));
+        });
         buffer.addCompletionListener(ActionListener.running(() -> {
             final ActionListener<Void> listener = ActionListener.assertAtLeastOnce(completionListener);
             try (RefCountingRunnable refs = new RefCountingRunnable(() -> {
@@ -257,7 +267,11 @@ void onSinkComplete() {
      * @see ExchangeSinkHandler#fetchPageAsync(boolean, ActionListener)
      */
     public void addRemoteSink(RemoteSink remoteSink, boolean failFast, int instances, ActionListener<Void> listener) {
-        final ActionListener<Void> sinkListener = ActionListener.assertAtLeastOnce(ActionListener.notifyOnce(listener));
+        final int sinkId = nextSinkId.incrementAndGet();
+        remoteSinks.put(sinkId, remoteSink);
+        final ActionListener<Void> sinkListener = ActionListener.assertAtLeastOnce(
+            ActionListener.notifyOnce(ActionListener.runBefore(listener, () -> remoteSinks.remove(sinkId)))
+        );
         fetchExecutor.execute(new AbstractRunnable() {
             @Override
             public void onFailure(Exception e) {
@@ -291,6 +305,19 @@ public Releasable addEmptySink() {
         return outstandingSinks::finishInstance;
     }
 
+    /**
+     * Gracefully terminates the exchange source early by instructing all remote exchange sinks to stop their computations.
+     * This can happen when the exchange source has accumulated enough data (e.g., reaching the LIMIT) or when users want to
+     * see the current result immediately.
+     */
+    public void finishEarly(ActionListener<Void> listener) {
+        try (RefCountingListener refs = new RefCountingListener(listener)) {
+            for (RemoteSink remoteSink : remoteSinks.values()) {
+                remoteSink.close(refs.acquire());
+            }
+        }
+    }
+
     private static class PendingInstances {
         private final AtomicInteger instances = new AtomicInteger();
         private final SubscribableListener<Void> completion = new SubscribableListener<>();

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/exchange/RemoteSink.java

@@ -8,18 +8,19 @@
 package org.elasticsearch.compute.operator.exchange;
 
 import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.compute.data.Page;
 
 public interface RemoteSink {
 
     void fetchPageAsync(boolean allSourcesFinished, ActionListener<ExchangeResponse> listener);
 
     default void close(ActionListener<Void> listener) {
         fetchPageAsync(true, listener.delegateFailure((l, r) -> {
-            try {
-                r.close();
-            } finally {
-                l.onResponse(null);
+            final Page page = r.takePage();
+            if (page != null) {
+                page.releaseBlocks();
             }
+            l.onResponse(null);
         }));
     }
 }

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/exchange/ExchangeServiceTests.java

@@ -55,7 +55,9 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.Queue;
 import java.util.Set;
+import java.util.concurrent.Semaphore;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -421,7 +423,7 @@ public void testExchangeSourceContinueOnFailure() {
         }
     }
 
-    public void testEarlyTerminate() {
+    public void testClosingSinks() {
         BlockFactory blockFactory = blockFactory();
         IntBlock block1 = blockFactory.newConstantIntBlockWith(1, 2);
         IntBlock block2 = blockFactory.newConstantIntBlockWith(1, 2);
@@ -441,6 +443,55 @@ public void testEarlyTerminate() {
         assertTrue(sink.isFinished());
     }
 
+    public void testFinishEarly() throws Exception {
+        ExchangeSourceHandler sourceHandler = new ExchangeSourceHandler(20, threadPool.generic(), ActionListener.noop());
+        Semaphore permits = new Semaphore(between(1, 5));
+        BlockFactory blockFactory = blockFactory();
+        Queue<Page> pages = ConcurrentCollections.newQueue();
+        ExchangeSource exchangeSource = sourceHandler.createExchangeSource();
+        AtomicBoolean sinkClosed = new AtomicBoolean();
+        PlainActionFuture<Void> sinkCompleted = new PlainActionFuture<>();
+        sourceHandler.addRemoteSink((allSourcesFinished, listener) -> {
+            if (allSourcesFinished) {
+                sinkClosed.set(true);
+                permits.release(10);
+                listener.onResponse(new ExchangeResponse(blockFactory, null, sinkClosed.get()));
+            } else {
+                try {
+                    if (permits.tryAcquire(between(0, 100), TimeUnit.MICROSECONDS)) {
+                        boolean closed = sinkClosed.get();
+                        final Page page;
+                        if (closed) {
+                            page = new Page(blockFactory.newConstantIntBlockWith(1, 1));
+                            pages.add(page);
+                        } else {
+                            page = null;
+                        }
+                        listener.onResponse(new ExchangeResponse(blockFactory, page, closed));
+                    } else {
+                        listener.onResponse(new ExchangeResponse(blockFactory, null, sinkClosed.get()));
+                    }
+                } catch (Exception e) {
+                    throw new AssertionError(e);
+                }
+            }
+        }, false, between(1, 3), sinkCompleted);
+        threadPool.schedule(
+            () -> sourceHandler.finishEarly(ActionListener.noop()),
+            TimeValue.timeValueMillis(between(0, 10)),
+            threadPool.generic()
+        );
+        sinkCompleted.actionGet();
+        Page p;
+        while ((p = exchangeSource.pollPage()) != null) {
+            assertSame(p, pages.poll());
+            p.releaseBlocks();
+        }
+        assertTrue(exchangeSource.isFinished());
+        assertNull(pages.poll());
+        exchangeSource.finish();
+    }
+
     public void testConcurrentWithTransportActions() {
         MockTransportService node0 = newTransportService();
         ExchangeService exchange0 = new ExchangeService(Settings.EMPTY, threadPool, ESQL_TEST_EXECUTOR, blockFactory());