making RestGetSampleStatsAction cancellable

Author: masseyke

server/src/main/java/org/elasticsearch/action/admin/indices/sampling/GetSampleStatsAction.java

@@ -22,19 +22,22 @@
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.io.stream.Writeable;
+import org.elasticsearch.common.xcontent.ChunkedToXContent;
 import org.elasticsearch.ingest.SamplingService;
 import org.elasticsearch.tasks.CancellableTask;
 import org.elasticsearch.tasks.Task;
 import org.elasticsearch.tasks.TaskId;
 import org.elasticsearch.transport.AbstractTransportRequest;
-import org.elasticsearch.xcontent.ToXContentObject;
-import org.elasticsearch.xcontent.XContentBuilder;
+import org.elasticsearch.xcontent.ToXContent;
 
 import java.io.IOException;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 
+import static org.elasticsearch.common.xcontent.ChunkedToXContentHelper.chunk;
+
 public class GetSampleStatsAction extends ActionType<GetSampleStatsAction.Response> {
 
     public static final GetSampleStatsAction INSTANCE = new GetSampleStatsAction();
@@ -131,7 +134,7 @@ public int hashCode() {
         }
     }
 
-    public static class Response extends BaseNodesResponse<GetSampleStatsAction.NodeResponse> implements Writeable, ToXContentObject {
+    public static class Response extends BaseNodesResponse<GetSampleStatsAction.NodeResponse> implements Writeable, ChunkedToXContent {
         final int maxSize;
 
         public Response(StreamInput in) throws IOException {
@@ -192,8 +195,8 @@ public int hashCode() {
         }
 
         @Override
-        public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
-            return getSampleStats().toXContent(builder, params);
+        public Iterator<? extends ToXContent> toXContentChunked(ToXContent.Params params) {
+            return chunk(getSampleStats());
         }
     }
 

server/src/main/java/org/elasticsearch/action/admin/indices/sampling/RestGetSampleStatsAction.java

@@ -13,7 +13,8 @@
 import org.elasticsearch.client.internal.node.NodeClient;
 import org.elasticsearch.rest.BaseRestHandler;
 import org.elasticsearch.rest.RestRequest;
-import org.elasticsearch.rest.action.RestToXContentListener;
+import org.elasticsearch.rest.action.RestCancellableNodeClient;
+import org.elasticsearch.rest.action.RestRefCountedChunkedToXContentListener;
 
 import java.io.IOException;
 import java.util.Arrays;
@@ -44,6 +45,10 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli
             );
         }
         GetSampleStatsAction.Request getSampleStatsRequest = new GetSampleStatsAction.Request(indexNames[0]);
-        return channel -> client.execute(GetSampleStatsAction.INSTANCE, getSampleStatsRequest, new RestToXContentListener<>(channel));
+        return channel -> new RestCancellableNodeClient(client, request.getHttpChannel()).execute(
+            GetSampleStatsAction.INSTANCE,
+            getSampleStatsRequest,
+            new RestRefCountedChunkedToXContentListener<>(channel)
+        );
     }
 }