This is an automated email from the ASF dual-hosted git repository.

gortiz pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git


The following commit(s) were added to refs/heads/master by this push:
     new c6f3f791f5f Check HardLimitExecutor limit only once, on submit (#19864)
c6f3f791f5f is described below

commit c6f3f791f5fe4c7de10c074b2509870bec33cd86
Author: Gonzalo Ortiz Jaureguizar <[email protected]>
AuthorDate: Thu Oct 15 11:44:21 2026 +0200

    Check HardLimitExecutor limit only once, on submit (#19864)
    
    * Check HardLimitExecutor limit only once, on submit
    
    HardLimitExecutor.decorate checked the limit twice: synchronously on submit
    and again inside the wrapper right before the task ran. The second check
    could reject a task that had passed the first one when another task started
    in between. The wrapper then threw before the task ran. For
    OpChainSchedulerService this meant the ListenableFutureTask never ran, its
    completion callback never fired, the opchain was never closed, the
    per-request active counter never reached zero, and the downstream receiver
    waited until the query timed out. The exception was captured by the
    FutureTask that AbstractExecutorService.submit creates and nobody read it, 
so
    nothing was logged.
    
    Keep the synchronous check as the only rejection point. The wrapper now only
    increments and decrements the running counter. The counter can exceed the
    limit by the number of tasks accepted but not yet started, which is
    negligible on a cached pool.
    
    Add a regression test in pinot-spi with a deferred executor that controls
    when each task starts, and one in pinot-query-runtime that drives
    OpChainSchedulerService through the same race and asserts both opchains run,
    close, and release the request.
    
    * Clarify HardLimitExecutor overshoot bound for fixed thread pools
    
    The HardLimitExecutor Javadoc claimed the counter overshoot is negligible
    on a cached pool. That holds only when accepted tasks start promptly. On a
    fixed pool, accepted tasks wait in the queue, so the overshoot can be as
    large as the queue backlog. Reword the Javadoc to state this bound.
    
    ---------
    
    Co-authored-by: Gonzalo Ortiz <[email protected]>
---
 .../HardLimitExecutorOpChainRegressionTest.java    | 157 +++++++++++++++++++++
 .../pinot/spi/executor/HardLimitExecutor.java      |  14 +-
 .../pinot/spi/executor/HardLimitExecutorTest.java  |  85 +++++++++++
 3 files changed, 253 insertions(+), 3 deletions(-)

diff --git 
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/executor/HardLimitExecutorOpChainRegressionTest.java
 
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/executor/HardLimitExecutorOpChainRegressionTest.java
new file mode 100644
index 00000000000..ad5b8008286
--- /dev/null
+++ 
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/executor/HardLimitExecutorOpChainRegressionTest.java
@@ -0,0 +1,157 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.query.runtime.executor;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.AbstractExecutorService;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import org.apache.pinot.common.datatable.StatMap;
+import org.apache.pinot.query.mailbox.MailboxService;
+import org.apache.pinot.query.routing.StageMetadata;
+import org.apache.pinot.query.routing.WorkerMetadata;
+import org.apache.pinot.query.runtime.blocks.SuccessMseBlock;
+import org.apache.pinot.query.runtime.operator.MailboxSendOperator;
+import org.apache.pinot.query.runtime.operator.MultiStageOperator;
+import org.apache.pinot.query.runtime.operator.OpChain;
+import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
+import org.apache.pinot.spi.executor.HardLimitExecutor;
+import org.apache.pinot.spi.query.QueryThreadContext;
+import org.apache.pinot.spi.query.QueryThreadExceedStrategy;
+import org.mockito.Mockito;
+import org.testng.annotations.Test;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertTrue;
+
+
+/// Regression test for the opchain leak that [HardLimitExecutor] used to 
cause when it re-checked the limit at run
+/// time. The second check could reject a task that passed the submit-time 
check, throwing before the
+/// [OpChainSchedulerService] `ListenableFutureTask` ran. The completion 
callback that closes the opchain and
+/// decrements the per-request counter never fired, and the exception was 
swallowed by the `FutureTask` created by
+/// `AbstractExecutorService.submit`, so nothing was logged either.
+///
+/// Now the limit is checked only once, on submit. A task that passes that 
check always runs.
+public class HardLimitExecutorOpChainRegressionTest {
+
+  /// Executor that queues tasks instead of running them so the test controls 
when each one starts.
+  private static class DeferredExecutor extends AbstractExecutorService {
+    final List<Runnable> _queued = new ArrayList<>();
+
+    @Override
+    public void execute(Runnable command) {
+      _queued.add(command);
+    }
+
+    @Override
+    public void shutdown() {
+    }
+
+    @Override
+    public List<Runnable> shutdownNow() {
+      return List.of();
+    }
+
+    @Override
+    public boolean isShutdown() {
+      return false;
+    }
+
+    @Override
+    public boolean isTerminated() {
+      return false;
+    }
+
+    @Override
+    public boolean awaitTermination(long timeout, TimeUnit unit) {
+      return true;
+    }
+  }
+
+  private static OpChain getChain(MultiStageOperator operator) {
+    MailboxService mailboxService = mock(MailboxService.class);
+    when(mailboxService.getHostname()).thenReturn("localhost");
+    when(mailboxService.getPort()).thenReturn(1234);
+    WorkerMetadata workerMetadata = new WorkerMetadata(0, Map.of(), Map.of());
+    OpChainExecutionContext context = 
OpChainExecutionContext.fromQueryContext(mailboxService, Map.of(),
+        new StageMetadata(0, List.of(workerMetadata), Map.of()), 
workerMetadata, null, true, true);
+    return new OpChain(context, operator);
+  }
+
+  private static MultiStageOperator mockOperator() {
+    MultiStageOperator operator = Mockito.mock(MultiStageOperator.class);
+    when(operator.copyStatMaps()).thenAnswer(inv -> new 
StatMap<>(MailboxSendOperator.StatKey.class));
+    return operator;
+  }
+
+  @Test
+  public void acceptedOpChainAlwaysRunsAndCompletes()
+      throws Exception {
+    DeferredExecutor deferred = new DeferredExecutor();
+    ExecutorService limited = new HardLimitExecutor(1, 
QueryThreadContext.contextAwareExecutorService(deferred),
+        QueryThreadExceedStrategy.ERROR);
+    OpChainSchedulerService scheduler = new OpChainSchedulerService(limited);
+
+    CountDownLatch aStarted = new CountDownLatch(1);
+    CountDownLatch aRelease = new CountDownLatch(1);
+    MultiStageOperator operatorA = mockOperator();
+    when(operatorA.nextBlock()).thenAnswer(inv -> {
+      aStarted.countDown();
+      aRelease.await();
+      return SuccessMseBlock.INSTANCE;
+    });
+    MultiStageOperator operatorB = mockOperator();
+    when(operatorB.nextBlock()).thenReturn(SuccessMseBlock.INSTANCE);
+
+    // Both registrations pass the synchronous check: nothing has started yet, 
so _running == 0.
+    try (QueryThreadContext ignore = QueryThreadContext.openForMseTest()) {
+      scheduler.register(getChain(operatorA));
+      scheduler.register(getChain(operatorB));
+    }
+    assertEquals(deferred._queued.size(), 2, "both opchains must have been 
accepted by the executor");
+    assertEquals(scheduler.activeRequestCount(), 1);
+
+    // Start A: the HardLimitExecutor wrapper increments _running to 1 and A 
blocks inside nextBlock().
+    Thread threadA = new Thread(deferred._queued.get(0), "opchain-a");
+    threadA.start();
+    assertTrue(aStarted.await(10, TimeUnit.SECONDS), "A did not start");
+
+    // Start B: it passed the submit-time check, so it must run even though A 
pushed the counter to the limit.
+    Runnable taskB = deferred._queued.get(1);
+    taskB.run();
+    assertTrue(taskB instanceof Future, "AbstractExecutorService.submit wraps 
the task in a FutureTask");
+    ((Future<?>) taskB).get(10, TimeUnit.SECONDS);
+    verify(operatorB).nextBlock();
+    verify(operatorB).close();
+
+    // Let A finish normally.
+    aRelease.countDown();
+    threadA.join(TimeUnit.SECONDS.toMillis(10));
+    verify(operatorA).close();
+
+    assertEquals(scheduler.activeRequestCount(), 0, "every opchain completed, 
so no request may stay active");
+  }
+}
diff --git 
a/pinot-spi/src/main/java/org/apache/pinot/spi/executor/HardLimitExecutor.java 
b/pinot-spi/src/main/java/org/apache/pinot/spi/executor/HardLimitExecutor.java
index d73ca0aacf9..51719ebf18f 100644
--- 
a/pinot-spi/src/main/java/org/apache/pinot/spi/executor/HardLimitExecutor.java
+++ 
b/pinot-spi/src/main/java/org/apache/pinot/spi/executor/HardLimitExecutor.java
@@ -29,8 +29,17 @@ import org.slf4j.LoggerFactory;
 
 
 /// An Executor that allows a maximum of tasks running at the same time, 
rejecting immediately any excess.
+///
+/// The limit is checked once, synchronously, when the task is submitted. A 
task that passes that check is always
+/// run. The running counter is incremented when the task starts and 
decremented when it finishes, so the number of
+/// running tasks can exceed the limit by the number of tasks that were 
accepted but had not started yet when a new
+/// task was checked. The overshoot is negligible only when accepted tasks 
start promptly, as on a cached thread pool.
+/// On a fixed thread pool, accepted tasks can wait in the queue, so the 
overshoot can be as large as that backlog.
+///
+/// The limit is deliberately not re-checked when the task starts. A task that 
is accepted and then rejected at start
+/// time never runs and never completes its future. Callers that discard the 
future (for example a scheduler that
+/// relies on a completion callback) would never learn about the rejection and 
would leak the task's resources.
 public class HardLimitExecutor extends DecoratorExecutorService {
-
   private static final Logger LOGGER = 
LoggerFactory.getLogger(HardLimitExecutor.class);
 
   private final AtomicInteger _running;
@@ -74,6 +83,7 @@ public class HardLimitExecutor extends 
DecoratorExecutorService {
     }
   }
 
+  /// Checks the limit on the submitting thread. This is the only place where 
a task can be rejected.
   protected void checkTaskAllowed() {
     if (_running.get() >= _max) {
       if (_exceedStrategy == QueryThreadExceedStrategy.LOG) {
@@ -92,7 +102,6 @@ public class HardLimitExecutor extends 
DecoratorExecutorService {
   protected <T> Callable<T> decorate(Callable<T> task) {
     checkTaskAllowed();
     return () -> {
-      checkTaskAllowed();
       _running.getAndIncrement();
       try {
         return task.call();
@@ -106,7 +115,6 @@ public class HardLimitExecutor extends 
DecoratorExecutorService {
   protected Runnable decorate(Runnable task) {
     checkTaskAllowed();
     return () -> {
-      checkTaskAllowed();
       _running.getAndIncrement();
       try {
         task.run();
diff --git 
a/pinot-spi/src/test/java/org/apache/pinot/spi/executor/HardLimitExecutorTest.java
 
b/pinot-spi/src/test/java/org/apache/pinot/spi/executor/HardLimitExecutorTest.java
index 35834783d25..432d5360d6b 100644
--- 
a/pinot-spi/src/test/java/org/apache/pinot/spi/executor/HardLimitExecutorTest.java
+++ 
b/pinot-spi/src/test/java/org/apache/pinot/spi/executor/HardLimitExecutorTest.java
@@ -18,17 +18,23 @@
  */
 package org.apache.pinot.spi.executor;
 
+import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
+import java.util.concurrent.AbstractExecutorService;
 import java.util.concurrent.BrokenBarrierException;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.CyclicBarrier;
 import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
 import org.apache.pinot.spi.env.PinotConfiguration;
 import org.apache.pinot.spi.query.QueryThreadExceedStrategy;
 import org.apache.pinot.spi.utils.CommonConstants;
 import org.testng.annotations.Test;
 
 import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertTrue;
 import static org.testng.Assert.fail;
 
 
@@ -127,4 +133,83 @@ public class HardLimitExecutorTest {
     PinotConfiguration config5 = new PinotConfiguration(configMap5);
     assertEquals(HardLimitExecutor.getMultiStageExecutorHardLimit(config5), 
-1);
   }
+
+  /// A task that passes the submit-time check must run even if the limit is 
reached before it starts. The limit used
+  /// to be re-checked at start time, which silently dropped the task: the 
exception was captured by the FutureTask
+  /// created by AbstractExecutorService.submit and never reached the 
submitter.
+  @Test
+  public void testAcceptedTaskRunsEvenIfLimitIsReachedBeforeItStarts()
+      throws Exception {
+    DeferredExecutor deferred = new DeferredExecutor();
+    HardLimitExecutor ex = new HardLimitExecutor(1, deferred);
+    CountDownLatch firstStarted = new CountDownLatch(1);
+    CountDownLatch firstRelease = new CountDownLatch(1);
+    CountDownLatch secondRan = new CountDownLatch(1);
+
+    // Both submissions pass the synchronous check because nothing has started 
yet.
+    ex.execute(() -> {
+      firstStarted.countDown();
+      try {
+        firstRelease.await();
+      } catch (InterruptedException e) {
+        Thread.currentThread().interrupt();
+      }
+    });
+    ex.execute(secondRan::countDown);
+    assertEquals(deferred._queued.size(), 2);
+
+    Thread first = new Thread(deferred._queued.get(0));
+    first.start();
+    assertTrue(firstStarted.await(10, TimeUnit.SECONDS), "first task did not 
start");
+
+    // The running counter is at the limit now, but the second task was 
already accepted and must run.
+    deferred._queued.get(1).run();
+    assertTrue(secondRan.await(10, TimeUnit.SECONDS), "accepted task was 
dropped at start time");
+
+    // A third submission is rejected synchronously because the first task is 
still running.
+    try {
+      ex.execute(() -> {
+      });
+      fail("Should not allow more than 1 task");
+    } catch (IllegalStateException e) {
+      assertEquals(e.getMessage(), "Tasks limit exceeded.");
+    }
+
+    firstRelease.countDown();
+    first.join(TimeUnit.SECONDS.toMillis(10));
+  }
+
+  /// Executor that queues tasks instead of running them so the test controls 
when each one starts.
+  private static class DeferredExecutor extends AbstractExecutorService {
+    final List<Runnable> _queued = new ArrayList<>();
+
+    @Override
+    public void execute(Runnable command) {
+      _queued.add(command);
+    }
+
+    @Override
+    public void shutdown() {
+    }
+
+    @Override
+    public List<Runnable> shutdownNow() {
+      return List.of();
+    }
+
+    @Override
+    public boolean isShutdown() {
+      return false;
+    }
+
+    @Override
+    public boolean isTerminated() {
+      return false;
+    }
+
+    @Override
+    public boolean awaitTermination(long timeout, TimeUnit unit) {
+      return true;
+    }
+  }
 }


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to