Add RestChunkedToXContentListener version that releases transport res…

…ponses (elastic#102607)

This is needed as part of elastic#102030, we need to be able to release
search/multi-search and maybe other transport messages after they've
been serialized as REST responses.

server/src/main/java/org/elasticsearch/rest/action/RestChunkedToXContentListener.java

@@ -9,6 +9,7 @@
 package org.elasticsearch.rest.action;
 
 import org.elasticsearch.common.xcontent.ChunkedToXContent;
+import org.elasticsearch.core.Releasable;
 import org.elasticsearch.rest.ChunkedRestResponseBody;
 import org.elasticsearch.rest.RestChannel;
 import org.elasticsearch.rest.RestResponse;
@@ -37,10 +38,17 @@ public RestChunkedToXContentListener(RestChannel channel, ToXContent.Params para
     @Override
     protected void processResponse(Response response) throws IOException {
         channel.sendResponse(
-            RestResponse.chunked(getRestStatus(response), ChunkedRestResponseBody.fromXContent(response, params, channel, null))
+            RestResponse.chunked(
+                getRestStatus(response),
+                ChunkedRestResponseBody.fromXContent(response, params, channel, releasableFromResponse(response))
+            )
         );
     }
 
+    protected Releasable releasableFromResponse(Response response) {
+        return null;
+    }
+
     protected RestStatus getRestStatus(Response response) {
         return RestStatus.OK;
     }

server/src/main/java/org/elasticsearch/rest/action/RestRefCountedChunkedToXContentListener.java

@@ -0,0 +1,32 @@
+/*
+ * 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 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.rest.action;
+
+import org.elasticsearch.common.xcontent.ChunkedToXContent;
+import org.elasticsearch.core.RefCounted;
+import org.elasticsearch.core.Releasable;
+import org.elasticsearch.core.Releasables;
+import org.elasticsearch.rest.RestChannel;
+
+/**
+ * Same as {@link RestChunkedToXContentListener} but decrements the ref count on the response it receives by one after serialization of the
+ * response.
+ */
+public class RestRefCountedChunkedToXContentListener<Response extends ChunkedToXContent & RefCounted> extends RestChunkedToXContentListener<
+    Response> {
+    public RestRefCountedChunkedToXContentListener(RestChannel channel) {
+        super(channel);
+    }
+
+    @Override
+    protected Releasable releasableFromResponse(Response response) {
+        response.mustIncRef();
+        return Releasables.assertOnce(response::decRef);
+    }
+}

server/src/main/java/org/elasticsearch/rest/action/search/RestMultiSearchAction.java

@@ -26,7 +26,7 @@
 import org.elasticsearch.rest.Scope;
 import org.elasticsearch.rest.ServerlessScope;
 import org.elasticsearch.rest.action.RestCancellableNodeClient;
-import org.elasticsearch.rest.action.RestChunkedToXContentListener;
+import org.elasticsearch.rest.action.RestRefCountedChunkedToXContentListener;
 import org.elasticsearch.search.builder.SearchSourceBuilder;
 import org.elasticsearch.usage.SearchUsageHolder;
 import org.elasticsearch.xcontent.XContent;
@@ -82,7 +82,11 @@ public RestChannelConsumer prepareRequest(final RestRequest request, final NodeC
         );
         return channel -> {
             final RestCancellableNodeClient cancellableClient = new RestCancellableNodeClient(client, request.getHttpChannel());
-            cancellableClient.execute(MultiSearchAction.INSTANCE, multiSearchRequest, new RestChunkedToXContentListener<>(channel));
+            cancellableClient.execute(
+                MultiSearchAction.INSTANCE,
+                multiSearchRequest,
+                new RestRefCountedChunkedToXContentListener<>(channel)
+            );
         };
     }
 

server/src/main/java/org/elasticsearch/rest/action/search/RestSearchAction.java

@@ -29,7 +29,7 @@
 import org.elasticsearch.rest.ServerlessScope;
 import org.elasticsearch.rest.action.RestActions;
 import org.elasticsearch.rest.action.RestCancellableNodeClient;
-import org.elasticsearch.rest.action.RestChunkedToXContentListener;
+import org.elasticsearch.rest.action.RestRefCountedChunkedToXContentListener;
 import org.elasticsearch.search.Scroll;
 import org.elasticsearch.search.SearchService;
 import org.elasticsearch.search.builder.SearchSourceBuilder;
@@ -121,7 +121,7 @@ public RestChannelConsumer prepareRequest(final RestRequest request, final NodeC
 
         return channel -> {
             RestCancellableNodeClient cancelClient = new RestCancellableNodeClient(client, request.getHttpChannel());
-            cancelClient.execute(SearchAction.INSTANCE, searchRequest, new RestChunkedToXContentListener<>(channel));
+            cancelClient.execute(SearchAction.INSTANCE, searchRequest, new RestRefCountedChunkedToXContentListener<>(channel));
         };
     }