Refcount responses in TransportNodesAction (#103254)

DaveCTurner · web-flow · commit 6c13a815bdb3 · 2024-01-02T13:38:02.000Z
Today we `decRef()` the per-node responses just after adding them to the
`responses` collection, but in fact we should keep them alive until
we've constructed the final response. This commit does that.
diff --git a/libs/core/src/main/java/org/elasticsearch/core/Releasables.java b/libs/core/src/main/java/org/elasticsearch/core/Releasables.java
@@ -11,6 +11,7 @@
 import java.io.IOException;
 import java.io.UncheckedIOException;
 import java.util.Arrays;
+import java.util.Iterator;
 import java.util.concurrent.atomic.AtomicReference;
 
 /** Utility methods to work with {@link Releasable}s. */
@@ -103,6 +104,24 @@ public String toString() {
         };
     }
 
+    /**
+     * Similar to {@link #wrap(Iterable)} except that it accepts an {@link Iterator} of releasables. The resulting resource must therefore
+     * only be released once.
+     */
+    public static Releasable wrap(final Iterator<Releasable> releasables) {
+        return assertOnce(wrap(new Iterable<>() {
+            @Override
+            public Iterator<Releasable> iterator() {
+                return releasables;
+            }
+
+            @Override
+            public String toString() {
+                return releasables.toString();
+            }
+        }));
+    }
+
     /** @see #wrap(Iterable) */
     public static Releasable wrap(final Releasable... releasables) {
         return new Releasable() {
diff --git a/libs/core/src/test/java/org/elasticsearch/core/ReleasablesTests.java b/libs/core/src/test/java/org/elasticsearch/core/ReleasablesTests.java
@@ -107,5 +107,27 @@ public String toString() {
         assertEquals("wrapped[list]", wrapIterable.toString());
         wrapIterable.close();
         assertEquals(5, count.get());
+
+        final var wrapIterator = Releasables.wrap(new Iterator<>() {
+            final Iterator<Releasable> innerIterator = List.of(releasable, releasable, releasable).iterator();
+
+            @Override
+            public boolean hasNext() {
+                return innerIterator.hasNext();
+            }
+
+            @Override
+            public Releasable next() {
+                return innerIterator.next();
+            }
+
+            @Override
+            public String toString() {
+                return "iterator";
+            }
+        });
+        assertEquals("wrapped[iterator]", wrapIterator.toString());
+        wrapIterator.close();
+        assertEquals(8, count.get());
     }
 }
diff --git a/server/src/main/java/org/elasticsearch/action/support/nodes/TransportNodesAction.java b/server/src/main/java/org/elasticsearch/action/support/nodes/TransportNodesAction.java
@@ -25,6 +25,8 @@
 import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.common.util.concurrent.EsExecutors;
 import org.elasticsearch.core.CheckedConsumer;
+import org.elasticsearch.core.Releasables;
+import org.elasticsearch.tasks.CancellableTask;
 import org.elasticsearch.tasks.Task;
 import org.elasticsearch.transport.TransportChannel;
 import org.elasticsearch.transport.TransportRequest;
@@ -96,6 +98,23 @@ protected void doExecute(Task task, NodesRequest request, ActionListener<NodesRe
 
             final TransportRequestOptions transportRequestOptions = TransportRequestOptions.timeout(request.timeout());
 
+            {
+                addReleaseOnCancellationListener();
+            }
+
+            private void addReleaseOnCancellationListener() {
+                if (task instanceof CancellableTask cancellableTask) {
+                    cancellableTask.addListener(() -> {
+                        final List<NodeResponse> drainedResponses;
+                        synchronized (responses) {
+                            drainedResponses = List.copyOf(responses);
+                            responses.clear();
+                        }
+                        Releasables.wrap(Iterators.map(drainedResponses.iterator(), r -> r::decRef)).close();
+                    });
+                }
+            }
+
             @Override
             protected void sendItemRequest(DiscoveryNode discoveryNode, ActionListener<NodeResponse> listener) {
                 final var nodeRequest = newNodeRequest(request);
@@ -118,9 +137,14 @@ protected void sendItemRequest(DiscoveryNode discoveryNode, ActionListener<NodeR
 
             @Override
             protected void onItemResponse(DiscoveryNode discoveryNode, NodeResponse nodeResponse) {
+                nodeResponse.mustIncRef();
                 synchronized (responses) {
-                    responses.add(nodeResponse);
+                    if ((task instanceof CancellableTask cancellableTask && cancellableTask.isCancelled()) == false) {
+                        responses.add(nodeResponse);
+                        return;
+                    }
                 }
+                nodeResponse.decRef();
             }
 
             @Override
@@ -134,7 +158,11 @@ protected void onItemFailure(DiscoveryNode discoveryNode, Exception e) {
             @Override
             protected CheckedConsumer<ActionListener<NodesResponse>, Exception> onCompletion() {
                 // ref releases all happen-before here so no need to be synchronized
-                return l -> newResponseAsync(task, request, responses, exceptions, l);
+                return l -> {
+                    try (var ignored = Releasables.wrap(Iterators.map(responses.iterator(), r -> r::decRef))) {
+                        newResponseAsync(task, request, responses, exceptions, l);
+                    }
+                };
             }
 
             @Override
@@ -154,9 +182,11 @@ private Writeable.Reader<NodeResponse> nodeResponseReader(DiscoveryNode discover
     }
 
     /**
-     * Create a new {@link NodesResponse} (multi-node response).
+     * Create a new {@link NodesResponse}. This method is executed on {@link #finalExecutor}.
      *
-     * @param request The associated request.
+     * @param request The request whose response we are constructing. {@link TransportNodesAction} may have already released all its
+     *                references to this object before calling this method, so it's up to individual implementations to retain their own
+     *                reference to the request if still needed here.
      * @param responses All successful node-level responses.
      * @param failures All node-level failures.
      * @return Never {@code null}.
@@ -166,7 +196,11 @@ private Writeable.Reader<NodeResponse> nodeResponseReader(DiscoveryNode discover
 
     /**
      * Create a new {@link NodesResponse}, possibly asynchronously. The default implementation is synchronous and calls
-     * {@link #newResponse(BaseNodesRequest, List, List)}
+     * {@link #newResponse(BaseNodesRequest, List, List)}. This method is executed on {@link #finalExecutor}.
+     *
+     * @param request The request whose response we are constructing. {@link TransportNodesAction} may have already released all its
+     *                references to this object before calling this method, so it's up to individual implementations to retain their own
+     *                reference to the request if still needed here.
      */
     protected void newResponseAsync(
         Task task,
diff --git a/server/src/test/java/org/elasticsearch/action/support/nodes/TransportNodesActionTests.java b/server/src/test/java/org/elasticsearch/action/support/nodes/TransportNodesActionTests.java
@@ -9,10 +9,12 @@
 package org.elasticsearch.action.support.nodes;
 
 import org.elasticsearch.ElasticsearchException;
+import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.FailedNodeException;
 import org.elasticsearch.action.support.ActionFilters;
 import org.elasticsearch.action.support.PlainActionFuture;
-import org.elasticsearch.action.support.broadcast.node.TransportBroadcastByNodeActionTests;
+import org.elasticsearch.action.support.RefCountingListener;
+import org.elasticsearch.action.support.SubscribableListener;
 import org.elasticsearch.cluster.ClusterName;
 import org.elasticsearch.cluster.ClusterState;
 import org.elasticsearch.cluster.node.DiscoveryNode;
@@ -55,6 +57,9 @@
 import java.util.Set;
 import java.util.concurrent.Executor;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Function;
+import java.util.function.ObjLongConsumer;
 
 import static java.util.Collections.emptyMap;
 import static org.elasticsearch.test.ClusterServiceUtils.createClusterService;
@@ -118,7 +123,10 @@ public void testResponseAggregation() {
         final TestTransportNodesAction action = getTestTransportNodesAction();
 
         final PlainActionFuture<TestNodesResponse> listener = new PlainActionFuture<>();
-        action.execute(null, new TestNodesRequest(), listener);
+        action.execute(null, new TestNodesRequest(), listener.delegateFailure((l, response) -> {
+            assertTrue(response.getNodes().stream().allMatch(TestNodeResponse::hasReferences));
+            l.onResponse(response);
+        }));
         assertFalse(listener.isDone());
 
         final Set<String> failedNodeIds = new HashSet<>();
@@ -127,7 +135,9 @@ public void testResponseAggregation() {
         for (CapturingTransport.CapturedRequest capturedRequest : transport.getCapturedRequestsAndClear()) {
             if (randomBoolean()) {
                 successfulNodes.add(capturedRequest.node());
-                transport.handleResponse(capturedRequest.requestId(), new TestNodeResponse(capturedRequest.node()));
+                final var response = new TestNodeResponse(capturedRequest.node());
+                transport.handleResponse(capturedRequest.requestId(), response);
+                assertFalse(response.hasReferences()); // response is copied (via the wire protocol) so this instance is released
             } else {
                 failedNodeIds.add(capturedRequest.node().getId());
                 if (randomBoolean()) {
@@ -138,7 +148,16 @@ public void testResponseAggregation() {
             }
         }
 
-        TestNodesResponse response = listener.actionGet(10, TimeUnit.SECONDS);
+        final TestNodesResponse response = listener.actionGet(10, TimeUnit.SECONDS);
+
+        final var allResponsesReleasedListener = new SubscribableListener<Void>();
+        try (var listeners = new RefCountingListener(allResponsesReleasedListener)) {
+            for (final var nodeResponse : response.getNodes()) {
+                nodeResponse.addCloseListener(listeners.acquire());
+            }
+        }
+        safeAwait(allResponsesReleasedListener);
+        assertTrue(response.getNodes().stream().noneMatch(TestNodeResponse::hasReferences));
 
         for (TestNodeResponse nodeResponse : response.getNodes()) {
             assertThat(successfulNodes, Matchers.hasItem(nodeResponse.getNode()));
@@ -164,7 +183,7 @@ public void testResponsesReleasedOnCancellation() {
         final CancellableTask cancellableTask = new CancellableTask(randomLong(), "transport", "action", "", null, emptyMap());
         final PlainActionFuture<TestNodesResponse> listener = new PlainActionFuture<>();
         action.execute(cancellableTask, new TestNodesRequest(), listener.delegateResponse((l, e) -> {
-            assert Thread.currentThread().getName().contains("[" + ThreadPool.Names.GENERIC + "]");
+            assert ThreadPool.assertCurrentThreadPool(ThreadPool.Names.GENERIC);
             l.onFailure(e);
         }));
 
@@ -173,13 +192,31 @@ public void testResponsesReleasedOnCancellation() {
         );
         Randomness.shuffle(capturedRequests);
 
+        final AtomicInteger liveResponseCount = new AtomicInteger();
+        final Function<DiscoveryNode, TestNodeResponse> responseCreator = node -> {
+            liveResponseCount.incrementAndGet();
+            final var testNodeResponse = new TestNodeResponse(node);
+            testNodeResponse.addCloseListener(ActionListener.running(liveResponseCount::decrementAndGet));
+            return testNodeResponse;
+        };
+
+        final ObjLongConsumer<TestNodeResponse> responseSender = (response, requestId) -> {
+            try {
+                // transport.handleResponse may de/serialize the response, releasing it early, so send the response straight to the handler
+                transport.getTransportResponseHandler(requestId).handleResponse(response);
+            } finally {
+                response.decRef();
+            }
+        };
+
         final ReachabilityChecker reachabilityChecker = new ReachabilityChecker();
         final Runnable nextRequestProcessor = () -> {
             var capturedRequest = capturedRequests.remove(0);
             if (randomBoolean()) {
-                // transport.handleResponse may de/serialize the response, releasing it early, so send the response straight to the handler
-                transport.getTransportResponseHandler(capturedRequest.requestId())
-                    .handleResponse(reachabilityChecker.register(new TestNodeResponse(capturedRequest.node())));
+                responseSender.accept(
+                    reachabilityChecker.register(responseCreator.apply(capturedRequest.node())),
+                    capturedRequest.requestId()
+                );
             } else {
                 // handleRemoteError may de/serialize the exception, releasing it early, so just use handleLocalError
                 transport.handleLocalError(
@@ -200,20 +237,23 @@ public void testResponsesReleasedOnCancellation() {
 
         // responses captured before cancellation are now unreachable
         reachabilityChecker.ensureUnreachable();
+        assertEquals(0, liveResponseCount.get());
 
         while (capturedRequests.size() > 0) {
             // a response sent after cancellation is dropped immediately
             assertFalse(listener.isDone());
             nextRequestProcessor.run();
             reachabilityChecker.ensureUnreachable();
+            assertEquals(0, liveResponseCount.get());
         }
 
         expectThrows(TaskCancelledException.class, () -> listener.actionGet(10, TimeUnit.SECONDS));
+        assertTrue(cancellableTask.isCancelled()); // keep task alive
     }
 
     @BeforeClass
     public static void startThreadPool() {
-        THREAD_POOL = new TestThreadPool(TransportBroadcastByNodeActionTests.class.getSimpleName());
+        THREAD_POOL = new TestThreadPool(TransportNodesActionTests.class.getSimpleName());
     }
 
     @AfterClass
@@ -268,11 +308,9 @@ public void tearDown() throws Exception {
 
     public TestTransportNodesAction getTestTransportNodesAction() {
         return new TestTransportNodesAction(
-            THREAD_POOL,
             clusterService,
             transportService,
             new ActionFilters(Collections.emptySet()),
-            TestNodesRequest::new,
             TestNodeRequest::new,
             THREAD_POOL.executor(ThreadPool.Names.GENERIC)
         );
@@ -302,11 +340,9 @@ private static class TestTransportNodesAction extends TransportNodesAction<
         TestNodeResponse> {
 
         TestTransportNodesAction(
-            ThreadPool threadPool,
             ClusterService clusterService,
             TransportService transportService,
             ActionFilters actionFilters,
-            Writeable.Reader<TestNodesRequest> request,
             Writeable.Reader<TestNodeRequest> nodeRequest,
             Executor nodeExecutor
         ) {
@@ -319,7 +355,7 @@ protected TestNodesResponse newResponse(
             List<TestNodeResponse> responses,
             List<FailedNodeException> failures
         ) {
-            return new TestNodesResponse(clusterService.getClusterName(), request, responses, failures);
+            return new TestNodesResponse(clusterService.getClusterName(), responses, failures);
         }
 
         @Override
@@ -350,7 +386,7 @@ private static class DataNodesOnlyTransportNodesAction extends TestTransportNode
             Writeable.Reader<TestNodeRequest> nodeRequest,
             Executor nodeExecutor
         ) {
-            super(threadPool, clusterService, transportService, actionFilters, request, nodeRequest, nodeExecutor);
+            super(clusterService, transportService, actionFilters, nodeRequest, nodeExecutor);
         }
 
         @Override
@@ -371,16 +407,8 @@ private static class TestNodesRequest extends BaseNodesRequest<TestNodesRequest>
 
     private static class TestNodesResponse extends BaseNodesResponse<TestNodeResponse> {
 
-        private final TestNodesRequest request;
-
-        TestNodesResponse(
-            ClusterName clusterName,
-            TestNodesRequest request,
-            List<TestNodeResponse> nodeResponses,
-            List<FailedNodeException> failures
-        ) {
+        TestNodesResponse(ClusterName clusterName, List<TestNodeResponse> nodeResponses, List<FailedNodeException> failures) {
             super(clusterName, nodeResponses, failures);
-            this.request = request;
         }
 
         @Override
@@ -425,6 +453,10 @@ public boolean hasReferences() {
     }
 
     private static class TestNodeResponse extends BaseNodeResponse {
+
+        private final SubscribableListener<Void> onClose = new SubscribableListener<>();
+        private final RefCounted refCounted = AbstractRefCounted.of(() -> onClose.onResponse(null));
+
         TestNodeResponse() {
             this(mock(DiscoveryNode.class));
         }
@@ -436,6 +468,30 @@ private static class TestNodeResponse extends BaseNodeResponse {
         protected TestNodeResponse(StreamInput in) throws IOException {
             super(in);
         }
+
+        @Override
+        public void incRef() {
+            refCounted.incRef();
+        }
+
+        @Override
+        public boolean tryIncRef() {
+            return refCounted.tryIncRef();
+        }
+
+        @Override
+        public boolean decRef() {
+            return refCounted.decRef();
+        }
+
+        @Override
+        public boolean hasReferences() {
+            return refCounted.hasReferences();
+        }
+
+        void addCloseListener(ActionListener<Void> listener) {
+            onClose.addListener(listener);
+        }
     }
 
 }
