Keep outstanding pages when finish buffer early (elastic#121857)

Today, the exchange buffer of an exchange source is finished in two
cases: (1) when the downstream pipeline has received enough data and (2)
when all remote sinks have completed. In the first case, outstanding
pages could be safely discarded. In the second case, no new pages should
be received after finishing. In both scenarios, discarding all
outstanding pages was safe if noMoreInputs was switched while adding
pages.

However, with the stop API, the buffer may now finish while keeping
outstanding pages, and new pages may still be received. This change
updates the exchange buffer to discard only the incoming page when
noMoreInputs is switched, rather than all pages in the buffer.

Closes elastic#120757

Author: dnhatn

muted-tests.yml

@@ -420,9 +420,6 @@ tests:
 - class: org.elasticsearch.xpack.test.rest.XPackRestIT
   method: test {p0=transform/transforms_start_stop/Test schedule_now on an already started transform}
   issue: https://github.com/elastic/elasticsearch/issues/120720
-- class: org.elasticsearch.xpack.esql.action.CrossClusterAsyncEnrichStopIT
-  method: testEnrichAfterStop
-  issue: https://github.com/elastic/elasticsearch/issues/120757
 - class: org.elasticsearch.xpack.shutdown.NodeShutdownIT
   method: testStalledShardMigrationProperlyDetected
   issue: https://github.com/elastic/elasticsearch/issues/115697

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

@@ -47,7 +47,17 @@ void addPage(Page page) {
             notifyNotEmpty();
         }
         if (noMoreInputs) {
-            discardPages();
+            // O(N) but acceptable because it only occurs with the stop API, and the queue size should be very small.
+            if (queue.removeIf(p -> p == page)) {
+                page.releaseBlocks();
+                final int size = queueSize.decrementAndGet();
+                if (size == maxSize - 1) {
+                    notifyNotFull();
+                }
+                if (size == 0) {
+                    completionFuture.onResponse(null);
+                }
+            }
         }
     }
 

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

@@ -66,6 +66,25 @@ public void testDrainPages() throws Exception {
         blockFactory.ensureAllBlocksAreReleased();
     }
 
+    public void testOutstandingPages() throws Exception {
+        ExchangeBuffer buffer = new ExchangeBuffer(randomIntBetween(1000, 10000));
+        var blockFactory = blockFactory();
+        Page p1 = randomPage(blockFactory);
+        Page p2 = randomPage(blockFactory);
+        buffer.addPage(p1);
+        buffer.addPage(p2);
+        buffer.finish(false);
+        buffer.addPage(randomPage(blockFactory));
+        assertThat(buffer.size(), equalTo(2));
+        assertSame(buffer.pollPage(), p1);
+        p1.releaseBlocks();
+        assertSame(buffer.pollPage(), p2);
+        p2.releaseBlocks();
+        assertNull(buffer.pollPage());
+        assertTrue(buffer.isFinished());
+        blockFactory.ensureAllBlocksAreReleased();
+    }
+
     private static MockBlockFactory blockFactory() {
         BigArrays bigArrays = new MockBigArrays(PageCacheRecycler.NON_RECYCLING_INSTANCE, ByteSizeValue.ofGb(1)).withCircuitBreaking();
         CircuitBreaker breaker = bigArrays.breakerService().getBreaker(CircuitBreaker.REQUEST);