Add timeout overflow break handling at leaf stage

Author: xiangfu0

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;
@@ -138,6 +139,8 @@ public class QueryRunner {
   @Nullable
   private WindowOverFlowMode _windowOverflowMode;
   @Nullable
+  private TimeoutOverflowMode _timeoutOverflowMode;
+  @Nullable
   private PhysicalTimeSeriesServerPlanVisitor _timeSeriesPhysicalPlanVisitor;
   private BooleanSupplier _sendStats;
 
@@ -191,6 +194,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 +278,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 +516,45 @@ 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);
+    if (leafTimeoutMs == null && leafExtraPassiveTimeoutMs == null) {
+      return;
+    }
+    QueryThreadContext threadContext = QueryThreadContext.getIfAvailable();
+    if (threadContext == null) {
+      return;
+    }
+    QueryExecutionContext executionContext = threadContext.getExecutionContext();
+    long startTimeMs = executionContext.getStartTimeMs();
+    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 +568,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

pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/BaseJoinOperator.java

@@ -39,9 +39,11 @@
 import org.apache.pinot.query.runtime.operator.operands.TransformOperandFactory;
 import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
 import org.apache.pinot.spi.exception.QueryErrorCode;
+import org.apache.pinot.spi.exception.QueryException;
 import org.apache.pinot.spi.utils.BooleanUtils;
 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.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -83,10 +85,12 @@ public abstract class BaseJoinOperator extends MultiStageOperator {
    *   BREAK: Break right table build process, continue to perform JOIN operation, results might be partial.
    */
   protected final JoinOverFlowMode _joinOverflowMode;
+  protected final TimeoutOverflowMode _timeoutOverflowMode;
 
   protected boolean _isRightTableBuilt;
   @Nullable
   protected MseBlock.Eos _eos;
+  private boolean _timeoutOverflowTriggered;
 
   public BaseJoinOperator(OpChainExecutionContext context, MultiStageOperator leftInput, DataSchema leftSchema,
       MultiStageOperator rightInput, JoinNode node) {
@@ -106,6 +110,8 @@ public BaseJoinOperator(OpChainExecutionContext context, MultiStageOperator left
     PlanNode.NodeHint nodeHint = node.getNodeHint();
     _maxRowsInJoin = getMaxRowsInJoin(metadata, nodeHint);
     _joinOverflowMode = getJoinOverflowMode(metadata, nodeHint);
+    TimeoutOverflowMode timeoutOverflowMode = QueryOptionsUtils.getTimeoutOverflowMode(metadata);
+    _timeoutOverflowMode = timeoutOverflowMode != null ? timeoutOverflowMode : TimeoutOverflowMode.THROW;
   }
 
   /// Constructor that takes the schema for NonEquiEvaluator as an argument
@@ -127,6 +133,8 @@ public BaseJoinOperator(OpChainExecutionContext context, MultiStageOperator left
     PlanNode.NodeHint nodeHint = node.getNodeHint();
     _maxRowsInJoin = getMaxRowsInJoin(metadata, nodeHint);
     _joinOverflowMode = getJoinOverflowMode(metadata, nodeHint);
+    TimeoutOverflowMode timeoutOverflowMode = QueryOptionsUtils.getTimeoutOverflowMode(metadata);
+    _timeoutOverflowMode = timeoutOverflowMode != null ? timeoutOverflowMode : TimeoutOverflowMode.THROW;
   }
 
   protected static int getMaxRowsInJoin(Map<String, String> opChainMetadata, @Nullable PlanNode.NodeHint nodeHint) {
@@ -158,6 +166,14 @@ protected static JoinOverFlowMode getJoinOverflowMode(Map<String, String> contex
     return joinOverflowMode != null ? joinOverflowMode : DEFAULT_JOIN_OVERFLOW_MODE;
   }
 
+  @Override
+  protected long getDeadlineMs() {
+    if (_timeoutOverflowTriggered && _timeoutOverflowMode == TimeoutOverflowMode.BREAK) {
+      return Long.MAX_VALUE;
+    }
+    return super.getDeadlineMs();
+  }
+
   @Override
   public void registerExecution(long time, int numRows, long memoryUsedBytes, long gcTimeMs) {
     _statMap.merge(StatKey.EXECUTION_TIME_MS, time);
@@ -283,6 +299,14 @@ protected MseBlock buildJoinedDataBlock() {
 
   protected abstract List<Object[]> buildNonMatchRightRows();
 
+  @Override
+  protected MseBlock handleException(Exception e) {
+    if (maybeHandleTimeoutOverflow(e)) {
+      return SuccessMseBlock.INSTANCE;
+    }
+    return super.handleException(e);
+  }
+
   // TODO: Optimize this to avoid unnecessary object copy.
   protected Object[] joinRow(@Nullable Object[] leftRow, @Nullable Object[] rightRow) {
     Object[] resultRow = new Object[_resultColumnSize];
@@ -329,6 +353,29 @@ protected void earlyTerminateLeftInput() {
     _eos = (MseBlock.Eos) leftBlock;
   }
 
+  private boolean maybeHandleTimeoutOverflow(Exception e) {
+    if (_timeoutOverflowMode != TimeoutOverflowMode.BREAK) {
+      return false;
+    }
+    if (e instanceof QueryException) {
+      QueryException queryException = (QueryException) e;
+      if (queryException.getErrorCode() == QueryErrorCode.EXECUTION_TIMEOUT) {
+        if (!_timeoutOverflowTriggered) {
+          logger().warn("Join operator {} returning partial results after timeout: {}", _operatorId,
+              queryException.getMessage());
+        }
+        _timeoutOverflowTriggered = true;
+        _statMap.merge(StatKey.TIMEOUT_OVERFLOW_REACHED, true);
+        _leftInput.earlyTerminate();
+        _rightInput.earlyTerminate();
+        _eos = SuccessMseBlock.INSTANCE;
+        onEosProduced();
+        return true;
+      }
+    }
+    return false;
+  }
+
   @Override
   protected StatMap<?> copyStatMaps() {
     return new StatMap<>(_statMap);
@@ -397,7 +444,8 @@ public boolean includeDefaultInJson() {
     /**
      * Time spent on GC while this operator or its children in the same stage were running.
      */
-    GC_TIME_MS(StatMap.Type.LONG);
+    GC_TIME_MS(StatMap.Type.LONG),
+    TIMEOUT_OVERFLOW_REACHED(StatMap.Type.BOOLEAN);
 
     private final StatMap.Type _type;