Re: [PR] feat: Segment Timout Per Datasource (druid)

["jtuglu1 (via GitHub)" <gi...@apache.org>]

jtuglu1 commented on code in PR #19221:
URL: https://github.com/apache/druid/pull/19221#discussion_r3019609150


##########
server/src/main/java/org/apache/druid/server/broker/BrokerDynamicConfig.java:
##########
@@ -52,14 +54,24 @@ public class BrokerDynamicConfig
    */
   private final QueryContext queryContext;
 
+  /**
+   * Per-datasource per-segment timeout configuration. Maps datasource name to its
+   * {@link PerSegmentTimeoutConfig}. When a query targets a datasource in this map,
+   * the broker injects the configured {@code perSegmentTimeout} into the query context
+   * (unless the caller already set it explicitly).
+   */
+  private final Map<String, PerSegmentTimeoutConfig> perSegmentTimeoutConfig;
+
   @JsonCreator
   public BrokerDynamicConfig(
       @JsonProperty("queryBlocklist") @Nullable List<QueryBlocklistRule> queryBlocklist,
-      @JsonProperty("queryContext") @Nullable QueryContext queryContext
+      @JsonProperty("queryContext") @Nullable QueryContext queryContext,
+      @JsonProperty("perSegmentTimeoutConfig") @Nullable Map<String, PerSegmentTimeoutConfig> perSegmentTimeoutConfig

Review Comment:
   It'd be great if we could enforce this as an immutable map (either with an immutable copy or otherwise) as AFAIK this deserde into a mutable LinkedHashMap.



##########
server/src/main/java/org/apache/druid/server/broker/BrokerDynamicConfig.java:
##########
@@ -74,6 +86,13 @@ public QueryContext getQueryContext()
     return queryContext;
   }
 
+  @JsonProperty
+  @JsonInclude(JsonInclude.Include.NON_EMPTY)

Review Comment:
   Have we tested with a null/empty map value? I think this will cause a bug in the UI. We want to allow people to "reset" the dynamic config, and this will exclude from the (de)serde I believe.



##########
server/src/main/java/org/apache/druid/server/QueryLifecycle.java:
##########
@@ -216,15 +221,49 @@ public void initialize(final Query<?> baseQuery)
       queryId = UUID.randomUUID().toString();
     }
 
-    Map<String, Object> mergedUserAndConfigContext = QueryContexts.override(
-        queryConfigProvider.getContext(),
-        baseQuery.getContext()
-    );
-    mergedUserAndConfigContext.put(BaseQuery.QUERY_ID, queryId);
-    this.baseQuery = baseQuery.withOverriddenContext(mergedUserAndConfigContext);
+    // Start with system defaults, apply per-datasource override, then user context wins
+    Map<String, Object> contextWithDefaults = new HashMap<>(queryConfigProvider.getContext());
+    applyPerDatasourcePerSegmentTimeout(baseQuery, contextWithDefaults, queryId);
+    Map<String, Object> finalContext = QueryContexts.override(contextWithDefaults, baseQuery.getContext());
+    finalContext.put(BaseQuery.QUERY_ID, queryId);
+
+    this.baseQuery = baseQuery.withOverriddenContext(finalContext);
     this.toolChest = conglomerate.getToolChest(this.baseQuery);
   }
 
+  /**
+   * If a per-datasource per-segment timeout is configured, injects it into the context defaults.
+   * User context (applied later via {@link QueryContexts#override}) will override this if set explicitly.
+   * In monitorOnly mode, logs the configured timeout but does not inject it.
+   */
+  private void applyPerDatasourcePerSegmentTimeout(
+      final Query<?> query,
+      final Map<String, Object> contextWithDefaults,
+      final String queryId
+  )
+  {
+    if (perSegmentTimeoutConfig.isEmpty()) {
+      return;
+    }
+
+    for (String tableName : query.getDataSource().getTableNames()) {

Review Comment:
   Let's leave a call-out in the docs, etc. saying that we will apply the *first* timeout (not the max or min) across things like a join or UnionDatasource where `getDataSource().getTableNames()` returns multiple values. This might come as unexpected for some users who see non-deterministic timeouts (depending on the order of the elements in the Set).



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org