Move headers collection to AsyncOperator

Author: ivancea

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/AsyncOperator.java

@@ -17,6 +17,7 @@
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.util.concurrent.ConcurrentCollections;
+import org.elasticsearch.common.util.concurrent.ThreadContext;
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.index.seqno.LocalCheckpointTracker;
@@ -34,6 +35,11 @@
  * to reduce communication overhead and fetches a {@code Fetched} at a time.
  * It's the responsibility of subclasses to transform that {@code Fetched} into
  * output.
+ * <p>
+ *     This operator will also take care of merging response headers from the thread context into the main thread,
+ *     which <b>must</b> be the that closes this.
+ * </p>
+ *
  * @see #performAsync(Page, ActionListener)
  */
 public abstract class AsyncOperator<Fetched> implements Operator {
@@ -45,6 +51,7 @@ public abstract class AsyncOperator<Fetched> implements Operator {
     private final DriverContext driverContext;
 
     private final int maxOutstandingRequests;
+    private final ResponseHeadersCollector responseHeadersCollector;
     private final LongAdder totalTimeInNanos = new LongAdder();
 
     private boolean finished = false;
@@ -66,9 +73,10 @@ public abstract class AsyncOperator<Fetched> implements Operator {
      *
      * @param maxOutstandingRequests the maximum number of outstanding requests
      */
-    public AsyncOperator(DriverContext driverContext, int maxOutstandingRequests) {
+    public AsyncOperator(DriverContext driverContext, ThreadContext threadContext, int maxOutstandingRequests) {
         this.driverContext = driverContext;
         this.maxOutstandingRequests = maxOutstandingRequests;
+        this.responseHeadersCollector = new ResponseHeadersCollector(threadContext);
     }
 
     @Override
@@ -97,6 +105,7 @@ public void addInput(Page input) {
             });
             final long startNanos = System.nanoTime();
             performAsync(input, ActionListener.runAfter(listener, () -> {
+                responseHeadersCollector.collect();
                 driverContext.removeAsyncAction();
                 totalTimeInNanos.add(System.nanoTime() - startNanos);
             }));
@@ -172,6 +181,7 @@ public final void close() {
         finish();
         closed = true;
         discardResults();
+        responseHeadersCollector.finish();
         doClose();
     }
 

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/enrich/EnrichLookupOperator.java

@@ -16,7 +16,6 @@
 import org.elasticsearch.compute.operator.AsyncOperator;
 import org.elasticsearch.compute.operator.DriverContext;
 import org.elasticsearch.compute.operator.Operator;
-import org.elasticsearch.compute.operator.ResponseHeadersCollector;
 import org.elasticsearch.core.CheckedFunction;
 import org.elasticsearch.tasks.CancellableTask;
 import org.elasticsearch.xcontent.XContentBuilder;
@@ -38,7 +37,6 @@ public final class EnrichLookupOperator extends AsyncOperator<Page> {
     private final String matchType;
     private final String matchField;
     private final List<NamedExpression> enrichFields;
-    private final ResponseHeadersCollector responseHeadersCollector;
     private final Source source;
     private long totalTerms = 0L;
 
@@ -101,7 +99,7 @@ public EnrichLookupOperator(
         List<NamedExpression> enrichFields,
         Source source
     ) {
-        super(driverContext, maxOutstandingRequests);
+        super(driverContext, enrichLookupService.getThreadContext(), maxOutstandingRequests);
         this.sessionId = sessionId;
         this.parentTask = parentTask;
         this.inputChannel = inputChannel;
@@ -112,7 +110,6 @@ public EnrichLookupOperator(
         this.matchField = matchField;
         this.enrichFields = enrichFields;
         this.source = source;
-        this.responseHeadersCollector = new ResponseHeadersCollector(enrichLookupService.getThreadContext());
     }
 
     @Override
@@ -135,11 +132,7 @@ protected void performAsync(Page inputPage, ActionListener<Page> listener) {
             }
             return inputPage.appendPage(pages.getFirst());
         };
-        enrichLookupService.lookupAsync(
-            request,
-            parentTask,
-            ActionListener.runBefore(listener.map(handleResponse), responseHeadersCollector::collect)
-        );
+        enrichLookupService.lookupAsync(request, parentTask, listener.map(handleResponse));
     }
 
     @Override
@@ -171,7 +164,6 @@ public String toString() {
     protected void doClose() {
         // TODO: Maybe create a sub-task as the parent task of all the lookup tasks
         // then cancel it when this operator terminates early (e.g., have enough result).
-        responseHeadersCollector.finish();
     }
 
     @Override

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/enrich/LookupFromIndexOperator.java

@@ -18,7 +18,6 @@
 import org.elasticsearch.compute.operator.DriverContext;
 import org.elasticsearch.compute.operator.IsBlockedResult;
 import org.elasticsearch.compute.operator.Operator;
-import org.elasticsearch.compute.operator.ResponseHeadersCollector;
 import org.elasticsearch.compute.operator.lookup.RightChunkedLeftJoin;
 import org.elasticsearch.core.Releasable;
 import org.elasticsearch.core.Releasables;
@@ -94,7 +93,6 @@ public Operator get(DriverContext driverContext) {
     private final String lookupIndex;
     private final String matchField;
     private final List<NamedExpression> loadFields;
-    private final ResponseHeadersCollector responseHeadersCollector;
     private final Source source;
     private long totalTerms = 0L;
     /**
@@ -120,7 +118,7 @@ public LookupFromIndexOperator(
         List<NamedExpression> loadFields,
         Source source
     ) {
-        super(driverContext, maxOutstandingRequests);
+        super(driverContext, lookupService.getThreadContext(), maxOutstandingRequests);
         this.sessionId = sessionId;
         this.configuration = configuration;
         this.parentTask = parentTask;
@@ -131,7 +129,6 @@ public LookupFromIndexOperator(
         this.matchField = matchField;
         this.loadFields = loadFields;
         this.source = source;
-        this.responseHeadersCollector = new ResponseHeadersCollector(lookupService.getThreadContext());
     }
 
     @Override
@@ -151,10 +148,7 @@ protected void performAsync(Page inputPage, ActionListener<OngoingJoin> listener
         lookupService.lookupAsync(
             request,
             parentTask,
-            ActionListener.runBefore(
-                listener.map(pages -> new OngoingJoin(new RightChunkedLeftJoin(inputPage, loadFields.size()), pages.iterator())),
-                responseHeadersCollector::collect
-            )
+            listener.map(pages -> new OngoingJoin(new RightChunkedLeftJoin(inputPage, loadFields.size()), pages.iterator()))
         );
     }
 
@@ -226,7 +220,6 @@ public IsBlockedResult isBlocked() {
     protected void doClose() {
         // TODO: Maybe create a sub-task as the parent task of all the lookup tasks
         // then cancel it when this operator terminates early (e.g., have enough result).
-        responseHeadersCollector.finish();
         Releasables.close(ongoing);
     }