Use existing ResponseHeadersCollector instead of manually taking warnings

Author: ivancea

server/src/main/java/org/elasticsearch/common/logging/HeaderWarning.java

@@ -16,7 +16,6 @@
 import java.nio.charset.Charset;
 import java.nio.charset.StandardCharsets;
 import java.util.BitSet;
-import java.util.HashSet;
 import java.util.Iterator;
 import java.util.Locale;
 import java.util.Objects;
@@ -409,17 +408,6 @@ public static void addWarning(String message, Object... params) {
         addWarning(THREAD_CONTEXT, message, params);
     }
 
-    /**
-     * Adds a warning header without any formatting or prefix
-     */
-    public static void addRawWarning(String warningHeader) {
-        addRawWarning(THREAD_CONTEXT, warningHeader);
-    }
-
-    public static Set<String> getWarnings() {
-        return getWarnings(THREAD_CONTEXT);
-    }
-
     // package scope for testing
     static void addWarning(Set<ThreadContext> threadContexts, String message, Object... params) {
         final Iterator<ThreadContext> iterator = threadContexts.iterator();
@@ -444,19 +432,4 @@ static void addRawWarning(Set<ThreadContext> threadContexts, String warningHeade
             threadContext.addResponseHeader("Warning", warningHeader);
         }
     }
-
-    static Set<String> getWarnings(Set<ThreadContext> threadContexts) {
-        final Iterator<ThreadContext> iterator = threadContexts.iterator();
-        if (iterator.hasNext()) {
-            final Set<String> warnings = new HashSet<>();
-            while (iterator.hasNext()) {
-                final ThreadContext next = iterator.next();
-                final Set<String> contextWarnings = next.getResponseHeader("Warning");
-                warnings.addAll(contextWarnings);
-
-            }
-            return warnings;
-        }
-        return Set.of();
-    }
 }

server/src/main/java/org/elasticsearch/common/util/concurrent/ThreadContext.java

@@ -581,23 +581,6 @@ public Map<String, List<String>> getResponseHeaders() {
         return Collections.unmodifiableMap(map);
     }
 
-    /**
-     * Get a copy of the specified <em>response</em> header.
-     *
-     * @param header The header to get.
-     * @return Never {@code null}.
-     */
-    public Set<String> getResponseHeader(String header) {
-        Map<String, Set<String>> responseHeaders = threadLocal.get().responseHeaders;
-        Set<String> values = responseHeaders.get(header);
-
-        if (values == null) {
-            return Set.of();
-        }
-
-        return Set.copyOf(values);
-    }
-
     /**
      * Copies all header key, value pairs into the current context
      */

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

@@ -16,7 +16,6 @@
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
-import org.elasticsearch.common.logging.HeaderWarning;
 import org.elasticsearch.common.util.concurrent.ConcurrentCollections;
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.core.TimeValue;
@@ -27,7 +26,6 @@
 import java.io.IOException;
 import java.util.Map;
 import java.util.Objects;
-import java.util.Set;
 import java.util.concurrent.atomic.LongAdder;
 
 /**
@@ -43,7 +41,6 @@ public abstract class AsyncOperator<Fetched> implements Operator {
     private volatile SubscribableListener<Void> blockedFuture;
 
     private final Map<Long, Fetched> buffers = ConcurrentCollections.newConcurrentMap();
-    private final Set<String> storedWarningHeaders = ConcurrentCollections.newConcurrentSet();
     private final FailureCollector failureCollector = new FailureCollector();
     private final DriverContext driverContext;
 
@@ -92,11 +89,9 @@ public void addInput(Page input) {
         try {
             final ActionListener<Fetched> listener = ActionListener.wrap(output -> {
                 buffers.put(seqNo, output);
-                storeWarnings();
                 onSeqNoCompleted(seqNo);
             }, e -> {
                 releasePageOnAnyThread(input);
-                storeWarnings();
                 failureCollector.unwrapAndCollect(e);
                 onSeqNoCompleted(seqNo);
             });
@@ -189,7 +184,6 @@ public void finish() {
     public boolean isFinished() {
         if (finished && checkpoint.getPersistedCheckpoint() == checkpoint.getMaxSeqNo()) {
             checkFailure();
-            restoreWarnings();
             return true;
         } else {
             return false;
@@ -203,7 +197,6 @@ public boolean isFinished() {
     public final Fetched fetchFromBuffer() {
         checkFailure();
         long persistedCheckpoint = checkpoint.getPersistedCheckpoint();
-        restoreWarnings();
         if (persistedCheckpoint < checkpoint.getProcessedCheckpoint()) {
             persistedCheckpoint++;
             Fetched result = buffers.remove(persistedCheckpoint);
@@ -214,18 +207,6 @@ public final Fetched fetchFromBuffer() {
         }
     }
 
-    private void storeWarnings() {
-        Set<String> warnings = HeaderWarning.getWarnings();
-        storedWarningHeaders.addAll(warnings);
-    }
-
-    private void restoreWarnings() {
-        for (String warning : storedWarningHeaders) {
-            HeaderWarning.addRawWarning(warning);
-        }
-        // TODO: Remove warnings?
-    }
-
     @Override
     public IsBlockedResult isBlocked() {
         // TODO: Add an exchange service between async operation instead?

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

@@ -18,6 +18,7 @@
 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;
@@ -93,6 +94,7 @@ 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;
     /**
@@ -129,6 +131,7 @@ public LookupFromIndexOperator(
         this.matchField = matchField;
         this.loadFields = loadFields;
         this.source = source;
+        this.responseHeadersCollector = new ResponseHeadersCollector(lookupService.getThreadContext());
     }
 
     @Override
@@ -148,7 +151,10 @@ protected void performAsync(Page inputPage, ActionListener<OngoingJoin> listener
         lookupService.lookupAsync(
             request,
             parentTask,
-            listener.map(pages -> new OngoingJoin(new RightChunkedLeftJoin(inputPage, loadFields.size()), pages.iterator()))
+            ActionListener.runBefore(
+                listener.map(pages -> new OngoingJoin(new RightChunkedLeftJoin(inputPage, loadFields.size()), pages.iterator())),
+                responseHeadersCollector::collect
+            )
         );
     }
 
@@ -220,6 +226,7 @@ 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);
     }