ESQL: Fix missing Lookup Join Warnings (elastic#143585)

Fix a race condition where the Lookup Join Warnings are not always propagated to the StreamingLookupFromIndexOperator.

Closes elastic#143350, elastic#143378, elastic#143388, elastic#143427, elastic#143432, elastic#143434, elastic#143442, elastic#143443, elastic#143552, elastic#143553

Assisted by Cursor

Author: julian-elastic

muted-tests.yml

@@ -378,18 +378,9 @@ tests:
 - class: org.elasticsearch.xpack.esql.heap_attack.HeapAttackLookupJoinIT
   method: testLookupExplosionManyMatchesManyFields
   issue: https://github.com/elastic/elasticsearch/issues/143347
-- class: org.elasticsearch.xpack.esql.action.LookupJoinIT
-  method: testMultiValueJoinKeyWarnings
-  issue: https://github.com/elastic/elasticsearch/issues/143350
-- class: org.elasticsearch.xpack.esql.qa.mixed.MixedClusterEsqlSpecIT
-  method: test {csv-spec:lookup-join.LookupJoinWithMixPushableAndUnpushableFilters}
-  issue: https://github.com/elastic/elasticsearch/issues/143378
 - class: org.elasticsearch.snapshots.ConcurrentSnapshotsIT
   method: testBackToBackQueuedDeletes
   issue: https://github.com/elastic/elasticsearch/issues/143387
-- class: org.elasticsearch.xpack.esql.qa.mixed.MixedClusterEsqlSpecIT
-  method: test {csv-spec:lookup-join.MvJoinKeyOnTheLookupIndexAfterStats}
-  issue: https://github.com/elastic/elasticsearch/issues/143388
 - class: org.elasticsearch.xpack.search.CrossClusterAsyncSearchIT
   method: testCancelViaExpirationOnRemoteResultsWithMinimizeRoundtrips
   issue: https://github.com/elastic/elasticsearch/issues/143407
@@ -405,33 +396,12 @@ tests:
 - class: org.elasticsearch.xpack.esql.qa.mixed.EsqlClientYamlIT
   method: test {p0=esql/40_tsdb/TS Command grouping on text field}
   issue: https://github.com/elastic/elasticsearch/issues/142544
-- class: org.elasticsearch.xpack.esql.CsvIT
-  method: test {csv-spec:lookup-join.MvJoinKeyOnFromAfterStats}
-  issue: https://github.com/elastic/elasticsearch/issues/143427
 - class: org.elasticsearch.multiproject.test.CoreWithMultipleProjectsClientYamlTestSuiteIT
   method: test {yaml=search.retrievers/result-diversification/10_mmr_result_diversification_retriever/Test MMR result diversification single index float type}
   issue: https://github.com/elastic/elasticsearch/issues/143430
-- class: org.elasticsearch.xpack.esql.CsvIT
-  method: test {csv-spec:lookup-join.MvJoinKeyFromRow}
-  issue: https://github.com/elastic/elasticsearch/issues/143432
-- class: org.elasticsearch.xpack.esql.CsvIT
-  method: test {csv-spec:lookup-join.MvJoinKeyOnTheLookupIndexAfterStats}
-  issue: https://github.com/elastic/elasticsearch/issues/143434
-- class: org.elasticsearch.xpack.esql.CsvIT
-  method: test {csv-spec:lookup-join.MvJoinKeyOnTheLookupIndex}
-  issue: https://github.com/elastic/elasticsearch/issues/143442
-- class: org.elasticsearch.xpack.esql.CsvIT
-  method: test {csv-spec:lookup-join.LookupJoinWithCoalesceFilterOnRight}
-  issue: https://github.com/elastic/elasticsearch/issues/143443
 - class: org.elasticsearch.repositories.azure.AzureBlobContainerRetriesTests
   method: testWriteLargeBlob
   issue: https://github.com/elastic/elasticsearch/issues/143551
-- class: org.elasticsearch.xpack.esql.qa.single_node.EsqlSpecIT
-  method: test {csv-spec:lookup-join.MvJoinKeyOnTheLookupIndexAfterStats}
-  issue: https://github.com/elastic/elasticsearch/issues/143552
-- class: org.elasticsearch.xpack.esql.qa.single_node.EsqlSpecForceStoredLoadingIT
-  method: test {csv-spec:lookup-join.LookupJoinOnTwoFieldsAfterTop}
-  issue: https://github.com/elastic/elasticsearch/issues/143553
 
 # Examples:
 #

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

@@ -17,6 +17,7 @@
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.FailureCollector;
 import org.elasticsearch.compute.operator.IsBlockedResult;
+import org.elasticsearch.compute.operator.ResponseHeadersCollector;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.tasks.Task;
 import org.elasticsearch.tasks.TaskCancelledException;
@@ -118,6 +119,8 @@ public interface ServerSetupCallback {
     @Nullable
     private final BiConsumer<String, String> lookupPlanConsumer;
 
+    private final ResponseHeadersCollector responseHeadersCollector;
+
     /**
      * Create a new BidirectionalBatchExchangeClient.
      *
@@ -155,6 +158,7 @@ public BidirectionalBatchExchangeClient(
         this.lookupPlanConsumer = lookupPlanConsumer;
         this.maxWorkers = maxWorkers;
         this.serverNodeSupplier = serverNodeSupplier;
+        this.responseHeadersCollector = new ResponseHeadersCollector(transportService.getThreadPool().getThreadContext());
         logger.debug(
             "[LookupJoinClient] Created BidirectionalBatchExchangeClient: sharedExchangeId={}, maxBufferSize={}, maxWorkers={}",
             sharedExchangeId,
@@ -368,6 +372,7 @@ private void sendBatchExchangeStatusRequest(Worker worker) {
                 worker.serverToClientId,
                 executor,
                 ActionListener.<BatchExchangeStatusResponse>wrap(response -> {
+                    responseHeadersCollector.collect();
                     logger.debug(
                         "[LookupJoinClient] Received batch exchange status response for worker={}, success={}",
                         worker.workerId,
@@ -386,6 +391,7 @@ private void sendBatchExchangeStatusRequest(Worker worker) {
                         worker.statusRef.onFailure(failure);
                     }
                 }, failure -> {
+                    responseHeadersCollector.collect();
                     logger.error(
                         "[LookupJoinClient] Failed to receive batch exchange status response for worker={}: {}",
                         worker.workerId,
@@ -830,6 +836,15 @@ private void stop() {
         }
     }
 
+    /**
+     * Applies collected response headers (e.g. warnings from server-side drivers) to the current
+     * thread's context. Must be called on the coordinator driver thread (typically from the
+     * operator's {@code close()} method) so that {@code DriverRunner} can propagate them.
+     */
+    public void finishCollectingResponseHeaders() {
+        responseHeadersCollector.finish();
+    }
+
     /**
      * Closes the client. The driver guarantees all workers have completed (via
      * {@link #waitForServerResponse()} blocking on {@link #allWorkersCompleted}) before

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

@@ -20,6 +20,7 @@
 import org.elasticsearch.compute.operator.Driver;
 import org.elasticsearch.compute.operator.DriverContext;
 import org.elasticsearch.compute.operator.Operator;
+import org.elasticsearch.compute.operator.ResponseHeadersCollector;
 import org.elasticsearch.compute.operator.SinkOperator;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.Releasable;
@@ -68,6 +69,7 @@ public final class BidirectionalBatchExchangeServer extends BidirectionalBatchEx
     private final DiscoveryNode clientNode; // Client node for transport connection
     private PlainActionFuture<Void> driverFuture; // Future for driver completion
     private ThreadContext threadContext; // Thread context for starting driver
+    private ResponseHeadersCollector responseHeadersCollector;
     private volatile boolean driverPrepared = false; // Whether driver has been prepared but not started
     private volatile boolean driverStarted = false; // Whether driver has been started (client sent BatchExchangeStatusRequest)
     private ScheduledFuture<?> clientReadyTimeoutFuture; // Timeout for client to send BatchExchangeStatusRequest
@@ -269,6 +271,7 @@ private void onClientReady() {
     private ActionListener<Void> createDriverCompletionListener() {
         return ActionListener.wrap(ignored -> {
             logger.debug("[LookupJoinServer] Driver completion listener onResponse called (success) for exchangeId={}", serverToClientId);
+            responseHeadersCollector.collect();
             driverFuture.onResponse(null);
             logger.debug("[LookupJoinServer] Batch processing completed successfully for exchangeId={}", serverToClientId);
             // Close server resources BEFORE releasing the driver ref
@@ -292,6 +295,7 @@ private ActionListener<Void> createDriverCompletionListener() {
                 serverToClientId,
                 failure != null ? failure.getMessage() : "unknown"
             );
+            responseHeadersCollector.collect();
             // Complete the future first so close() won't throw
             driverFuture.onFailure(failure);
             // Close server resources BEFORE releasing the driver ref
@@ -313,6 +317,7 @@ private ActionListener<Void> createDriverCompletionListener() {
      * The listener is stored when BatchExchangeStatusRequest is received, before processing starts.
      */
     private void sendBatchExchangeStatusResponse(@Nullable Exception failure) {
+        responseHeadersCollector.finish();
         ActionListener<BatchExchangeStatusResponse> listener = batchExchangeStatusListener;
         if (listener != null) {
             logger.debug(
@@ -453,6 +458,7 @@ private void startBatchProcessing(
 
         // Store thread context for later driver startup
         this.threadContext = threadContext;
+        this.responseHeadersCollector = new ResponseHeadersCollector(threadContext);
 
         // Handler was already registered in initialize(), no need to register again
         logger.debug(

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

@@ -643,6 +643,7 @@ public void close() {
             } catch (Exception e) {
                 logger.error("Error finishing client", e);
             }
+            client.finishCollectingResponseHeaders();
             try {
                 client.close();
             } catch (Exception e) {