Indices segments: bg serialize, make cancellable (#68965)

The response to an `IndicesSegmentsAction` might be large, perhaps 10s
of MBs of JSON, and today it is serialized on a transport thread. It
also might take so long to respond that the client times out, resulting
in the work needed to compute the response being wasted.

This commit introduces the `DispatchingRestToXContentListener` which
dispatches the work of serializing an `XContent` response to a
non-transport thread, and also makes `TransportBroadcastByNodeAction`
sensitive to the cancellability of its tasks.

It uses these two features to make the `RestIndicesSegmentsAction`
serialize its response on a `MANAGEMENT` thread, and to abort its work
more promptly if the client's channel is closed before the response is
sent.

Author: DaveCTurner

qa/smoke-test-http/src/test/java/org/elasticsearch/http/IndicesSegmentsRestCancellationIT.java

@@ -0,0 +1,181 @@
+/*
+ * 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.http;
+
+import org.apache.http.client.methods.HttpGet;
+import org.elasticsearch.action.admin.indices.segments.IndicesSegmentsAction;
+import org.elasticsearch.action.support.PlainActionFuture;
+import org.elasticsearch.client.Cancellable;
+import org.elasticsearch.client.Request;
+import org.elasticsearch.client.Response;
+import org.elasticsearch.client.ResponseListener;
+import org.elasticsearch.common.lease.Releasable;
+import org.elasticsearch.common.lease.Releasables;
+import org.elasticsearch.common.settings.Setting;
+import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.common.util.CollectionUtils;
+import org.elasticsearch.index.IndexService;
+import org.elasticsearch.index.IndexSettings;
+import org.elasticsearch.index.engine.Engine;
+import org.elasticsearch.index.engine.EngineConfig;
+import org.elasticsearch.index.engine.EngineException;
+import org.elasticsearch.index.engine.EngineFactory;
+import org.elasticsearch.index.engine.ReadOnlyEngine;
+import org.elasticsearch.index.shard.IndexShard;
+import org.elasticsearch.index.shard.IndexShardTestCase;
+import org.elasticsearch.index.translog.TranslogStats;
+import org.elasticsearch.indices.IndicesService;
+import org.elasticsearch.plugins.EnginePlugin;
+import org.elasticsearch.plugins.Plugin;
+import org.elasticsearch.tasks.CancellableTask;
+import org.elasticsearch.tasks.TaskInfo;
+import org.elasticsearch.transport.TransportService;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+import java.util.Optional;
+import java.util.concurrent.CancellationException;
+import java.util.concurrent.Semaphore;
+import java.util.function.Function;
+
+import static java.util.Collections.singletonList;
+import static org.hamcrest.Matchers.empty;
+import static org.hamcrest.Matchers.not;
+
+public class IndicesSegmentsRestCancellationIT extends HttpSmokeTestCase {
+
+    public static final Setting<Boolean> BLOCK_SEARCHER_SETTING
+            = Setting.boolSetting("index.block_searcher", false, Setting.Property.IndexScope);
+
+    @Override
+    protected Collection<Class<? extends Plugin>> nodePlugins() {
+        return CollectionUtils.appendToCopy(super.nodePlugins(), SearcherBlockingPlugin.class);
+    }
+
+    @Override
+    protected boolean addMockInternalEngine() {
+        return false;
+    }
+
+    public void testClusterStateRestCancellation() throws Exception {
+
+        createIndex("test", Settings.builder().put(BLOCK_SEARCHER_SETTING.getKey(), true).build());
+        ensureGreen("test");
+
+        final List<Semaphore> searcherBlocks = new ArrayList<>();
+        for (final IndicesService indicesService : internalCluster().getInstances(IndicesService.class)) {
+            for (final IndexService indexService : indicesService) {
+                for (final IndexShard indexShard : indexService) {
+                    final Engine engine = IndexShardTestCase.getEngine(indexShard);
+                    if (engine instanceof SearcherBlockingEngine) {
+                        searcherBlocks.add(((SearcherBlockingEngine) engine).searcherBlock);
+                    }
+                }
+            }
+        }
+        assertThat(searcherBlocks, not(empty()));
+
+        final List<Releasable> releasables = new ArrayList<>();
+        try {
+            for (final Semaphore searcherBlock : searcherBlocks) {
+                searcherBlock.acquire();
+                releasables.add(searcherBlock::release);
+            }
+
+            final Request indicesSegments = new Request(HttpGet.METHOD_NAME, "/_segments");
+
+            final PlainActionFuture<Void> future = new PlainActionFuture<>();
+            logger.info("--> sending indices segments request");
+            final Cancellable cancellable = getRestClient().performRequestAsync(indicesSegments, new ResponseListener() {
+                @Override
+                public void onSuccess(Response response) {
+                    future.onResponse(null);
+                }
+
+                @Override
+                public void onFailure(Exception exception) {
+                    future.onFailure(exception);
+                }
+            });
+
+            logger.info("--> waiting for task to start");
+            assertBusy(() -> {
+                final List<TaskInfo> tasks = client().admin().cluster().prepareListTasks().get().getTasks();
+                assertTrue(tasks.toString(), tasks.stream().anyMatch(t -> t.getAction().startsWith(IndicesSegmentsAction.NAME)));
+            });
+
+            logger.info("--> waiting for at least one task to hit a block");
+            assertBusy(() -> assertTrue(searcherBlocks.stream().anyMatch(Semaphore::hasQueuedThreads)));
+
+            logger.info("--> cancelling request");
+            cancellable.cancel();
+            expectThrows(CancellationException.class, future::actionGet);
+
+            logger.info("--> checking that all indices segments tasks are marked as cancelled");
+            assertBusy(() -> {
+                boolean foundTask = false;
+                for (TransportService transportService : internalCluster().getInstances(TransportService.class)) {
+                    for (CancellableTask cancellableTask : transportService.getTaskManager().getCancellableTasks().values()) {
+                        if (cancellableTask.getAction().startsWith(IndicesSegmentsAction.NAME)) {
+                            foundTask = true;
+                            assertTrue("task " + cancellableTask.getId() + " not cancelled", cancellableTask.isCancelled());
+                        }
+                    }
+                }
+                assertTrue("found no cancellable tasks", foundTask);
+            });
+        } finally {
+            Releasables.close(releasables);
+        }
+
+        logger.info("--> checking that all indices segments tasks have finished");
+        assertBusy(() -> {
+            final List<TaskInfo> tasks = client().admin().cluster().prepareListTasks().get().getTasks();
+            assertTrue(tasks.toString(), tasks.stream().noneMatch(t -> t.getAction().startsWith(IndicesSegmentsAction.NAME)));
+        });
+    }
+
+    public static class SearcherBlockingPlugin extends Plugin implements EnginePlugin {
+
+        @Override
+        public Optional<EngineFactory> getEngineFactory(IndexSettings indexSettings) {
+            if (BLOCK_SEARCHER_SETTING.get(indexSettings.getSettings())) {
+                return Optional.of(SearcherBlockingEngine::new);
+            }
+            return Optional.empty();
+        }
+
+        @Override
+        public List<Setting<?>> getSettings() {
+            return singletonList(BLOCK_SEARCHER_SETTING);
+        }
+    }
+
+    private static class SearcherBlockingEngine extends ReadOnlyEngine {
+
+        final Semaphore searcherBlock = new Semaphore(1);
+
+        SearcherBlockingEngine(EngineConfig config) {
+            super(config, null, new TranslogStats(), true, Function.identity(), true);
+        }
+
+        @Override
+        public Searcher acquireSearcher(String source, SearcherScope scope, Function<Searcher, Searcher> wrapper) throws EngineException {
+            try {
+                searcherBlock.acquire();
+            } catch (InterruptedException e) {
+                throw new AssertionError(e);
+            }
+            searcherBlock.release();
+            return super.acquireSearcher(source, scope, wrapper);
+        }
+    }
+
+}

server/src/main/java/org/elasticsearch/action/ActionModule.java

@@ -653,7 +653,7 @@ public void initRestHandlers(Supplier<DiscoveryNodes> nodesInCluster) {
         registerHandler.accept(new RestSnapshottableFeaturesAction());
         registerHandler.accept(new RestGetIndicesAction());
         registerHandler.accept(new RestIndicesStatsAction());
-        registerHandler.accept(new RestIndicesSegmentsAction());
+        registerHandler.accept(new RestIndicesSegmentsAction(threadPool));
         registerHandler.accept(new RestIndicesShardStoresAction());
         registerHandler.accept(new RestGetAliasesAction());
         registerHandler.accept(new RestIndexDeleteAliasesAction());

server/src/main/java/org/elasticsearch/action/admin/indices/cache/clear/TransportClearIndicesCacheAction.java

@@ -21,6 +21,7 @@
 import org.elasticsearch.common.inject.Inject;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.indices.IndicesService;
+import org.elasticsearch.tasks.Task;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.transport.TransportService;
 
@@ -63,7 +64,7 @@ protected ClearIndicesCacheRequest readRequestFrom(StreamInput in) throws IOExce
     }
 
     @Override
-    protected EmptyResult shardOperation(ClearIndicesCacheRequest request, ShardRouting shardRouting) {
+    protected EmptyResult shardOperation(ClearIndicesCacheRequest request, ShardRouting shardRouting, Task task) {
         indicesService.clearIndexShardCache(shardRouting.shardId(), request.queryCache(), request.fieldDataCache(), request.requestCache(),
             request.fields());
         return EmptyResult.INSTANCE;

server/src/main/java/org/elasticsearch/action/admin/indices/forcemerge/TransportForceMergeAction.java

@@ -22,6 +22,7 @@
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.index.shard.IndexShard;
 import org.elasticsearch.indices.IndicesService;
+import org.elasticsearch.tasks.Task;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.transport.TransportService;
 
@@ -62,7 +63,7 @@ protected ForceMergeRequest readRequestFrom(StreamInput in) throws IOException {
     }
 
     @Override
-    protected EmptyResult shardOperation(ForceMergeRequest request, ShardRouting shardRouting) throws IOException {
+    protected EmptyResult shardOperation(ForceMergeRequest request, ShardRouting shardRouting, Task task) throws IOException {
         IndexShard indexShard = indicesService.indexServiceSafe(shardRouting.shardId().getIndex()).getShard(shardRouting.shardId().id());
         indexShard.forceMerge(request);
         return EmptyResult.INSTANCE;

server/src/main/java/org/elasticsearch/action/admin/indices/recovery/TransportRecoveryAction.java

@@ -24,6 +24,7 @@
 import org.elasticsearch.index.shard.IndexShard;
 import org.elasticsearch.indices.IndicesService;
 import org.elasticsearch.indices.recovery.RecoveryState;
+import org.elasticsearch.tasks.Task;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.transport.TransportService;
 
@@ -86,7 +87,7 @@ protected RecoveryRequest readRequestFrom(StreamInput in) throws IOException {
     }
 
     @Override
-    protected RecoveryState shardOperation(RecoveryRequest request, ShardRouting shardRouting) {
+    protected RecoveryState shardOperation(RecoveryRequest request, ShardRouting shardRouting, Task task) {
         IndexService indexService = indicesService.indexServiceSafe(shardRouting.shardId().getIndex());
         IndexShard indexShard = indexService.getShard(shardRouting.shardId().id());
         return indexShard.recoveryState();

server/src/main/java/org/elasticsearch/action/admin/indices/segments/IndicesSegmentResponse.java

@@ -20,6 +20,7 @@
 import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.index.engine.Segment;
+import org.elasticsearch.transport.Transports;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -80,6 +81,8 @@ public void writeTo(StreamOutput out) throws IOException {
 
     @Override
     protected void addCustomXContentFields(XContentBuilder builder, Params params) throws IOException {
+        assert Transports.assertNotTransportThread("segments are very numerous, too expensive to serialize on a transport thread");
+
         builder.startObject(Fields.INDICES);
 
         for (IndexSegments indexSegments : getIndices().values()) {

server/src/main/java/org/elasticsearch/action/admin/indices/segments/IndicesSegmentsRequest.java

@@ -12,8 +12,12 @@
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.tasks.CancellableTask;
+import org.elasticsearch.tasks.Task;
+import org.elasticsearch.tasks.TaskId;
 
 import java.io.IOException;
+import java.util.Map;
 
 public class IndicesSegmentsRequest extends BroadcastRequest<IndicesSegmentsRequest> {
 
@@ -52,6 +56,10 @@ public void verbose(boolean v) {
     public void writeTo(StreamOutput out) throws IOException {
         super.writeTo(out);
         out.writeBoolean(verbose);
+    }
 
+    @Override
+    public Task createTask(long id, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
+        return new CancellableTask(id, type, action, "", parentTaskId, headers);
     }
 }

server/src/main/java/org/elasticsearch/action/admin/indices/segments/TransportIndicesSegmentsAction.java

@@ -23,6 +23,8 @@
 import org.elasticsearch.index.IndexService;
 import org.elasticsearch.index.shard.IndexShard;
 import org.elasticsearch.indices.IndicesService;
+import org.elasticsearch.tasks.CancellableTask;
+import org.elasticsearch.tasks.Task;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.transport.TransportService;
 
@@ -80,7 +82,8 @@ protected IndicesSegmentsRequest readRequestFrom(StreamInput in) throws IOExcept
     }
 
     @Override
-    protected ShardSegments shardOperation(IndicesSegmentsRequest request, ShardRouting shardRouting) {
+    protected ShardSegments shardOperation(IndicesSegmentsRequest request, ShardRouting shardRouting, Task task) {
+        assert task instanceof CancellableTask;
         IndexService indexService = indicesService.indexServiceSafe(shardRouting.index());
         IndexShard indexShard = indexService.getShard(shardRouting.id());
         return new ShardSegments(indexShard.routingEntry(), indexShard.segments(request.verbose()));

server/src/main/java/org/elasticsearch/action/admin/indices/stats/TransportIndicesStatsAction.java

@@ -28,6 +28,7 @@
 import org.elasticsearch.index.shard.IndexShard;
 import org.elasticsearch.index.shard.ShardNotFoundException;
 import org.elasticsearch.indices.IndicesService;
+import org.elasticsearch.tasks.Task;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.transport.TransportService;
 
@@ -83,7 +84,7 @@ protected IndicesStatsRequest readRequestFrom(StreamInput in) throws IOException
     }
 
     @Override
-    protected ShardStats shardOperation(IndicesStatsRequest request, ShardRouting shardRouting) {
+    protected ShardStats shardOperation(IndicesStatsRequest request, ShardRouting shardRouting, Task task) {
         IndexService indexService = indicesService.indexServiceSafe(shardRouting.shardId().getIndex());
         IndexShard indexShard = indexService.getShard(shardRouting.shardId().id());
         // if we don't have the routing entry yet, we need it stats wise, we treat it as if the shard is not ready yet