Add timeout overflow break handling at leaf stage

Author: xiangfu0

pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java

@@ -975,6 +975,10 @@ private void fillOldBrokerResponseStats(BrokerResponseNativeV2 brokerResponse,
             (int) numSegmentsPrunedByBroker);
         brokerResponse.addBrokerStats(brokerPruningStats);
       }
+
+      if (isLeafStageMissingStats(queryStageMap, queryStats)) {
+        brokerResponse.mergeTimeoutOverflowReached(true);
+      }
     } catch (Exception e) {
       LOGGER.warn("Error encountered while collecting multi-stage stats", e);
       brokerResponse.setStageStats(JsonNodeFactory.instance.objectNode()
@@ -996,6 +1000,31 @@ private void fillOldBrokerResponseStats(BrokerResponseNativeV2 brokerResponse,
     }
   }
 
+  private boolean isLeafStageMissingStats(Map<Integer, DispatchablePlanFragment> queryStageMap,
+      List<MultiStageQueryStats.StageStats.Closed> queryStats) {
+    for (Map.Entry<Integer, DispatchablePlanFragment> entry : queryStageMap.entrySet()) {
+      int stageId = entry.getKey();
+      DispatchablePlanFragment fragment = entry.getValue();
+      if (!isLeafStage(fragment)) {
+        continue;
+      }
+      MultiStageQueryStats.StageStats.Closed stageStats =
+          stageId < queryStats.size() ? queryStats.get(stageId) : null;
+      if (stageStats == null) {
+        return true;
+      }
+    }
+    return false;
+  }
+
+  private boolean isLeafStage(DispatchablePlanFragment fragment) {
+    Map<String, String> customProperties = fragment.getCustomProperties();
+    if (customProperties.containsKey(DispatchablePlanFragment.TABLE_NAME_KEY)) {
+      return true;
+    }
+    return !fragment.getWorkerIdToSegmentsMap().isEmpty();
+  }
+
   private BrokerResponse constructMultistageExplainPlan(String sql, String plan, Map<String, String> extraFields) {
     BrokerResponseNative brokerResponse = BrokerResponseNative.empty();
     int totalFieldCount = extraFields.size() + 2;

pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseNativeV2.java

@@ -88,6 +88,7 @@ public class BrokerResponseNativeV2 implements BrokerResponse {
   private String _requestId;
   private String _clientRequestId;
   private String _brokerId;
+  private boolean _timeoutOverflowReached;
   private int _numServersQueried;
   private int _numServersResponded;
   private long _brokerReduceTimeMs;
@@ -126,7 +127,7 @@ public int getNumRowsResultSet() {
   @Override
   public boolean isPartialResult() {
     return getExceptionsSize() > 0 || isNumGroupsLimitReached() || !getEarlyTerminationReasons().isEmpty()
-        || isMaxRowsInJoinReached() || isMseLiteLeafStageLimitReached();
+        || isMaxRowsInJoinReached() || isMseLiteLeafStageLimitReached() || isTimeoutOverflowReached();
   }
 
   @Override
@@ -211,6 +212,18 @@ public List<String> getEarlyTerminationReasons() {
     return List.copyOf(_brokerStats.getStringSet(StatKey.EARLY_TERMINATION_REASONS));
   }
 
+  @JsonProperty("timeoutOverflowReached")
+  public boolean isTimeoutOverflowReached() {
+    return _timeoutOverflowReached;
+  }
+
+  public void mergeTimeoutOverflowReached(boolean timeoutOverflowReached) {
+    if (timeoutOverflowReached) {
+      _timeoutOverflowReached = true;
+    }
+    _brokerStats.merge(StatKey.TIMEOUT_OVERFLOW_REACHED, timeoutOverflowReached);
+  }
+
   @Override
   public boolean isMaxRowsInJoinReached() {
     return _maxRowsInJoinReached;
@@ -515,6 +528,9 @@ public boolean getRLSFiltersApplied() {
 
   public void addBrokerStats(StatMap<StatKey> brokerStats) {
     _brokerStats.merge(brokerStats);
+    if (brokerStats.getBoolean(StatKey.TIMEOUT_OVERFLOW_REACHED)) {
+      _timeoutOverflowReached = true;
+    }
   }
 
   // NOTE: The following keys should match the keys in the leaf-stage operator.
@@ -550,7 +566,8 @@ public long merge(long value1, long value2) {
       }
     },
     EARLY_TERMINATION_REASONS(StatMap.Type.STRING_SET),
-    LITE_MODE_LEAF_STAGE_LIMIT_REACHED(StatMap.Type.BOOLEAN);
+    LITE_MODE_LEAF_STAGE_LIMIT_REACHED(StatMap.Type.BOOLEAN),
+    TIMEOUT_OVERFLOW_REACHED(StatMap.Type.BOOLEAN);
 
     private final StatMap.Type _type;
 

pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java

@@ -33,6 +33,7 @@
 import org.apache.pinot.spi.utils.CommonConstants;
 import org.apache.pinot.spi.utils.CommonConstants.Broker.Request.QueryOptionKey;
 import org.apache.pinot.spi.utils.CommonConstants.MultiStageQueryRunner.JoinOverFlowMode;
+import org.apache.pinot.spi.utils.CommonConstants.MultiStageQueryRunner.TimeoutOverflowMode;
 import org.apache.pinot.spi.utils.CommonConstants.MultiStageQueryRunner.WindowOverFlowMode;
 
 
@@ -511,6 +512,24 @@ public static WindowOverFlowMode getWindowOverflowMode(Map<String, String> query
     return windowOverflowModeStr != null ? WindowOverFlowMode.valueOf(windowOverflowModeStr) : null;
   }
 
+  @Nullable
+  public static TimeoutOverflowMode getTimeoutOverflowMode(Map<String, String> queryOptions) {
+    String timeoutOverflowModeStr = queryOptions.get(QueryOptionKey.TIMEOUT_OVERFLOW_MODE);
+    return timeoutOverflowModeStr != null ? TimeoutOverflowMode.valueOf(timeoutOverflowModeStr) : null;
+  }
+
+  @Nullable
+  public static Long getLeafTimeoutMs(Map<String, String> queryOptions) {
+    String leafTimeoutMs = queryOptions.get(QueryOptionKey.LEAF_TIMEOUT_MS);
+    return checkedParseLongPositive(QueryOptionKey.LEAF_TIMEOUT_MS, leafTimeoutMs);
+  }
+
+  @Nullable
+  public static Long getLeafExtraPassiveTimeoutMs(Map<String, String> queryOptions) {
+    String leafExtraPassiveTimeoutMs = queryOptions.get(QueryOptionKey.LEAF_EXTRA_PASSIVE_TIMEOUT_MS);
+    return checkedParseLong(QueryOptionKey.LEAF_EXTRA_PASSIVE_TIMEOUT_MS, leafExtraPassiveTimeoutMs, 0);
+  }
+
   public static boolean isSkipUnavailableServers(Map<String, String> queryOptions) {
     return Boolean.parseBoolean(queryOptions.get(QueryOptionKey.SKIP_UNAVAILABLE_SERVERS));
   }

pinot-common/src/test/java/org/apache/pinot/common/response/broker/BrokerResponseNativeV2Test.java

@@ -18,10 +18,13 @@
  */
 package org.apache.pinot.common.response.broker;
 
+import com.fasterxml.jackson.databind.JsonNode;
+import java.io.IOException;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Set;
 import org.apache.pinot.common.datatable.StatMap;
+import org.apache.pinot.spi.utils.JsonUtils;
 import org.testng.annotations.Test;
 
 import static org.testng.Assert.assertEquals;
@@ -47,6 +50,23 @@ public void testEarlyTerminationReasonsMarkPartialResponse() {
     assertTrue(brokerResponse.isPartialResult());
   }
 
+  @Test
+  public void testTimeoutOverflowFlagSerialized()
+      throws IOException {
+    BrokerResponseNativeV2 response = new BrokerResponseNativeV2();
+    assertFalse(response.isTimeoutOverflowReached());
+    assertFalse(response.isPartialResult());
+
+    response.mergeTimeoutOverflowReached(true);
+
+    assertTrue(response.isTimeoutOverflowReached());
+    assertTrue(response.isPartialResult());
+
+    JsonNode json = JsonUtils.objectToJsonNode(response);
+    assertTrue(json.get("timeoutOverflowReached").asBoolean());
+    assertTrue(json.get("partialResult").asBoolean());
+  }
+
   private static Set<String> stringSet(String... values) {
     return new LinkedHashSet<>(List.of(values));
   }

pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/QueryRunner.java

@@ -77,6 +77,7 @@
 import org.apache.pinot.spi.utils.CommonConstants.Helix;
 import org.apache.pinot.spi.utils.CommonConstants.MultiStageQueryRunner;
 import org.apache.pinot.spi.utils.CommonConstants.MultiStageQueryRunner.JoinOverFlowMode;
+import org.apache.pinot.spi.utils.CommonConstants.MultiStageQueryRunner.TimeoutOverflowMode;
 import org.apache.pinot.spi.utils.CommonConstants.MultiStageQueryRunner.WindowOverFlowMode;
 import org.apache.pinot.spi.utils.CommonConstants.Query.Request;
 import org.apache.pinot.spi.utils.CommonConstants.Query.Response;
@@ -98,6 +99,7 @@
  */
 public class QueryRunner {
   private static final Logger LOGGER = LoggerFactory.getLogger(QueryRunner.class);
+  private static final double DEFAULT_LEAF_TIMEOUT_RATIO = 0.8;
 
   private ExecutorService _executorService;
   private OpChainSchedulerService _opChainScheduler;
@@ -130,6 +132,8 @@ public class QueryRunner {
   @Nullable
   private WindowOverFlowMode _windowOverflowMode;
   @Nullable
+  private TimeoutOverflowMode _timeoutOverflowMode;
+  @Nullable
   private PhysicalTimeSeriesServerPlanVisitor _timeSeriesPhysicalPlanVisitor;
   /// Cluster-level decision on whether to send stats over the mailbox path, driven by the {@code SendStatsPredicate}
   /// at startup time. <b>May be overridden per-request</b> via the {@code KEY_OF_STATS_REPORTING_MODE} metadata key —
@@ -191,6 +195,10 @@ public void init(PinotConfiguration serverConf, String instanceId, @Nullable Ins
     String windowOverflowModeStr = serverConf.getProperty(MultiStageQueryRunner.KEY_OF_WINDOW_OVERFLOW_MODE);
     _windowOverflowMode = windowOverflowModeStr != null ? WindowOverFlowMode.valueOf(windowOverflowModeStr) : null;
 
+    String timeoutOverflowModeStr = serverConf.getProperty(MultiStageQueryRunner.KEY_OF_TIMEOUT_OVERFLOW_MODE);
+    _timeoutOverflowMode =
+        timeoutOverflowModeStr != null ? TimeoutOverflowMode.valueOf(timeoutOverflowModeStr) : null;
+
     ExecutorService baseExecutorService =
         ExecutorServiceUtils.create(serverConf, Server.MULTISTAGE_EXECUTOR_CONFIG_PREFIX, "query-runner-on-" + port,
             Server.DEFAULT_MULTISTAGE_EXECUTOR_TYPE);
@@ -301,6 +309,7 @@ private void processQueryBlocking(WorkerMetadata workerMetadata, StagePlan stage
       Map<String, String> requestMetadata) {
     StageMetadata stageMetadata = stagePlan.getStageMetadata();
     Map<String, String> opChainMetadata = consolidateMetadata(stageMetadata.getCustomProperties(), requestMetadata);
+    applyLeafDeadlineOverrides(workerMetadata, opChainMetadata);
 
     // The cluster-level _sendStats decision can be overridden per-request by the SubmitWithStream RPC handler via
     // MultiStageQueryRunner.KEY_OF_STATS_REPORTING_MODE; in stream mode stats travel out-of-band
@@ -557,6 +566,14 @@ private Map<String, String> consolidateMetadata(Map<String, String> customProper
       opChainMetadata.put(QueryOptionKey.WINDOW_OVERFLOW_MODE, windowOverflowMode.name());
     }
 
+    TimeoutOverflowMode timeoutOverflowMode = QueryOptionsUtils.getTimeoutOverflowMode(opChainMetadata);
+    if (timeoutOverflowMode == null) {
+      timeoutOverflowMode = _timeoutOverflowMode;
+    }
+    if (timeoutOverflowMode != null) {
+      opChainMetadata.put(QueryOptionKey.TIMEOUT_OVERFLOW_MODE, timeoutOverflowMode.name());
+    }
+
     return opChainMetadata;
   }
 
@@ -604,6 +621,53 @@ public void unregisterOpChainCompletionListener(long requestId) {
     _opChainScheduler.unregisterCompletionListener(requestId);
   }
 
+  private void applyLeafDeadlineOverrides(WorkerMetadata workerMetadata, Map<String, String> opChainMetadata) {
+    if (!workerMetadata.isLeafStageWorker()) {
+      return;
+    }
+    Long leafTimeoutMs = QueryOptionsUtils.getLeafTimeoutMs(opChainMetadata);
+    Long leafExtraPassiveTimeoutMs = QueryOptionsUtils.getLeafExtraPassiveTimeoutMs(opChainMetadata);
+    TimeoutOverflowMode timeoutOverflowMode = _timeoutOverflowMode;
+    String timeoutOverflowModeStr = opChainMetadata.get(QueryOptionKey.TIMEOUT_OVERFLOW_MODE);
+    if (timeoutOverflowModeStr != null) {
+      try {
+        timeoutOverflowMode = TimeoutOverflowMode.valueOf(timeoutOverflowModeStr);
+      } catch (IllegalArgumentException e) {
+        LOGGER.warn("Invalid timeout overflow mode: {}", timeoutOverflowModeStr, e);
+      }
+    }
+    QueryThreadContext threadContext = QueryThreadContext.getIfAvailable();
+    if (threadContext == null) {
+      return;
+    }
+    QueryExecutionContext executionContext = threadContext.getExecutionContext();
+    long startTimeMs = executionContext.getStartTimeMs();
+    if (leafTimeoutMs == null && timeoutOverflowMode == TimeoutOverflowMode.BREAK) {
+      long globalActiveWindowMs = Math.max(1L, executionContext.getActiveDeadlineMs() - startTimeMs);
+      leafTimeoutMs = Math.max(1L, (long) Math.floor(globalActiveWindowMs * DEFAULT_LEAF_TIMEOUT_RATIO));
+    }
+    if (leafExtraPassiveTimeoutMs == null && timeoutOverflowMode == TimeoutOverflowMode.BREAK) {
+      long globalPassiveWindowMs =
+          Math.max(0L, executionContext.getPassiveDeadlineMs() - executionContext.getActiveDeadlineMs());
+      leafExtraPassiveTimeoutMs =
+          Math.max(0L, (long) Math.floor(globalPassiveWindowMs * DEFAULT_LEAF_TIMEOUT_RATIO));
+    }
+    if (leafTimeoutMs == null && leafExtraPassiveTimeoutMs == null) {
+      return;
+    }
+    long activeDeadlineMs = executionContext.getActiveDeadlineMs();
+    long passiveDeadlineMs = executionContext.getPassiveDeadlineMs();
+    if (leafTimeoutMs != null) {
+      activeDeadlineMs = Math.min(activeDeadlineMs, startTimeMs + leafTimeoutMs);
+    }
+    if (leafExtraPassiveTimeoutMs != null) {
+      passiveDeadlineMs = Math.min(passiveDeadlineMs, activeDeadlineMs + leafExtraPassiveTimeoutMs);
+    }
+    passiveDeadlineMs = Math.max(passiveDeadlineMs, activeDeadlineMs);
+    opChainMetadata.put(LeafOperator.LEAF_ACTIVE_DEADLINE_METADATA_KEY, Long.toString(activeDeadlineMs));
+    opChainMetadata.put(LeafOperator.LEAF_PASSIVE_DEADLINE_METADATA_KEY, Long.toString(passiveDeadlineMs));
+  }
+
   public StagePlan explainQuery(WorkerMetadata workerMetadata, StagePlan stagePlan,
       Map<String, String> requestMetadata) {
     if (!workerMetadata.isLeafStageWorker()) {
@@ -613,6 +677,7 @@ public StagePlan explainQuery(WorkerMetadata workerMetadata, StagePlan stagePlan
 
     StageMetadata stageMetadata = stagePlan.getStageMetadata();
     Map<String, String> opChainMetadata = consolidateMetadata(stageMetadata.getCustomProperties(), requestMetadata);
+    applyLeafDeadlineOverrides(workerMetadata, opChainMetadata);
 
     if (PipelineBreakerExecutor.hasPipelineBreakers(stagePlan)) {
       //TODO: See https://github.com/apache/pinot/pull/13733#discussion_r1752031714