Changed all to work with Iterators and return them instead of calling sendChunks

Author: ivancea

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/stream/DefaultEsqlQueryResponseStream.java

@@ -36,7 +36,7 @@ public class DefaultEsqlQueryResponseStream extends EsqlQueryResponseStream {
     @Nullable
     private List<ColumnInfoImpl> columns;
 
-    private boolean dropNullColumns;
+    private final boolean dropNullColumns;
 
     DefaultEsqlQueryResponseStream(RestChannel restChannel, RestRequest restRequest, EsqlQueryRequest esqlRequest) throws IOException {
         super(restChannel, restRequest, esqlRequest);
@@ -50,7 +50,7 @@ protected boolean canBeStreamed() {
     }
 
     @Override
-    protected void doStartResponse(List<ColumnInfoImpl> columns) {
+    protected Iterator<? extends ToXContent> doStartResponse(List<ColumnInfoImpl> columns) {
         assert dropNullColumns == false : "this method doesn't support dropping null columns";
 
         this.columns = columns;
@@ -63,18 +63,18 @@ protected void doStartResponse(List<ColumnInfoImpl> columns) {
         // Start the values array, to be filled in
         content.add(ChunkedToXContentHelper.startArray("values"));
 
-        sendChunks(content);
+        return asIterator(content);
     }
 
     @Override
-    protected void doSendPages(Iterable<Page> pages) {
+    protected Iterator<? extends ToXContent> doSendPages(Iterable<Page> pages) {
         assert columns != null : "columns must be set before sending pages";
 
-        sendChunks(List.of(ResponseXContentUtils.rowValues(columns, pages, null)));
+        return ResponseXContentUtils.rowValues(columns, pages, null);
     }
 
     @Override
-    protected void doFinishResponse(EsqlQueryResponse response) {
+    protected Iterator<? extends ToXContent> doFinishResponse(EsqlQueryResponse response) {
         var content = new ArrayList<Iterator<? extends ToXContent>>(10);
 
         // End the values array
@@ -120,35 +120,36 @@ protected void doFinishResponse(EsqlQueryResponse response) {
 
         content.add(ChunkedToXContentHelper.endObject());
 
-        sendChunks(content);
+        return asIterator(content);
     }
 
     @Override
-    protected void doSendEverything(EsqlQueryResponse response) {
-        // TODO: Close the response
+    protected Iterator<? extends ToXContent> doSendEverything(EsqlQueryResponse response) {
         // final Releasable releasable = releasableFromResponse(esqlResponse);
 
         // TODO: Instead of sendChunks, implement a flush() to attach the response to it? Or pass the response to the methods
         // TODO: Or make "doX" methods return an Iterator<? extends ToXContent> and then concat them all together
 
+        var content = new ArrayList<Iterator<? extends ToXContent>>(3);
         boolean[] nullColumns = null;
         if (dropNullColumns) {
             nullColumns = nullColumns(response.columns(), response.pages());
-            sendStartResponseDroppingNullColumns(response.columns(), nullColumns);
+            content.add(sendStartResponseDroppingNullColumns(response.columns(), nullColumns));
         } else {
-            doStartResponse(response.columns());
+            content.add(doStartResponse(response.columns()));
         }
         // doSendPages doesn't work with nullColumns or columnar, so we generate them here directly
-        sendChunks(List.of(ResponseXContentUtils.columnValues(response.columns(), response.pages(), esqlRequest.columnar(), nullColumns)));
-        doFinishResponse(response);
+        content.add(ResponseXContentUtils.columnValues(response.columns(), response.pages(), esqlRequest.columnar(), nullColumns));
+        content.add(doFinishResponse(response));
+        return asIterator(content);
     }
 
     @Override
-    protected void doHandleException(Exception e) {
-        // TODO: Implement this
+    protected Iterator<? extends ToXContent> doHandleException(Exception e) {
+        throw new UnsupportedOperationException("Not implemented yet");
     }
 
-    private void sendStartResponseDroppingNullColumns(List<ColumnInfoImpl> columns, boolean[] nullColumns) {
+    private Iterator<? extends ToXContent> sendStartResponseDroppingNullColumns(List<ColumnInfoImpl> columns, boolean[] nullColumns) {
         assert dropNullColumns : "this method should only be called when dropping null columns";
 
         var content = new ArrayList<Iterator<? extends ToXContent>>(3);
@@ -157,7 +158,7 @@ private void sendStartResponseDroppingNullColumns(List<ColumnInfoImpl> columns,
         content.add(ResponseXContentUtils.allColumns(columns, "all_columns"));
         content.add(ResponseXContentUtils.nonNullColumns(columns, nullColumns, "columns"));
 
-        sendChunks(content);
+        return asIterator(content);
     }
 
     private boolean[] nullColumns(List<ColumnInfoImpl> columns, List<Page> pages) {

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/stream/EsqlQueryResponseStream.java

@@ -120,28 +120,35 @@ public final void startResponse(List<Attribute> schema) {
             return new ColumnInfoImpl(c.name(), c.dataType().outputType(), originalTypes);
         }).toList();
 
-        doStartResponse(columns);
+        sendChunks(doStartResponse(columns));
+
         initialStreamChunkSent = true;
     }
 
     public final void sendPages(Iterable<Page> pages) {
         assert finished == false : "sendPages() called on a finished stream";
 
         if (initialStreamChunkSent) {
-            doSendPages(pages);
+            sendChunks(doSendPages(pages));
         }
     }
 
     public final void finishResponse(EsqlQueryResponse response) {
         assert finished == false : "finishResponse() called more than once";
 
-        try (response) {
+        // TODO: Also, is this closing right? EsqlResponseListener uses releasableFromResponse(), which increments the ref first
+        boolean success = false;
+        try {
             if (initialStreamChunkSent) {
-                doFinishResponse(response);
+                sendChunks(doFinishResponse(response), response);
             } else {
-                doSendEverything(response);
+                sendChunks(doSendEverything(response), response);
             }
+            success = true;
         } finally {
+            if (success == false) {
+                response.close();
+            }
             finished = true;
         }
     }
@@ -152,6 +159,8 @@ public final void handleException(Exception e) {
         // TODO: To be overridden by subclasses. This should append the error to the stream, if possible
         LOGGER.error("Error while streaming response", e);
 
+        sendChunks(doHandleException(e));
+
         finished = true;
     }
 
@@ -185,25 +194,65 @@ public void onFailure(Exception e) {
      */
     protected abstract boolean canBeStreamed();
 
-    protected abstract void doStartResponse(List<ColumnInfoImpl> columns);
+    /**
+     * Returns the chunks to be sent at the beginning of the response. Called once, at the start.
+     * <p>
+     *     Only called if {@link #canBeStreamed()} returns {@code true}.
+     * </p>
+     */
+    protected abstract Iterator<? extends ToXContent> doStartResponse(List<ColumnInfoImpl> columns);
 
-    protected abstract void doSendPages(Iterable<Page> pages);
+    /**
+     * Returns the chunks for the given page. Called 0 to N times, after {@link #doStartResponse} and before {@link #doFinishResponse}.
+     * <p>
+     *     Only called if {@link #canBeStreamed()} returns {@code true}.
+     * </p>
+     */
+    protected abstract Iterator<? extends ToXContent> doSendPages(Iterable<Page> pages);
 
-    protected abstract void doFinishResponse(EsqlQueryResponse response);
+    /**
+     * Returns the remaining chunks of the response. Called once, at the end of the response.
+     * <p>
+     *     Only called if {@link #canBeStreamed()} returns {@code true}.
+     * </p>
+     */
+    protected abstract Iterator<? extends ToXContent> doFinishResponse(EsqlQueryResponse response);
 
-    protected abstract void doHandleException(Exception e);
+    /**
+     * Returns the chunks to be sent for the given exception.
+     * <p>
+     *     This may be called at any time, so the code should track what was sent already
+     *     and how to send a meaningful response given the chunks sent in previous calls.
+     * </p>
+     */
+    protected abstract Iterator<? extends ToXContent> doHandleException(Exception e);
 
-    protected void doSendEverything(EsqlQueryResponse response) {
-        // TODO: Is this safe? Should this be abstract to ensure proper implementation? Add tests for both cases
-        doStartResponse(response.columns());
-        doSendPages(response.pages());
-        doFinishResponse(response);
+    /**
+     * Returns the chunks of the full response. Called once for the full response.
+     * <p>
+     *     Only called if {@link #canBeStreamed()} returns {@code false}.
+     * </p>
+     */
+    protected Iterator<? extends ToXContent> doSendEverything(EsqlQueryResponse response) {
+        // TODO: Is this safe? Should this be abstract to ensure proper implementation? Add tests for both streamed and "everything" cases
+        return Iterators.concat(
+            doStartResponse(response.columns()),
+            doSendPages(response.pages()),
+            doFinishResponse(response)
+        );
     }
 
     @SuppressWarnings("unchecked")
-    protected final void sendChunks(List<Iterator<? extends ToXContent>> chunkedContent) {
-        // TODO: Maybe accept a single chunk here, and do a flush() inside of each method?
-        streamingXContentResponse.writeFragment(p0 -> Iterators.concat(chunkedContent.toArray(Iterator[]::new)), () -> {});
+    protected static Iterator<? extends ToXContent> asIterator(List<Iterator<? extends ToXContent>> chunks) {
+        return Iterators.concat(chunks.toArray(Iterator[]::new));
+    }
+
+    private void sendChunks(Iterator<? extends ToXContent> chunks) {
+        sendChunks(chunks, () -> {});
+    }
+
+    private void sendChunks(Iterator<? extends ToXContent> chunks, Releasable releasable) {
+        streamingXContentResponse.writeFragment(p0 -> chunks, releasable);
     }
 
     @Override

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/stream/NonStreamingEsqlQueryResponseStream.java

@@ -11,12 +11,15 @@
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.rest.RestChannel;
 import org.elasticsearch.rest.RestRequest;
+import org.elasticsearch.xcontent.ToXContent;
 import org.elasticsearch.xpack.esql.action.ColumnInfoImpl;
 import org.elasticsearch.xpack.esql.action.EsqlQueryRequest;
 import org.elasticsearch.xpack.esql.action.EsqlQueryResponse;
 import org.elasticsearch.xpack.esql.action.EsqlResponseListener;
 
 import java.io.IOException;
+import java.util.Collections;
+import java.util.Iterator;
 import java.util.List;
 
 /**
@@ -44,32 +47,34 @@ protected boolean canBeStreamed() {
     }
 
     @Override
-    protected void doStartResponse(List<ColumnInfoImpl> columns) {
+    protected Iterator<? extends ToXContent> doStartResponse(List<ColumnInfoImpl> columns) {
         throw new UnsupportedOperationException("This class does not support streaming");
     }
 
     @Override
-    protected void doSendPages(Iterable<Page> pages) {
+    protected Iterator<? extends ToXContent> doSendPages(Iterable<Page> pages) {
         throw new UnsupportedOperationException("This class does not support streaming");
     }
 
     @Override
-    protected void doFinishResponse(EsqlQueryResponse response) {
+    protected Iterator<? extends ToXContent> doFinishResponse(EsqlQueryResponse response) {
         throw new UnsupportedOperationException("This class does not support streaming");
     }
 
     @Override
-    protected void doHandleException(Exception e) {
+    protected Iterator<? extends ToXContent> doHandleException(Exception e) {
         listener.onFailure(e);
+        return Collections.emptyIterator();
     }
 
     @Override
-    protected void doSendEverything(EsqlQueryResponse response) {
+    protected Iterator<? extends ToXContent> doSendEverything(EsqlQueryResponse response) {
         // The base class will close the response, and the listener will do so too.
         // So we need to increment the reference count here
         // TODO: Check this
         response.incRef();
         listener.onResponse(response);
+        return Collections.emptyIterator();
     }
 
     @Override