Add timeout overflow break handling at leaf stage

Author: xiangfu0

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

@@ -708,13 +708,42 @@ private void fillOldBrokerResponseStats(BrokerResponseNativeV2 brokerResponse,
           stageStats.forEach((type, stats) -> type.mergeInto(brokerResponse, stats));
         }
       }
+
+      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()
           .put("error", "Error encountered while collecting multi-stage stats - " + e));
     }
   }
 
+  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

@@ -76,6 +76,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;
@@ -109,7 +110,8 @@ public int getNumRowsResultSet() {
   @JsonProperty(access = JsonProperty.Access.READ_ONLY)
   @Override
   public boolean isPartialResult() {
-    return getExceptionsSize() > 0 || isNumGroupsLimitReached() || isMaxRowsInJoinReached();
+    return getExceptionsSize() > 0 || isNumGroupsLimitReached() || isMaxRowsInJoinReached()
+        || isTimeoutOverflowReached();
   }
 
   @Override
@@ -152,6 +154,18 @@ public void mergeNumGroupsWarningLimitReached(boolean numGroupsWarningLimitReach
     _brokerStats.merge(StatKey.NUM_GROUPS_WARNING_LIMIT_REACHED, numGroupsWarningLimitReached);
   }
 
+  @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;
@@ -427,6 +441,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.
@@ -453,7 +470,8 @@ public long merge(long value1, long value2) {
     NUM_SEGMENTS_PRUNED_BY_VALUE(StatMap.Type.INT),
     GROUPS_TRIMMED(StatMap.Type.BOOLEAN),
     NUM_GROUPS_LIMIT_REACHED(StatMap.Type.BOOLEAN),
-    NUM_GROUPS_WARNING_LIMIT_REACHED(StatMap.Type.BOOLEAN);
+    NUM_GROUPS_WARNING_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

@@ -34,6 +34,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;
 
 
@@ -432,6 +433,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

@@ -0,0 +1,46 @@
+/**
+ * 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.common.response.broker;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import java.io.IOException;
+import org.apache.pinot.spi.utils.JsonUtils;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+
+public class BrokerResponseNativeV2Test {
+
+  @Test
+  public void testTimeoutOverflowFlagSerialized()
+      throws IOException {
+    BrokerResponseNativeV2 response = new BrokerResponseNativeV2();
+    Assert.assertFalse(response.isTimeoutOverflowReached());
+    Assert.assertFalse(response.isPartialResult());
+
+    response.mergeTimeoutOverflowReached(true);
+
+    Assert.assertTrue(response.isTimeoutOverflowReached());
+    Assert.assertTrue(response.isPartialResult());
+
+    JsonNode json = JsonUtils.objectToJsonNode(response);
+    Assert.assertTrue(json.get("timeoutOverflowReached").asBoolean());
+    Assert.assertTrue(json.get("partialResult").asBoolean());
+  }
+}

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

@@ -86,6 +86,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;
@@ -107,6 +108,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;
@@ -138,6 +140,8 @@ public class QueryRunner {
   @Nullable
   private WindowOverFlowMode _windowOverflowMode;
   @Nullable
+  private TimeoutOverflowMode _timeoutOverflowMode;
+  @Nullable
   private PhysicalTimeSeriesServerPlanVisitor _timeSeriesPhysicalPlanVisitor;
   private BooleanSupplier _sendStats;
 
@@ -191,6 +195,10 @@ public void init(PinotConfiguration serverConf, InstanceDataManager instanceData
     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);
@@ -271,6 +279,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);
 
     // run pre-stage execution for all pipeline breakers
     PipelineBreakerResult pipelineBreakerResult =
@@ -508,9 +517,64 @@ 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;
   }
 
+  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 Map<Integer, MultiStageQueryStats.StageStats.Closed> cancel(long requestId) {
     return _opChainScheduler.cancel(requestId);
   }
@@ -524,6 +588,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