Resume Driver on cancelled or early finished

Author: dnhatn

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

@@ -11,6 +11,7 @@
 import org.elasticsearch.action.support.ContextPreservingActionListener;
 import org.elasticsearch.action.support.SubscribableListener;
 import org.elasticsearch.common.util.concurrent.AbstractRunnable;
+import org.elasticsearch.common.util.concurrent.EsExecutors;
 import org.elasticsearch.common.util.concurrent.ThreadContext;
 import org.elasticsearch.compute.Describable;
 import org.elasticsearch.compute.data.Page;
@@ -74,10 +75,12 @@ public class Driver implements Releasable, Describable {
  private final long statusNanos;
 
  private final AtomicReference<String> cancelReason = new AtomicReference<>();
- private final AtomicReference<SubscribableListener<Void>> blocked = new AtomicReference<>();
 
  private final AtomicBoolean started = new AtomicBoolean();
+ private final AtomicBoolean earlyFinished = new AtomicBoolean();
  private final SubscribableListener<Void> completionListener = new SubscribableListener<>();
+ private final AtomicReference<SubscribableListener<Void>> blocked = new AtomicReference<>();
+ private final AtomicReference<Runnable> scheduledTask = new AtomicReference<>();
 
  /**
  * Status reported to the tasks API. We write the status at most once every
@@ -215,6 +218,23 @@ SubscribableListener<Void> run(TimeValue maxTime, int maxIterations, LongSupplie
  }
  }
 
+ private void earlyFinish() {
+ earlyFinished.set(true);
+ tryResumeDriver();
+ }
+
+ private void tryResumeDriver() {
+ // Attempt to run the scheduled task on the current thread to finish the driver by closing operators
+ SubscribableListener<Void> blockingFuture = blocked.get();
+ if (blockingFuture != null) {
+ blockingFuture.onResponse(null);
+ }
+ Runnable task = scheduledTask.getAndSet(null);
+ if (task != null) {
+ task.run();
+ }
+ }
+
  /**
  * Whether the driver has run the chain of operators to completion.
  */
@@ -245,31 +265,33 @@ private IsBlockedResult runSingleLoopIteration() {
  ensureNotCancelled();
  boolean movedPage = false;
 
- for (int i = 0; i < activeOperators.size() - 1; i++) {
- Operator op = activeOperators.get(i);
- Operator nextOp = activeOperators.get(i + 1);
+ if (earlyFinished.get() == false) {
+ for (int i = 0; i < activeOperators.size() - 1; i++) {
+ Operator op = activeOperators.get(i);
+ Operator nextOp = activeOperators.get(i + 1);
 
- // skip blocked operator
- if (op.isBlocked().listener().isDone() == false) {
- continue;
- }
+  // skip blocked operator
+  if (op.isBlocked().listener().isDone() == false) {
+  continue;
+  }
 
- if (op.isFinished() == false && nextOp.needsInput()) {
- Page page = op.getOutput();
- if (page == null) {
- // No result, just move to the next iteration
- } else if (page.getPositionCount() == 0) {
- // Empty result, release any memory it holds immediately and move to the next iteration
- page.releaseBlocks();
- } else {
- // Non-empty result from the previous operation, move it to the next operation
- nextOp.addInput(page);
- movedPage = true;
+ if (op.isFinished() == false && nextOp.needsInput()) {
+ Page page = op.getOutput();
+ if (page == null) {
+ // No result, just move to the next iteration
+ } else if (page.getPositionCount() == 0) {
+ // Empty result, release any memory it holds immediately and move to the next iteration
+ page.releaseBlocks();
+ } else {
+ // Non-empty result from the previous operation, move it to the next operation
+ nextOp.addInput(page);
+ movedPage = true;
+ }
  }
- }
 
- if (op.isFinished()) {
- nextOp.finish();
+ if (op.isFinished()) {
+ nextOp.finish();
+ }
  }
  }
 
@@ -312,19 +334,10 @@ private IsBlockedResult runSingleLoopIteration() {
 
  public void cancel(String reason) {
  if (cancelReason.compareAndSet(null, reason)) {
- synchronized (this) {
- SubscribableListener<Void> fut = this.blocked.get();
- if (fut != null) {
- fut.onFailure(new TaskCancelledException(reason));
- }
- }
+ tryResumeDriver();
  }
  }
 
- private boolean isCancelled() {
- return cancelReason.get() != null;
- }
-
  private void ensureNotCancelled() {
  String reason = cancelReason.get();
  if (reason != null) {
@@ -342,6 +355,12 @@ public static void start(
  driver.completionListener.addListener(listener);
  if (driver.started.compareAndSet(false, true)) {
  driver.updateStatus(0, 0, DriverStatus.Status.STARTING, "driver starting");
+ if (driver.activeOperators.isEmpty() == false) {
+ var onFinishedListener = driver.activeOperators.getLast().onFinishedListener();
+ if (onFinishedListener != null) {
+ onFinishedListener.addListener(ActionListener.running(driver::earlyFinish));
+ }
+ }
  schedule(DEFAULT_TIME_BEFORE_YIELDING, maxIterations, threadContext, executor, driver, driver.completionListener);
  }
  }
@@ -371,7 +390,7 @@ private static void schedule(
  Driver driver,
  ActionListener<Void> listener
  ) {
- executor.execute(new AbstractRunnable() {
+ final var task = new AbstractRunnable() {
 
  @Override
  protected void doRun() {
@@ -383,16 +402,12 @@ protected void doRun() {
  if (fut.isDone()) {
  schedule(maxTime, maxIterations, threadContext, executor, driver, listener);
  } else {
- synchronized (driver) {
- if (driver.isCancelled() == false) {
- driver.blocked.set(fut);
- }
- }
  ActionListener<Void> readyListener = ActionListener.wrap(
  ignored -> schedule(maxTime, maxIterations, threadContext, executor, driver, listener),
  this::onFailure
  );
  fut.addListener(ContextPreservingActionListener.wrapPreservingContext(readyListener, threadContext));
+ driver.blocked.set(fut);
  }
  }
 
@@ -405,6 +420,17 @@ public void onFailure(Exception e) {
  void onComplete(ActionListener<Void> listener) {
  driver.driverContext.waitForAsyncActions(ContextPreservingActionListener.wrapPreservingContext(listener, threadContext));
  }
+ };
+ final Runnable existing = driver.scheduledTask.getAndSet(task);
+ assert existing == null : existing;
+ final boolean runOnCurrentThread = driver.earlyFinished.get() || driver.cancelReason.get() != null;
+ final Executor executorToUse = runOnCurrentThread ? EsExecutors.DIRECT_EXECUTOR_SERVICE : executor;
+ executorToUse.execute(() -> {
+ final Runnable next = driver.scheduledTask.getAndSet(null);
+ if (next != null) {
+ assert next == task;
+ next.run();
+ }
  });
  }
 

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

@@ -13,6 +13,7 @@
 import org.elasticsearch.compute.Describable;
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.Page;
+import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.Releasable;
 import org.elasticsearch.xcontent.ToXContentObject;
 
@@ -62,6 +63,14 @@ public interface Operator extends Releasable {
  */
  boolean isFinished();
 
+ /**
+ * Returns a listener that is notified when the operator is finished. This returns null if the operator doesn't support this feature.
+ */
+ @Nullable
+ default SubscribableListener<Void> onFinishedListener() {
+ return null;
+ }
+
  /**
  * returns non-null if output page available. Only called when isFinished() == false
  * @throws UnsupportedOperationException if the operator is a {@link SinkOperator}

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

@@ -7,6 +7,7 @@
 
 package org.elasticsearch.compute.operator.exchange;
 
+import org.elasticsearch.action.support.SubscribableListener;
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.IsBlockedResult;
 
@@ -26,12 +27,12 @@ public interface ExchangeSink {
  void finish();
 
  /**
- * Whether the sink has received all pages
+ * Whether the sink is blocked on adding more pages
  */
- boolean isFinished();
+ IsBlockedResult waitForWriting();
 
  /**
- * Whether the sink is blocked on adding more pages
+ * Returns a listener that is notified when the sink is finished
  */
- IsBlockedResult waitForWriting();
+ SubscribableListener<Void> onFinished();
 }

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

@@ -52,9 +52,11 @@ public ExchangeSinkHandler(BlockFactory blockFactory, int maxBufferSize, LongSup
 
  private class ExchangeSinkImpl implements ExchangeSink {
  boolean finished;
+ private final SubscribableListener<Void> onFinished = new SubscribableListener<>();
 
  ExchangeSinkImpl() {
  onChanged();
+ buffer.addCompletionListener(onFinished);
  outstandingSinks.incrementAndGet();
  }
 
@@ -68,6 +70,7 @@ public void addPage(Page page) {
  public void finish() {
  if (finished == false) {
  finished = true;
+ onFinished.onResponse(null);
  onChanged();
  if (outstandingSinks.decrementAndGet() == 0) {
  buffer.finish(false);
@@ -77,8 +80,8 @@ public void finish() {
  }
 
  @Override
- public boolean isFinished() {
- return finished || buffer.isFinished();
+ public SubscribableListener<Void> onFinished() {
+ return onFinished;
  }
 
  @Override

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

@@ -9,6 +9,7 @@
 
 import org.elasticsearch.TransportVersion;
 import org.elasticsearch.TransportVersions;
+import org.elasticsearch.action.support.SubscribableListener;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.io.stream.StreamInput;
@@ -56,7 +57,12 @@ public ExchangeSinkOperator(ExchangeSink sink, Function<Page, Page> transformer)
 
  @Override
  public boolean isFinished() {
- return sink.isFinished();
+ return sink.onFinished().isDone();
+ }
+
+ @Override
+ public SubscribableListener<Void> onFinishedListener() {
+ return sink.onFinished();
  }
 
  @Override

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/DriverTests.java

@@ -9,6 +9,7 @@
 
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.ActionRunnable;
+import org.elasticsearch.action.support.PlainActionFuture;
 import org.elasticsearch.common.breaker.CircuitBreaker;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.unit.ByteSizeValue;
@@ -21,6 +22,10 @@
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.ElementType;
 import org.elasticsearch.compute.data.Page;
+import org.elasticsearch.compute.operator.exchange.ExchangeSinkHandler;
+import org.elasticsearch.compute.operator.exchange.ExchangeSinkOperator;
+import org.elasticsearch.compute.operator.exchange.ExchangeSourceHandler;
+import org.elasticsearch.compute.operator.exchange.ExchangeSourceOperator;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.threadpool.FixedExecutorBuilder;
@@ -35,8 +40,10 @@
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.CyclicBarrier;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Function;
 import java.util.function.LongSupplier;
 
+import static org.hamcrest.Matchers.either;
 import static org.hamcrest.Matchers.equalTo;
 
 public class DriverTests extends ESTestCase {
@@ -273,6 +280,33 @@ public Page getOutput() {
  }
  }
 
+ public void testResumeOnEarlyFinish() throws Exception {
+ DriverContext driverContext = driverContext();
+ ThreadPool threadPool = threadPool();
+ try {
+ PlainActionFuture<Void> sourceFuture = new PlainActionFuture<>();
+ var sourceHandler = new ExchangeSourceHandler(between(1, 5), threadPool.executor("esql"), sourceFuture);
+ var sinkHandler = new ExchangeSinkHandler(driverContext.blockFactory(), between(1, 5), System::currentTimeMillis);
+ var sourceOperator = new ExchangeSourceOperator(sourceHandler.createExchangeSource());
+ var sinkOperator = new ExchangeSinkOperator(sinkHandler.createExchangeSink(), Function.identity());
+ Driver driver = new Driver(driverContext, sourceOperator, List.of(), sinkOperator, () -> {});
+ PlainActionFuture<Void> future = new PlainActionFuture<>();
+ Driver.start(threadPool.getThreadContext(), threadPool.executor("esql"), driver, between(1, 1000), future);
+ assertBusy(
+ () -> assertThat(
+ driver.status().status(),
+ either(equalTo(DriverStatus.Status.ASYNC)).or(equalTo(DriverStatus.Status.STARTING))
+ )
+ );
+ sinkHandler.fetchPageAsync(true, ActionListener.noop());
+ future.actionGet(5, TimeUnit.SECONDS);
+ assertThat(driver.status().status(), equalTo(DriverStatus.Status.DONE));
+ sourceFuture.actionGet(5, TimeUnit.SECONDS);
+ } finally {
+ terminate(threadPool);
+ }
+ }
+
  private static void assertRunningWithRegularUser(ThreadPool threadPool) {
  String user = threadPool.getThreadContext().getHeader("user");
  assertThat(user, equalTo("user1"));

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/exchange/ExchangeServiceTests.java

@@ -129,7 +129,7 @@ public void testBasic() throws Exception {
  // sink buffer is full
  assertFalse(randomFrom(sink1, sink2).waitForWriting().listener().isDone());
  sink1.finish();
- assertTrue(sink1.isFinished());
+ assertTrue(sink1.onFinished().isDone());
  for (int i = 0; i < 5; i++) {
  assertBusy(() -> assertTrue(source.waitForReading().listener().isDone()));
  assertEquals(pages[2 + i], source.pollPage());
@@ -138,7 +138,7 @@ public void testBasic() throws Exception {
  assertFalse(source.waitForReading().listener().isDone());
  assertBusy(() -> assertTrue(sink2.waitForWriting().listener().isDone()));
  sink2.finish();
- assertTrue(sink2.isFinished());
+ assertTrue(sink2.onFinished().isDone());
  assertTrue(source.isFinished());
  assertFalse(sourceCompletion.isDone());
  source.finish();
@@ -440,7 +440,7 @@ public void testClosingSinks() {
  assertTrue(resp.finished());
  assertNull(resp.takePage());
  assertTrue(sink.waitForWriting().listener().isDone());
- assertTrue(sink.isFinished());
+ assertTrue(sink.onFinished().isDone());
  }
 
  public void testFinishEarly() throws Exception {