Make MutableSearchResponse ref-counted to prevent use-after-close in async search

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,266 @@
+/*
+ * 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.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.CancellationException;
+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.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.concurrent.atomic.LongAdder;
+
+@SuiteScopeTestCase
+public class AsyncSearchConcurrentStatusIT extends AsyncSearchIntegTestCase {
+    private static String indexName;
+    private static int numShards;
+
+    private static int numKeywords;
+    private static Map<String, AtomicInteger> keywordFreqs;
+    private static float maxMetric = Float.NEGATIVE_INFINITY;
+    private static float minMetric = Float.POSITIVE_INFINITY;
+
+    @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);
+        keywordFreqs = new HashMap<>();
+        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();
+            maxMetric = Math.max(metric, maxMetric);
+            minMetric = Math.min(metric, minMetric);
+            String keyword = keywords[randomIntBetween(0, numKeywords - 1)];
+            keywordFreqs.compute(keyword, (k, v) -> {
+                if (v == null) {
+                    return new AtomicInteger(1);
+                }
+                v.incrementAndGet();
+                return v;
+            });
+            reqs.add(prepareIndex(indexName).setSource("terms", keyword, "metric", metric));
+        }
+        indexRandom(true, true, reqs);
+    }
+
+    /**
+     * Tests that concurrent async search status requests behave correctly
+     * while the underlying async search task is still executing and during its close/cleanup.
+     */
+    public void testConcurrentStatusFetchWhileTaskCloses() throws Exception {
+        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();
+
+            int statusThreads = randomIntBetween(1, Math.max(2, 4 * numShards));
+            StartableThreadGroup pollers = startGetStatusThreadsHot(id, statusThreads, aggName, stats);
+            pollers.startHotThreads.countDown(); // release pollers
+
+            // Finish consumption 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);
+                }
+            });
+
+            Thread.sleep(randomIntBetween(100, 200));
+            pollers.stopAndAwait(TimeValue.timeValueMillis(randomIntBetween(500, 1000)));
+
+            // Join consumer & surface errors
+            try {
+                consumer.get();
+            } catch (Exception ignored) {} finally {
+                consumerExec.shutdown();
+            }
+
+            assertNull(consumerError.get());
+            assertNoWorkerFailures(pollers);
+            assertStats(stats);
+        }
+    }
+
+    private void assertNoWorkerFailures(StartableThreadGroup pollers) {
+        List<Throwable> failures = pollers.getFailures();
+        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 StartableThreadGroup startGetStatusThreadsHot(String id, int threads, String aggName, PollStats stats) {
+        final ExecutorService exec = Executors.newFixedThreadPool(threads);
+        final List<Future<?>> futures = new ArrayList<>(threads);
+        final AtomicBoolean running = new AtomicBoolean(true);
+        final CountDownLatch start = new CountDownLatch(1);
+
+        for (int i = 0; i < threads; i++) {
+            futures.add(exec.submit(() -> {
+                start.await();
+                while (running.get()) {
+                    AsyncSearchResponse resp = null;
+                    try {
+                        resp = getAsyncSearch(id);
+                        stats.totalCalls.increment();
+
+                        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(aggName)
+                            );
+                            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();
+                        }
+                    } finally {
+                        if (resp != null) {
+                            resp.decRef();
+                        }
+                    }
+                }
+                return null;
+            }));
+        }
+        return new StartableThreadGroup(exec, futures, running, start);
+    }
+
+    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();
+    }
+
+    static class StartableThreadGroup extends ThreadGroup {
+        private final CountDownLatch startHotThreads;
+
+        StartableThreadGroup(ExecutorService exec, List<Future<?>> futures, AtomicBoolean running, CountDownLatch startHotThreads) {
+            super(exec, futures, running);
+            this.startHotThreads = startHotThreads;
+        }
+    }
+
+    static class ThreadGroup {
+        private final ExecutorService exec;
+        private final List<Future<?>> futures;
+        private final AtomicBoolean running;
+
+        private ThreadGroup(ExecutorService exec, List<Future<?>> futures, AtomicBoolean running) {
+            this.exec = exec;
+            this.futures = futures;
+            this.running = running;
+        }
+
+        void stopAndAwait(TimeValue timeout) throws InterruptedException {
+            running.set(false);
+            exec.shutdown();
+            if (exec.awaitTermination(timeout.millis(), TimeUnit.MILLISECONDS) == false) {
+                exec.shutdownNow();
+                exec.awaitTermination(timeout.millis(), TimeUnit.MILLISECONDS);
+            }
+        }
+
+        List<Throwable> getFailures() {
+            List<Throwable> failures = new ArrayList<>();
+            for (Future<?> f : futures) {
+                try {
+                    f.get();
+                } catch (CancellationException ignored) {} catch (ExecutionException ee) {
+                    failures.add(ExceptionsHelper.unwrapCause(ee.getCause()));
+                } catch (InterruptedException ie) {
+                    Thread.currentThread().interrupt();
+                    if (failures.isEmpty()) failures.add(ie);
+                }
+            }
+            return failures;
+        }
+    }
+}