Make MutableSearchResponse ref-counted to prevent use-after-close in async search (#134359)

Author: drempapis

docs/changelog/134359.yaml

@@ -0,0 +1,6 @@
+pr: 134359
+summary: Make `MutableSearchResponse` ref-counted to prevent use-after-close in async
+  search
+area: Search
+type: bug
+issues: []

x-pack/plugin/async-search/src/internalClusterTest/java/org/elasticsearch/xpack/search/AsyncSearchConcurrentStatusIT.java

@@ -0,0 +1,271 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.search;
+
+import org.elasticsearch.ElasticsearchStatusException;
+import org.elasticsearch.ExceptionsHelper;
+import org.elasticsearch.action.index.IndexRequestBuilder;
+import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.core.TimeValue;
+import org.elasticsearch.rest.RestStatus;
+import org.elasticsearch.search.aggregations.AggregationBuilders;
+import org.elasticsearch.search.builder.SearchSourceBuilder;
+import org.elasticsearch.test.ESIntegTestCase.SuiteScopeTestCase;
+import org.elasticsearch.xpack.core.search.action.AsyncSearchResponse;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Queue;
+import java.util.Set;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.concurrent.atomic.LongAdder;
+import java.util.stream.IntStream;
+
+@SuiteScopeTestCase
+public class AsyncSearchConcurrentStatusIT extends AsyncSearchIntegTestCase {
+    private static String indexName;
+    private static int numShards;
+
+    private static int numKeywords;
+
+    @Override
+    public void setupSuiteScopeCluster() {
+        indexName = "test-async";
+        numShards = randomIntBetween(1, 20);
+        int numDocs = randomIntBetween(100, 1000);
+        createIndex(indexName, Settings.builder().put("index.number_of_shards", numShards).build());
+        numKeywords = randomIntBetween(50, 100);
+        Set<String> keywordSet = new HashSet<>();
+        for (int i = 0; i < numKeywords; i++) {
+            keywordSet.add(randomAlphaOfLengthBetween(10, 20));
+        }
+        numKeywords = keywordSet.size();
+        String[] keywords = keywordSet.toArray(String[]::new);
+        List<IndexRequestBuilder> reqs = new ArrayList<>();
+        for (int i = 0; i < numDocs; i++) {
+            float metric = randomFloat();
+            String keyword = keywords[randomIntBetween(0, numKeywords - 1)];
+            reqs.add(prepareIndex(indexName).setSource("terms", keyword, "metric", metric));
+        }
+        indexRandom(true, true, reqs);
+    }
+
+    /**
+     * This test spins up a set of poller threads that repeatedly call
+     * {@code _async_search/{id}}. Each poller starts immediately, and once enough
+     * requests have been issued they signal a latch to indicate the group is "warmed up".
+     * The test waits on this latch to deterministically ensure pollers are active.
+     * In parallel, a consumer thread drives the async search to completion using the
+     * blocking iterator. This coordinated overlap exercises the window where the task
+     * is closing and some status calls may return {@code 410 GONE}.
+     */
+    public void testConcurrentStatusFetchWhileTaskCloses() throws Exception {
+        final TimeValue timeout = TimeValue.timeValueSeconds(3);
+        final String aggName = "terms";
+        final SearchSourceBuilder source = new SearchSourceBuilder().aggregation(
+            AggregationBuilders.terms(aggName).field("terms.keyword").size(Math.max(1, numKeywords))
+        );
+
+        final int progressStep = (numShards > 2) ? randomIntBetween(2, numShards) : 2;
+        try (SearchResponseIterator it = assertBlockingIterator(indexName, numShards, source, 0, progressStep)) {
+            String id = getAsyncId(it);
+
+            PollStats stats = new PollStats();
+
+            // Pick a random number of status-poller threads, at least 1, up to (4×numShards)
+            int pollerThreads = randomIntBetween(1, 4 * numShards);
+
+            // Wait for pollers to be active
+            CountDownLatch warmed = new CountDownLatch(1);
+
+            // Executor and coordination for pollers
+            ExecutorService pollerExec = Executors.newFixedThreadPool(pollerThreads);
+            AtomicBoolean running = new AtomicBoolean(true);
+            Queue<Throwable> failures = new ConcurrentLinkedQueue<>();
+
+            CompletableFuture<Void> pollers = createPollers(id, pollerThreads, stats, warmed, pollerExec, running, failures);
+
+            // Wait until pollers are issuing requests (warming period)
+            assertTrue("pollers did not warm up in time", warmed.await(timeout.millis(), TimeUnit.MILLISECONDS));
+
+            // Start consumer on a separate thread and capture errors
+            var consumerExec = Executors.newSingleThreadExecutor();
+            AtomicReference<Throwable> consumerError = new AtomicReference<>();
+            Future<?> consumer = consumerExec.submit(() -> {
+                try {
+                    consumeAllResponses(it, aggName);
+                } catch (Throwable t) {
+                    consumerError.set(t);
+                }
+            });
+
+            // Join consumer & surface errors
+            try {
+                consumer.get(timeout.millis(), TimeUnit.MILLISECONDS);
+
+                if (consumerError.get() != null) {
+                    fail("consumeAllResponses failed: " + consumerError.get());
+                }
+            } catch (TimeoutException e) {
+                consumer.cancel(true);
+                fail(e, "Consumer thread did not finish within timeout");
+            } catch (Exception ignored) {
+                // ignored
+            } finally {
+                // Stop pollers
+                running.set(false);
+                try {
+                    pollers.get(timeout.millis(), TimeUnit.MILLISECONDS);
+                } catch (TimeoutException te) {
+                    // The finally block will shut down the pollers forcibly
+                } catch (ExecutionException ee) {
+                    failures.add(ExceptionsHelper.unwrapCause(ee.getCause()));
+                } catch (InterruptedException ie) {
+                    Thread.currentThread().interrupt();
+                } finally {
+                    pollerExec.shutdownNow();
+                    try {
+                        pollerExec.awaitTermination(timeout.millis(), TimeUnit.MILLISECONDS);
+                    } catch (InterruptedException ie) {
+                        Thread.currentThread().interrupt();
+                        fail("Interrupted while stopping pollers: " + ie.getMessage());
+                    }
+                }
+
+                // Shut down the consumer executor
+                consumerExec.shutdown();
+                try {
+                    consumerExec.awaitTermination(timeout.millis(), TimeUnit.MILLISECONDS);
+                } catch (InterruptedException ie) {
+                    Thread.currentThread().interrupt();
+                }
+            }
+
+            assertNoWorkerFailures(failures);
+            assertStats(stats);
+        }
+    }
+
+    private void assertNoWorkerFailures(Queue<Throwable> failures) {
+        assertTrue(
+            "Unexpected worker failures:\n" + failures.stream().map(ExceptionsHelper::stackTrace).reduce("", (a, b) -> a + "\n---\n" + b),
+            failures.isEmpty()
+        );
+    }
+
+    private void assertStats(PollStats stats) {
+        assertEquals(stats.totalCalls.sum(), stats.runningResponses.sum() + stats.completedResponses.sum());
+        assertEquals("There should be no exceptions other than GONE", 0, stats.exceptions.sum());
+    }
+
+    private String getAsyncId(SearchResponseIterator it) {
+        AsyncSearchResponse response = it.next();
+        try {
+            assertNotNull(response.getId());
+            return response.getId();
+        } finally {
+            response.decRef();
+        }
+    }
+
+    private void consumeAllResponses(SearchResponseIterator it, String aggName) throws Exception {
+        while (it.hasNext()) {
+            AsyncSearchResponse response = it.next();
+            try {
+                if (response.getSearchResponse() != null && response.getSearchResponse().getAggregations() != null) {
+                    assertNotNull(response.getSearchResponse().getAggregations().get(aggName));
+                }
+            } finally {
+                response.decRef();
+            }
+        }
+    }
+
+    private CompletableFuture<Void> createPollers(
+        String id,
+        int threads,
+        PollStats stats,
+        CountDownLatch warmed,
+        ExecutorService pollerExec,
+        AtomicBoolean running,
+        Queue<Throwable> failures
+    ) {
+        @SuppressWarnings("unchecked")
+        final CompletableFuture<Void>[] tasks = IntStream.range(0, threads).mapToObj(i -> CompletableFuture.runAsync(() -> {
+            while (running.get()) {
+                AsyncSearchResponse resp = null;
+                try {
+                    resp = getAsyncSearch(id);
+                    stats.totalCalls.increment();
+
+                    // Once enough requests have been sent, consider pollers "warmed".
+                    if (stats.totalCalls.sum() >= threads) {
+                        warmed.countDown();
+                    }
+
+                    if (resp.isRunning()) {
+                        stats.runningResponses.increment();
+                    } else {
+                        // Success-only assertions: if reported completed, we must have a proper search response
+                        assertNull("Async search reported completed with failure", resp.getFailure());
+                        assertNotNull("Completed async search must carry a SearchResponse", resp.getSearchResponse());
+                        assertNotNull("Completed async search must have aggregations", resp.getSearchResponse().getAggregations());
+                        assertNotNull(
+                            "Completed async search must contain the expected aggregation",
+                            resp.getSearchResponse().getAggregations().get("terms")
+                        );
+                        stats.completedResponses.increment();
+                    }
+                } catch (Exception e) {
+                    Throwable cause = ExceptionsHelper.unwrapCause(e);
+                    if (cause instanceof ElasticsearchStatusException) {
+                        RestStatus status = ExceptionsHelper.status(cause);
+                        if (status == RestStatus.GONE) {
+                            stats.gone410.increment();
+                        } else {
+                            stats.exceptions.increment();
+                            failures.add(cause);
+                        }
+                    } else {
+                        stats.exceptions.increment();
+                        failures.add(cause);
+                    }
+                } finally {
+                    if (resp != null) {
+                        resp.decRef();
+                    }
+                }
+            }
+        }, pollerExec).whenComplete((v, ex) -> {
+            if (ex != null) {
+                failures.add(ExceptionsHelper.unwrapCause(ex));
+            }
+        })).toArray(CompletableFuture[]::new);
+
+        return CompletableFuture.allOf(tasks);
+    }
+
+    static final class PollStats {
+        final LongAdder totalCalls = new LongAdder();
+        final LongAdder runningResponses = new LongAdder();
+        final LongAdder completedResponses = new LongAdder();
+        final LongAdder exceptions = new LongAdder();
+        final LongAdder gone410 = new LongAdder();
+    }
+}