Fix missing expired query callbacks and improve secondary workload queue thread wakeup

pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/AbstractSchedulerGroup.java

@@ -19,7 +19,9 @@
 package org.apache.pinot.core.query.scheduler;
 
 import com.google.common.base.Preconditions;
+import java.util.ArrayList;
 import java.util.Iterator;
+import java.util.List;
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.atomic.AtomicInteger;
 
@@ -64,14 +66,20 @@ public SchedulerQueryContext removeFirst() {
   }
 
   @Override
-  public void trimExpired(long deadlineMillis) {
+  public List<SchedulerQueryContext> trimExpired(long deadlineMillis) {
+    List<SchedulerQueryContext> expired = List.of();
     Iterator<SchedulerQueryContext> iter = _pendingQueries.iterator();
     while (iter.hasNext()) {
       SchedulerQueryContext next = iter.next();
       if (next.getArrivalTimeMs() < deadlineMillis) {
         iter.remove();
+        if (expired.isEmpty()) {
+          expired = new ArrayList<>();
+        }
+        expired.add(next);
       }
     }
+    return expired;
   }
 
   @Override

pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/BinaryWorkloadScheduler.java

@@ -35,6 +35,7 @@
 import org.apache.pinot.core.query.scheduler.resources.QueryExecutorService;
 import org.apache.pinot.spi.accounting.ThreadAccountant;
 import org.apache.pinot.spi.env.PinotConfiguration;
+import org.apache.pinot.spi.exception.QueryErrorCode;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -72,7 +73,7 @@ public BinaryWorkloadScheduler(PinotConfiguration config, String instanceId, Que
     super(config, instanceId, queryExecutor, threadAccountant, latestQueryTime,
         new BinaryWorkloadResourceManager(config));
 
-    _secondaryQueryQ = new SecondaryWorkloadQueue(config, _resourceManager);
+    _secondaryQueryQ = new SecondaryWorkloadQueue(config, _resourceManager, this::onQueryExpired);
     _numSecondaryRunners = config.getProperty(MAX_SECONDARY_QUERIES, DEFAULT_MAX_SECONDARY_QUERIES);
     LOGGER.info("numSecondaryRunners={}", _numSecondaryRunners);
     _secondaryRunnerSemaphore = new Semaphore(_numSecondaryRunners);
@@ -146,6 +147,9 @@ public void run() {
           try {
             SchedulerQueryContext request = _secondaryQueryQ.take();
             if (request == null) {
+              // mirrors query completion callback functionality.
+              _secondaryRunnerSemaphore.release();
+              checkStopResourceManager();
               continue;
             }
             ServerQueryRequest queryRequest = request.getQueryRequest();
@@ -157,6 +161,7 @@ public void run() {
               public void run() {
                 executorService.releaseWorkers();
                 schedulerGroup.endQuery();
+                _secondaryQueryQ.signalWorkersReleased();
                 _secondaryRunnerSemaphore.release();
                 checkStopResourceManager();
               }
@@ -209,8 +214,27 @@ private void checkStopResourceManager() {
   synchronized private void failAllPendingQueries() {
     List<SchedulerQueryContext> pending = _secondaryQueryQ.drain();
     for (SchedulerQueryContext queryContext : pending) {
-      ListenableFuture<byte[]> serverShuttingDown = shuttingDown(queryContext.getQueryRequest());
-      queryContext.setResultFuture(serverShuttingDown);
+      try {
+        ListenableFuture<byte[]> serverShuttingDown = shuttingDown(queryContext.getQueryRequest());
+        queryContext.setResultFuture(serverShuttingDown);
+      } catch (Throwable t) {
+        LOGGER.error("Failed to fail pending query: {} on shutdown", queryContext.getQueryRequest().getRequestId(), t);
+        // Last resort, so one bad query does not leave the remaining ones without a response.
+        // No-op if the handler already completed the future.
+        if (queryContext.getResultFuture() != null) {
+          queryContext.getResultFuture().setException(t);
+        }
+      }
     }
   }
+
+  private void onQueryExpired(SchedulerQueryContext queryContext) {
+    ServerQueryRequest queryRequest = queryContext.getQueryRequest();
+    long waitMs = System.currentTimeMillis() - queryRequest.getTimerContext().getQueryArrivalTimeMs();
+    LOGGER.warn("Dropping secondary query: {} for table: {} after waiting: {}ms in the queue",
+        queryRequest.getRequestId(), queryRequest.getTableNameWithType(), waitMs);
+    _serverMetrics.addMeteredTableValue(queryRequest.getTableNameWithType(), ServerMeter.SCHEDULING_TIMEOUT_EXCEPTIONS,
+        1L);
+    queryContext.setResultFuture(immediateErrorResponse(queryRequest, QueryErrorCode.QUERY_SCHEDULING_TIMEOUT));
+  }
 }

pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/SchedulerGroup.java

@@ -18,6 +18,9 @@
  */
 package org.apache.pinot.core.query.scheduler;
 
+import java.util.List;
+
+
 /// Scheduler group is a sub-queue in multi-level scheduling queues.
 /// This class maintains context information for each of the scheduling
 /// queues. Each SchedulerGroup is a queue of requests and related accounting.
@@ -41,7 +44,8 @@ public interface SchedulerGroup extends SchedulerGroupAccountant {
   SchedulerQueryContext removeFirst();
 
   /// Remove all the pending queries with arrival time earlier than the deadline
-  void trimExpired(long deadlineMillis);
+  /// @return the removed queries, so the caller can complete their result futures. Never null.
+  List<SchedulerQueryContext> trimExpired(long deadlineMillis);
 
   /// @return true if there are no pending queries for this group
   boolean isEmpty();

pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/SecondaryWorkloadQueue.java

@@ -26,6 +26,7 @@
 import java.util.concurrent.locks.Condition;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantLock;
+import java.util.function.Consumer;
 import javax.annotation.Nullable;
 import org.apache.pinot.core.query.request.ServerQueryRequest;
 import org.apache.pinot.core.query.scheduler.fcfs.FCFSSchedulerGroup;
@@ -58,10 +59,13 @@ public class SecondaryWorkloadQueue {
   private final Condition _queryReaderCondition = _queueLock.newCondition();
   private final ResourceManager _resourceManager;
   private final int _queryDeadlineMs;
+  private final Consumer<SchedulerQueryContext> _expiredQueryHandler;
 
-  public SecondaryWorkloadQueue(PinotConfiguration config, ResourceManager resourceManager) {
+  public SecondaryWorkloadQueue(PinotConfiguration config, ResourceManager resourceManager,
+      Consumer<SchedulerQueryContext> expiredQueryHandler) {
     Preconditions.checkNotNull(config);
     Preconditions.checkNotNull(resourceManager);
+    Preconditions.checkNotNull(expiredQueryHandler);
 
     _queryDeadlineMs =
         config.getProperty(SECONDARY_QUEUE_QUERY_TIMEOUT, DEFAULT_SECONDARY_QUEUE_QUERY_TIMEOUT_SEC) * 1000;
@@ -71,6 +75,7 @@ public SecondaryWorkloadQueue(PinotConfiguration config, ResourceManager resourc
         _maxPendingPerGroup);
     _schedulerGroup = new FCFSSchedulerGroup(SECONDARY_WORKLOAD_GROUP_NAME);
     _resourceManager = resourceManager;
+    _expiredQueryHandler = expiredQueryHandler;
   }
 
   /// Adds a query to the secondary workload queue.
@@ -94,21 +99,52 @@ public void put(SchedulerQueryContext query)
   /// @return
   @Nullable
   public SchedulerQueryContext take() {
-    _queueLock.lock();
-    try {
-      while (true) {
-        SchedulerQueryContext schedulerQueryContext;
-        while ((schedulerQueryContext = takeNextInternal()) == null) {
+    while (true) {
+      SchedulerQueryContext schedulerQueryContext = null;
+      List<SchedulerQueryContext> expired;
+      _queueLock.lock();
+      try {
+        expired = _schedulerGroup.trimExpired(System.currentTimeMillis() - _queryDeadlineMs);
+        if (!_schedulerGroup.isEmpty() && _resourceManager.canSchedule(_schedulerGroup)) {
+          schedulerQueryContext = _schedulerGroup.removeFirst();
+        } else if (expired.isEmpty()) {
+          // Nothing to dispatch and nothing to clean up, so wait for a signal.
           try {
-            _queryReaderCondition.await(_wakeUpTimeMs, TimeUnit.MILLISECONDS);
+            if (_schedulerGroup.isEmpty()) {
+              _queryReaderCondition.await(); // woken by put()
+            } else {
+              _queryReaderCondition.await(_wakeUpTimeMs, TimeUnit.MILLISECONDS); // TTL sweep safety net
+            }
           } catch (InterruptedException e) {
             return null;
           }
         }
+      } finally {
+        _queueLock.unlock();
+      }
+
+      // Complete expired queries outside the lock, since the handler writes to the request channel.
+      // A failure must not escape: it would drop the query dequeued above and leak the caller's runner permit.
+      for (SchedulerQueryContext queryContext : expired) {
+        try {
+          _expiredQueryHandler.accept(queryContext);
+        } catch (Throwable t) {
+          LOGGER.error("Failed to complete expired query: {}", queryContext.getQueryRequest().getRequestId(), t);
+          // Last resort, so the request always gets a response. No-op if the handler already completed the future.
+          if (queryContext.getResultFuture() != null) {
+            queryContext.getResultFuture().setException(t);
+          }
+        }
+      }
+      if (schedulerQueryContext != null) {
+        if (LOGGER.isDebugEnabled()) {
+          ServerQueryRequest queryRequest = schedulerQueryContext.getQueryRequest();
+          LOGGER.debug("Scheduling query: {} for group: {} with arrivalTimeMs: {} and numSegments: {}",
+              queryRequest.getRequestId(), _schedulerGroup.name(),
+              queryRequest.getTimerContext().getQueryArrivalTimeMs(), queryRequest.getSegmentsToQuery().size());
+        }
         return schedulerQueryContext;
       }
-    } finally {
-      _queueLock.unlock();
     }
   }
 
@@ -125,28 +161,6 @@ public List<SchedulerQueryContext> drain() {
     return pending;
   }
 
-  private SchedulerQueryContext takeNextInternal() {
-    long startTimeMs = System.currentTimeMillis();
-    long deadlineEpochMillis = startTimeMs - _queryDeadlineMs;
-
-    _schedulerGroup.trimExpired(deadlineEpochMillis);
-    if (_schedulerGroup.isEmpty() || !_resourceManager.canSchedule(_schedulerGroup)) {
-      return null;
-    }
-
-    if (LOGGER.isDebugEnabled()) {
-      StringBuilder sb = new StringBuilder("SchedulerInfo:");
-      sb.append(_schedulerGroup.toString());
-      ServerQueryRequest queryRequest = _schedulerGroup.peekFirst().getQueryRequest();
-      sb.append(" Group: " + _schedulerGroup.name() + ": [" + queryRequest.getTimerContext().getQueryArrivalTimeMs()
-          + "," + queryRequest.getRequestId() + "," + queryRequest.getSegmentsToQuery().size() + "," + startTimeMs
-          + "]");
-      LOGGER.debug(sb.toString());
-    }
-
-    return _schedulerGroup.removeFirst();
-  }
-
   private void checkSchedulerGroupCapacity(SchedulerQueryContext query)
       throws OutOfCapacityException {
     if (_schedulerGroup.numPending() >= _maxPendingPerGroup
@@ -158,4 +172,14 @@ private void checkSchedulerGroupCapacity(SchedulerQueryContext query)
               + _resourceManager.getTableThreadsHardLimit());
     }
   }
+
+  /// Signals the reader that reserved threads were released, so canSchedule() may now pass.
+  public void signalWorkersReleased() {
+    _queueLock.lock();
+    try {
+      _queryReaderCondition.signal();
+    } finally {
+      _queueLock.unlock();
+    }
+  }
 }