Fix tests

Author: dnhatn

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

@@ -366,13 +366,13 @@ public void close(ActionListener<Void> listener) {
             );
             actual.addListener(listener);
             if (candidate == actual) {
-                doFetchPageAsync(true, candidate.delegateFailure((l, r) -> {
+                doFetchPageAsync(true, ActionListener.wrap(r -> {
                     final Page page = r.takePage();
                     if (page != null) {
                         page.releaseBlocks();
                     }
-                    l.onResponse(null);
-                }));
+                    candidate.onResponse(null);
+                }, e -> candidate.onResponse(null)));
             }
         }
     }

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

@@ -60,10 +60,7 @@ public ExchangeSourceHandler(int maxBufferSize, Executor fetchExecutor, ActionLi
         this.outstandingSinks = new PendingInstances(() -> buffer.finish(false));
         final PendingInstances closingSinks = new PendingInstances(() -> {});
         closingSinks.trackNewInstance();
-        this.outstandingSources = new PendingInstances(() -> {
-            buffer.finish(true);
-            finishEarly(ActionListener.running(closingSinks::finishInstance));
-        });
+        this.outstandingSources = new PendingInstances(() -> finishEarly(true, ActionListener.running(closingSinks::finishInstance)));
         buffer.addCompletionListener(ActionListener.running(() -> {
             final ActionListener<Void> listener = ActionListener.assertAtLeastOnce(completionListener);
             try (RefCountingRunnable refs = new RefCountingRunnable(() -> {
@@ -309,9 +306,12 @@ public Releasable addEmptySink() {
      * 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.
+     *
+     * @param drainingPages whether to discard pages already fetched in the exchange
      */
-    public void finishEarly(ActionListener<Void> listener) {
-        try (RefCountingListener refs = new RefCountingListener(listener)) {
+    public void finishEarly(boolean drainingPages, ActionListener<Void> listener) {
+        buffer.finish(drainingPages);
+        try (EsqlRefCountingListener refs = new EsqlRefCountingListener(listener)) {
             for (RemoteSink remoteSink : remoteSinks.values()) {
                 remoteSink.close(refs.acquire());
             }

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

@@ -477,7 +477,7 @@ public void testFinishEarly() throws Exception {
             }
         }, false, between(1, 3), sinkCompleted);
         threadPool.schedule(
-            () -> sourceHandler.finishEarly(ActionListener.noop()),
+            () -> sourceHandler.finishEarly(randomBoolean(), ActionListener.noop()),
             TimeValue.timeValueMillis(between(0, 10)),
             threadPool.generic()
         );
@@ -487,8 +487,10 @@ public void testFinishEarly() throws Exception {
             assertSame(p, pages.poll());
             p.releaseBlocks();
         }
+        while ((p = pages.poll()) != null) {
+            p.releaseBlocks();
+        }
         assertTrue(exchangeSource.isFinished());
-        assertNull(pages.poll());
         exchangeSource.finish();
     }