Introduce compute listener

Author: dnhatn

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

@@ -21,13 +21,11 @@
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.index.seqno.LocalCheckpointTracker;
 import org.elasticsearch.index.seqno.SequenceNumbers;
-import org.elasticsearch.tasks.TaskCancelledException;
 import org.elasticsearch.xcontent.XContentBuilder;
 
 import java.io.IOException;
 import java.util.Map;
 import java.util.Objects;
-import java.util.concurrent.atomic.AtomicReference;
 import java.util.concurrent.atomic.LongAdder;
 
 /**
@@ -40,7 +38,7 @@ public abstract class AsyncOperator implements Operator {
     private volatile SubscribableListener<Void> blockedFuture;
 
     private final Map<Long, Page> buffers = ConcurrentCollections.newConcurrentMap();
-    private final AtomicReference<Exception> failure = new AtomicReference<>();
+    private final FailureCollector failure = new FailureCollector();
     private final DriverContext driverContext;
 
     private final int maxOutstandingRequests;
@@ -90,7 +88,7 @@ public void addInput(Page input) {
                 onSeqNoCompleted(seqNo);
             }, e -> {
                 releasePageOnAnyThread(input);
-                onFailure(e);
+                failure.unwrapAndCollect(e);
                 onSeqNoCompleted(seqNo);
             });
             final long startNanos = System.nanoTime();
@@ -121,25 +119,6 @@ private void releasePageOnAnyThread(Page page) {
 
     protected abstract void doClose();
 
-    private void onFailure(Exception e) {
-        failure.getAndUpdate(first -> {
-            if (first == null) {
-                return e;
-            }
-            // ignore subsequent TaskCancelledException exceptions as they don't provide useful info.
-            if (ExceptionsHelper.unwrap(e, TaskCancelledException.class) != null) {
-                return first;
-            }
-            if (ExceptionsHelper.unwrap(first, TaskCancelledException.class) != null) {
-                return e;
-            }
-            if (ExceptionsHelper.unwrapCause(first) != ExceptionsHelper.unwrapCause(e)) {
-                first.addSuppressed(e);
-            }
-            return first;
-        });
-    }
-
     private void onSeqNoCompleted(long seqNo) {
         checkpoint.markSeqNoAsProcessed(seqNo);
         if (checkpoint.getPersistedCheckpoint() < checkpoint.getProcessedCheckpoint()) {

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

@@ -0,0 +1,54 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.compute.operator;
+
+import org.elasticsearch.ElasticsearchException;
+import org.elasticsearch.ExceptionsHelper;
+import org.elasticsearch.tasks.TaskCancelledException;
+import org.elasticsearch.transport.TransportException;
+
+import java.util.concurrent.Semaphore;
+import java.util.concurrent.atomic.AtomicReference;
+
+/**
+ * Collects exceptions occurred in compute, specifically ignores trivial exceptions such as {@link TaskCancelledException},
+ * and selects the most appropriate exception to return to the caller.
+ * To bound the memory usage, this class only collects the first 10 non-trivial exceptions.
+ */
+public final class FailureCollector {
+    private final AtomicReference<Exception> exceptionRefs = new AtomicReference<>();
+    private final Semaphore exceptionPermits = new Semaphore(10);
+
+    public void unwrapAndCollect(Exception originEx) {
+        final Exception e = originEx instanceof TransportException
+            ? (originEx.getCause() instanceof Exception cause ? cause : new ElasticsearchException(originEx.getCause()))
+            : originEx;
+        exceptionRefs.getAndUpdate(first -> {
+            if (first == null) {
+                return e;
+            }
+            // ignore subsequent TaskCancelledException exceptions as they don't provide useful info.
+            if (ExceptionsHelper.unwrap(e, TaskCancelledException.class) != null) {
+                return first;
+            }
+            if (ExceptionsHelper.unwrap(first, TaskCancelledException.class) != null) {
+                return e;
+            }
+            if (ExceptionsHelper.unwrapCause(first) != ExceptionsHelper.unwrapCause(e)) {
+                if (exceptionPermits.tryAcquire()) {
+                    first.addSuppressed(e);
+                }
+            }
+            return first;
+        });
+    }
+
+    public Exception get() {
+        return exceptionRefs.get();
+    }
+}

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

@@ -7,21 +7,18 @@
 
 package org.elasticsearch.compute.operator.exchange;
 
-import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.ExceptionsHelper;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.support.RefCountingListener;
 import org.elasticsearch.action.support.SubscribableListener;
 import org.elasticsearch.common.util.concurrent.AbstractRunnable;
 import org.elasticsearch.compute.data.Page;
+import org.elasticsearch.compute.operator.FailureCollector;
 import org.elasticsearch.core.Releasable;
-import org.elasticsearch.tasks.TaskCancelledException;
-import org.elasticsearch.transport.TransportException;
 
 import java.util.List;
 import java.util.concurrent.Executor;
 import java.util.concurrent.atomic.AtomicInteger;
-import java.util.concurrent.atomic.AtomicReference;
 
 /**
  * An {@link ExchangeSourceHandler} asynchronously fetches pages and status from multiple {@link RemoteSink}s
@@ -37,7 +34,7 @@ public final class ExchangeSourceHandler {
 
     private final PendingInstances outstandingSinks;
     private final PendingInstances outstandingSources;
-    private final AtomicReference<Exception> failure = new AtomicReference<>();
+    private final FailureCollector failure = new FailureCollector();
 
     public ExchangeSourceHandler(int maxBufferSize, Executor fetchExecutor) {
         this.buffer = new ExchangeBuffer(maxBufferSize);
@@ -199,26 +196,8 @@ void fetchPage() {
             loopControl.exited();
         }
 
-        void onSinkFailed(Exception originEx) {
-            final Exception e = originEx instanceof TransportException
-                ? (originEx.getCause() instanceof Exception cause ? cause : new ElasticsearchException(originEx.getCause()))
-                : originEx;
-            failure.getAndUpdate(first -> {
-                if (first == null) {
-                    return e;
-                }
-                // ignore subsequent TaskCancelledException exceptions as they don't provide useful info.
-                if (ExceptionsHelper.unwrap(e, TaskCancelledException.class) != null) {
-                    return first;
-                }
-                if (ExceptionsHelper.unwrap(first, TaskCancelledException.class) != null) {
-                    return e;
-                }
-                if (ExceptionsHelper.unwrapCause(first) != ExceptionsHelper.unwrapCause(e)) {
-                    first.addSuppressed(e);
-                }
-                return first;
-            });
+        void onSinkFailed(Exception e) {
+            failure.unwrapAndCollect(e);
             buffer.waitForReading().onResponse(null); // resume the Driver if it is being blocked on reading
             onSinkComplete();
         }