Break out clear scroll logic from TransportClearScrollAction

core/src/main/java/org/elasticsearch/action/search/ClearScrollController.java

@@ -0,0 +1,141 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.elasticsearch.action.search;
+
+import org.apache.logging.log4j.Logger;
+import org.apache.logging.log4j.message.ParameterizedMessage;
+import org.apache.logging.log4j.util.Supplier;
+import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.cluster.node.DiscoveryNode;
+import org.elasticsearch.cluster.node.DiscoveryNodes;
+import org.elasticsearch.common.util.concurrent.CountDown;
+import org.elasticsearch.transport.Transport;
+import org.elasticsearch.transport.TransportResponse;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.elasticsearch.action.search.TransportSearchHelper.parseScrollId;
+
+final class ClearScrollController implements Runnable {
+    private final DiscoveryNodes nodes;
+    private final SearchTransportService searchTransportService;
+    private final CountDown expectedOps;
+    private final ActionListener<ClearScrollResponse> listener;
+    private final AtomicBoolean hasFailed = new AtomicBoolean(false);
+    private final AtomicInteger freedSearchContexts = new AtomicInteger(0);
+    private final Logger logger;
+    private final Runnable runner;
+
+    ClearScrollController(ClearScrollRequest request, ActionListener<ClearScrollResponse> listener, DiscoveryNodes nodes, Logger logger,
+                          SearchTransportService searchTransportService) {
+        this.nodes = nodes;
+        this.logger = logger;
+        this.searchTransportService = searchTransportService;
+        this.listener = listener;
+        List<String> scrollIds = request.getScrollIds();
+        final int expectedOps;
+        if (scrollIds.size() == 1 && "_all".equals(scrollIds.get(0))) {
+            expectedOps = nodes.getSize();
+            runner = this::cleanAllScrolls;
+        } else {
+            List<ScrollIdForNode> parsedScrollIds = new ArrayList<>();
+            for (String parsedScrollId : request.getScrollIds()) {
+                ScrollIdForNode[] context = parseScrollId(parsedScrollId).getContext();
+                for (ScrollIdForNode id : context) {
+                    parsedScrollIds.add(id);
+                }
+            }
+            if (parsedScrollIds.isEmpty()) {
+                expectedOps = 0;
+                runner = () -> listener.onResponse(new ClearScrollResponse(true, 0));
+            } else {
+                expectedOps = parsedScrollIds.size();
+                runner = () -> cleanScrollIds(parsedScrollIds);
+            }
+        }
+        this.expectedOps = new CountDown(expectedOps);
+
+    }
+
+    @Override
+    public void run() {
+        runner.run();
+    }
+
+    void cleanAllScrolls() {
+        for (final DiscoveryNode node : nodes) {
+            try {
+                Transport.Connection connection = searchTransportService.getConnection(null, node);
+                searchTransportService.sendClearAllScrollContexts(connection, new ActionListener<TransportResponse>() {
+                    @Override
+                    public void onResponse(TransportResponse response) {
+                        onFreedContext(true);
+                    }
+
+                    @Override
+                    public void onFailure(Exception e) {
+                        onFailedFreedContext(e, node);
+                    }
+                });
+            } catch (Exception e) {
+                onFailedFreedContext(e, node);
+            }
+        }
+    }
+
+    void cleanScrollIds(List<ScrollIdForNode> parsedScrollIds) {
+        for (ScrollIdForNode target : parsedScrollIds) {
+            final DiscoveryNode node = nodes.get(target.getNode());
+            if (node == null) {
+                onFreedContext(false);
+            } else {
+                try {
+                    Transport.Connection connection = searchTransportService.getConnection(null, node);
+                    searchTransportService.sendFreeContext(connection, target.getScrollId(),
+                        ActionListener.wrap(freed -> onFreedContext(freed.isFreed()),
+                            e -> onFailedFreedContext(e, node)));
+                } catch (Exception e) {
+                    onFailedFreedContext(e, node);
+                }
+            }
+        }
+    }
+
+    private void onFreedContext(boolean freed) {
+        if (freed) {
+            freedSearchContexts.incrementAndGet();
+        }
+        if (expectedOps.countDown()) {
+            boolean succeeded = hasFailed.get() == false;
+            listener.onResponse(new ClearScrollResponse(succeeded, freedSearchContexts.get()));
+        }
+    }
+
+    private void onFailedFreedContext(Throwable e, DiscoveryNode node) {
+        logger.warn((Supplier<?>) () -> new ParameterizedMessage("Clear SC failed on node[{}]", node), e);
+        if (expectedOps.countDown()) {
+            listener.onResponse(new ClearScrollResponse(false, freedSearchContexts.get()));
+        } else {
+            hasFailed.set(true);
+        }
+    }
+}

core/src/main/java/org/elasticsearch/action/search/SearchTransportService.java

@@ -98,14 +98,14 @@ public void onFailure(Exception e) {
             }, SearchFreeContextResponse::new));
     }
 
-    public void sendFreeContext(DiscoveryNode node, long contextId, final ActionListener<SearchFreeContextResponse> listener) {
-        transportService.sendRequest(node, FREE_CONTEXT_SCROLL_ACTION_NAME, new ScrollFreeContextRequest(contextId),
-            new ActionListenerResponseHandler<>(listener, SearchFreeContextResponse::new));
+    public void sendFreeContext(Transport.Connection connection, long contextId, final ActionListener<SearchFreeContextResponse> listener) {
+        transportService.sendRequest(connection, FREE_CONTEXT_SCROLL_ACTION_NAME, new ScrollFreeContextRequest(contextId),
+            TransportRequestOptions.EMPTY, new ActionListenerResponseHandler<>(listener, SearchFreeContextResponse::new));
     }
 
-    public void sendClearAllScrollContexts(DiscoveryNode node, final ActionListener<TransportResponse> listener) {
-        transportService.sendRequest(node, CLEAR_SCROLL_CONTEXTS_ACTION_NAME, TransportRequest.Empty.INSTANCE,
-            new ActionListenerResponseHandler<>(listener, () -> TransportResponse.Empty.INSTANCE));
+    public void sendClearAllScrollContexts(Transport.Connection connection, final ActionListener<TransportResponse> listener) {
+        transportService.sendRequest(connection, CLEAR_SCROLL_CONTEXTS_ACTION_NAME, TransportRequest.Empty.INSTANCE,
+            TransportRequestOptions.EMPTY, new ActionListenerResponseHandler<>(listener, () -> TransportResponse.Empty.INSTANCE));
     }
 
     public void sendExecuteDfs(Transport.Connection connection, final ShardSearchTransportRequest request, SearchTask task,

core/src/main/java/org/elasticsearch/action/search/TransportClearScrollAction.java

@@ -19,30 +19,16 @@
 
 package org.elasticsearch.action.search;
 
-import org.apache.logging.log4j.message.ParameterizedMessage;
-import org.apache.logging.log4j.util.Supplier;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.support.ActionFilters;
 import org.elasticsearch.action.support.HandledTransportAction;
-import org.elasticsearch.cluster.ClusterState;
 import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver;
-import org.elasticsearch.cluster.node.DiscoveryNode;
-import org.elasticsearch.cluster.node.DiscoveryNodes;
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.inject.Inject;
 import org.elasticsearch.common.settings.Settings;
-import org.elasticsearch.common.util.concurrent.CountDown;
 import org.elasticsearch.threadpool.ThreadPool;
-import org.elasticsearch.transport.TransportResponse;
 import org.elasticsearch.transport.TransportService;
 
-import java.util.ArrayList;
-import java.util.List;
-import java.util.concurrent.atomic.AtomicInteger;
-import java.util.concurrent.atomic.AtomicReference;
-
-import static org.elasticsearch.action.search.TransportSearchHelper.parseScrollId;
-
 public class TransportClearScrollAction extends HandledTransportAction<ClearScrollRequest, ClearScrollResponse> {
 
     private final ClusterService clusterService;
@@ -53,105 +39,16 @@ public TransportClearScrollAction(Settings settings, TransportService transportS
                                       ClusterService clusterService, ActionFilters actionFilters,
                                       IndexNameExpressionResolver indexNameExpressionResolver,
                                       SearchTransportService searchTransportService) {
-        super(settings, ClearScrollAction.NAME, threadPool, transportService, actionFilters, indexNameExpressionResolver, ClearScrollRequest::new);
+        super(settings, ClearScrollAction.NAME, threadPool, transportService, actionFilters, indexNameExpressionResolver,
+            ClearScrollRequest::new);
         this.clusterService = clusterService;
         this.searchTransportService = searchTransportService;
     }
 
     @Override
     protected void doExecute(ClearScrollRequest request, final ActionListener<ClearScrollResponse> listener) {
-        new Async(request, listener, clusterService.state()).run();
-    }
-
-    private class Async {
-        final DiscoveryNodes nodes;
-        final CountDown expectedOps;
-        final List<ScrollIdForNode[]> contexts = new ArrayList<>();
-        final ActionListener<ClearScrollResponse> listener;
-        final AtomicReference<Throwable> expHolder;
-        final AtomicInteger numberOfFreedSearchContexts = new AtomicInteger(0);
-
-        private Async(ClearScrollRequest request, ActionListener<ClearScrollResponse> listener, ClusterState clusterState) {
-            int expectedOps = 0;
-            this.nodes = clusterState.nodes();
-            if (request.getScrollIds().size() == 1 && "_all".equals(request.getScrollIds().get(0))) {
-                expectedOps = nodes.getSize();
-            } else {
-                for (String parsedScrollId : request.getScrollIds()) {
-                    ScrollIdForNode[] context = parseScrollId(parsedScrollId).getContext();
-                    expectedOps += context.length;
-                    this.contexts.add(context);
-                }
-            }
-            this.listener = listener;
-            this.expHolder = new AtomicReference<>();
-            this.expectedOps = new CountDown(expectedOps);
-        }
-
-        public void run() {
-            if (expectedOps.isCountedDown()) {
-                listener.onResponse(new ClearScrollResponse(true, 0));
-                return;
-            }
-
-            if (contexts.isEmpty()) {
-                for (final DiscoveryNode node : nodes) {
-                    searchTransportService.sendClearAllScrollContexts(node, new ActionListener<TransportResponse>() {
-                        @Override
-                        public void onResponse(TransportResponse response) {
-                            onFreedContext(true);
-                        }
-
-                        @Override
-                        public void onFailure(Exception e) {
-                            onFailedFreedContext(e, node);
-                        }
-                    });
-                }
-            } else {
-                for (ScrollIdForNode[] context : contexts) {
-                    for (ScrollIdForNode target : context) {
-                        final DiscoveryNode node = nodes.get(target.getNode());
-                        if (node == null) {
-                            onFreedContext(false);
-                            continue;
-                        }
-
-                        searchTransportService.sendFreeContext(node, target.getScrollId(), new ActionListener<SearchTransportService.SearchFreeContextResponse>() {
-                            @Override
-                            public void onResponse(SearchTransportService.SearchFreeContextResponse freed) {
-                                onFreedContext(freed.isFreed());
-                            }
-
-                            @Override
-                            public void onFailure(Exception e) {
-                                onFailedFreedContext(e, node);
-                            }
-                        });
-                    }
-                }
-            }
-        }
-
-        void onFreedContext(boolean freed) {
-            if (freed) {
-                numberOfFreedSearchContexts.incrementAndGet();
-            }
-            if (expectedOps.countDown()) {
-                boolean succeeded = expHolder.get() == null;
-                listener.onResponse(new ClearScrollResponse(succeeded, numberOfFreedSearchContexts.get()));
-            }
-        }
-
-        void onFailedFreedContext(Throwable e, DiscoveryNode node) {
-            logger.warn((Supplier<?>) () -> new ParameterizedMessage("Clear SC failed on node[{}]", node), e);
-            if (expectedOps.countDown()) {
-                listener.onResponse(new ClearScrollResponse(false, numberOfFreedSearchContexts.get()));
-            } else {
-                expHolder.set(e);
-            }
-        }
-
+        Runnable runnable = new ClearScrollController(request, listener, clusterService.state().nodes(), logger, searchTransportService);
+        runnable.run();
     }
 
 }