From 2e99ef5f69f2d0db8f2bf03dc4b56654bf41677c Mon Sep 17 00:00:00 2001
From: Armin <me@obrown.io>
Date: Mon, 28 Apr 2025 23:49:51 +0200
Subject: [PATCH] Release SearchContext referenced resources  via safer pattern

No need for two fields for this, using the future pattern is also much
safer and easier to reason about and automatically unlinks objects
on close.
---
 .../search/DefaultSearchContext.java          |  4 +++-
 .../search/internal/SearchContext.java        | 24 +++++++++----------
 2 files changed, 15 insertions(+), 13 deletions(-)

diff --git a/server/src/main/java/org/elasticsearch/search/DefaultSearchContext.java b/server/src/main/java/org/elasticsearch/search/DefaultSearchContext.java
index a6c15fc6b1a57..2bf284d2668f3 100644
--- a/server/src/main/java/org/elasticsearch/search/DefaultSearchContext.java
+++ b/server/src/main/java/org/elasticsearch/search/DefaultSearchContext.java
@@ -21,12 +21,14 @@
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.TotalHits;
 import org.apache.lucene.util.NumericUtils;
+import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.search.SearchType;
 import org.elasticsearch.cluster.routing.IndexRouting;
 import org.elasticsearch.common.breaker.CircuitBreaker;
 import org.elasticsearch.common.lucene.search.Queries;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.Releasable;
+import org.elasticsearch.core.Releasables;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.index.IndexMode;
 import org.elasticsearch.index.IndexService;
@@ -212,7 +214,7 @@ final class DefaultSearchContext extends SearchContext {
                     minimumDocsPerSlice
                 );
             }
-            releasables.addAll(List.of(engineSearcher, searcher));
+            closeFuture.addListener(ActionListener.releasing(Releasables.wrap(engineSearcher, searcher)));
             this.relativeTimeSupplier = relativeTimeSupplier;
             this.timeout = timeout;
             searchExecutionContext = indexService.newSearchExecutionContext(
diff --git a/server/src/main/java/org/elasticsearch/search/internal/SearchContext.java b/server/src/main/java/org/elasticsearch/search/internal/SearchContext.java
index d5e26b9c91a8d..9b1df285cd11c 100644
--- a/server/src/main/java/org/elasticsearch/search/internal/SearchContext.java
+++ b/server/src/main/java/org/elasticsearch/search/internal/SearchContext.java
@@ -11,12 +11,13 @@
 import org.apache.lucene.search.FieldDoc;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.TotalHits;
+import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.search.SearchType;
+import org.elasticsearch.action.support.SubscribableListener;
 import org.elasticsearch.common.breaker.CircuitBreaker;
 import org.elasticsearch.core.Assertions;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.Releasable;
-import org.elasticsearch.core.Releasables;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.index.cache.bitset.BitsetFilterCache;
 import org.elasticsearch.index.mapper.IdLoader;
@@ -56,8 +57,6 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
-import java.util.concurrent.CopyOnWriteArrayList;
-import java.util.concurrent.atomic.AtomicBoolean;
 
 /**
  * This class encapsulates the state needed to execute a search. It holds a reference to the
@@ -71,13 +70,16 @@ public abstract class SearchContext implements Releasable {
     public static final int TRACK_TOTAL_HITS_DISABLED = -1;
     public static final int DEFAULT_TRACK_TOTAL_HITS_UP_TO = 10000;
 
-    protected final List<Releasable> releasables = new CopyOnWriteArrayList<>();
-
-    private final AtomicBoolean closed = new AtomicBoolean(false);
+    protected final SubscribableListener<Void> closeFuture = new SubscribableListener<>();
 
     {
         if (Assertions.ENABLED) {
-            releasables.add(LeakTracker.wrap(() -> { assert closed.get(); }));
+            closeFuture.addListener(ActionListener.releasing(LeakTracker.wrap(new Releasable() {
+                @Override
+                public void close() {
+                    // empty instance that will actually get GC'ed so that the leak tracker works
+                }
+            })));
         }
     }
     private InnerHitsContext innerHitsContext;
@@ -109,9 +111,7 @@ public final List<Runnable> getCancellationChecks() {
 
     @Override
     public final void close() {
-        if (closed.compareAndSet(false, true)) {
-            Releasables.close(releasables);
-        }
+        closeFuture.onResponse(null);
     }
 
     /**
@@ -399,8 +399,8 @@ public final boolean checkRealMemoryCB(int locallyAccumulatedBytes, String label
      * Adds a releasable that will be freed when this context is closed.
      */
     public void addReleasable(Releasable releasable) {   // TODO most Releasables are managed by their callers. We probably don't need this.
-        assert closed.get() == false;
-        releasables.add(releasable);
+        assert closeFuture.isDone() == false;
+        closeFuture.addListener(ActionListener.releasing(releasable));
     }
 
     /**