fix: fix query context precedence layer for datasource-level per-segment timeout

processing/src/main/java/org/apache/druid/query/QueryConfigProvider.java

@@ -22,7 +22,7 @@
 import java.util.Map;
 
 /**
- * Provides the default query config applied to all incoming queries before per-query overrides are merged in.
+ * Provides the default query context applied to all incoming queries before per-query overrides are merged in.
  *
  * <p>On non-broker nodes this is backed by static runtime properties ({@link DefaultQueryConfig}).
  * On brokers, it is backed by {@code BrokerViewOfBrokerConfig}, which merges the static defaults with

server/src/main/java/org/apache/druid/client/BrokerViewOfBrokerConfig.java

@@ -36,6 +36,7 @@
 import org.apache.druid.rpc.ServiceLocator;
 import org.apache.druid.rpc.StandardRetryPolicy;
 import org.apache.druid.server.broker.BrokerDynamicConfig;
+import org.apache.druid.server.broker.QueryConfigSnapshot;
 
 import javax.validation.constraints.NotNull;
 import java.util.Map;
@@ -127,4 +128,13 @@ public Map<String, Object> getContext()
   {
     return resolvedDefaultQueryContext.asMap();
   }
+
+  /**
+   * Captures the resolved default context and the dynamic config atomically, so a single query resolves its context
+   * and blocklist against one consistent snapshot rather than re-reading the live config.
+   */
+  public synchronized QueryConfigSnapshot snapshotForQuery()
+  {
+    return new QueryConfigSnapshot(getContext(), getDynamicConfig());
+  }
 }

server/src/main/java/org/apache/druid/server/QueryLifecycle.java

@@ -37,9 +37,7 @@
 import org.apache.druid.query.DruidMetrics;
 import org.apache.druid.query.GenericQueryMetricsFactory;
 import org.apache.druid.query.Query;
-import org.apache.druid.query.QueryConfigProvider;
 import org.apache.druid.query.QueryContext;
-import org.apache.druid.query.QueryContexts;
 import org.apache.druid.query.QueryInterruptedException;
 import org.apache.druid.query.QueryMetrics;
 import org.apache.druid.query.QueryPlus;
@@ -49,7 +47,7 @@
 import org.apache.druid.query.QueryToolChest;
 import org.apache.druid.query.context.ResponseContext;
 import org.apache.druid.query.policy.PolicyEnforcer;
-import org.apache.druid.server.broker.PerSegmentTimeoutConfig;
+import org.apache.druid.server.broker.QueryConfigSnapshot;
 import org.apache.druid.server.log.RequestLogger;
 import org.apache.druid.server.security.Action;
 import org.apache.druid.server.security.AuthConfig;
@@ -64,7 +62,6 @@
 
 import javax.annotation.Nullable;
 import javax.servlet.http.HttpServletRequest;
-import java.util.HashMap;
 import java.util.HashSet;
 import java.util.LinkedHashMap;
 import java.util.List;
@@ -99,11 +96,9 @@ public class QueryLifecycle
   private final ServiceEmitter emitter;
   private final RequestLogger requestLogger;
   private final AuthorizerMapper authorizerMapper;
-  private final QueryConfigProvider queryConfigProvider;
   private final AuthConfig authConfig;
   private final PolicyEnforcer policyEnforcer;
-  private final List<QueryBlocklistRule> queryBlocklist;
-  private final Map<String, PerSegmentTimeoutConfig> perSegmentTimeoutConfig;
+  private final QueryConfigSnapshot configSnapshot;
   private final long startMs;
   private final long startNs;
 
@@ -113,8 +108,11 @@ public class QueryLifecycle
 
   @MonotonicNonNull
   private Query<?> baseQuery;
+  /**
+   * Keys present on the query context as received; the candidate set fed to context-key authorization.
+   */
   @MonotonicNonNull
-  private Set<String> userContextKeys;
+  private Set<String> authorizationContextKeys;
 
   public QueryLifecycle(
       final QueryRunnerFactoryConglomerate conglomerate,
@@ -123,11 +121,9 @@ public QueryLifecycle(
       final ServiceEmitter emitter,
       final RequestLogger requestLogger,
       final AuthorizerMapper authorizerMapper,
-      final QueryConfigProvider queryConfigProvider,
       final AuthConfig authConfig,
       final PolicyEnforcer policyEnforcer,
-      final List<QueryBlocklistRule> queryBlocklist,
-      final Map<String, PerSegmentTimeoutConfig> perSegmentTimeoutConfig,
+      final QueryConfigSnapshot configSnapshot,
       final long startMs,
       final long startNs
   )
@@ -138,11 +134,9 @@ public QueryLifecycle(
     this.emitter = emitter;
     this.requestLogger = requestLogger;
     this.authorizerMapper = authorizerMapper;
-    this.queryConfigProvider = queryConfigProvider;
     this.authConfig = authConfig;
     this.policyEnforcer = policyEnforcer;
-    this.queryBlocklist = queryBlocklist;
-    this.perSegmentTimeoutConfig = perSegmentTimeoutConfig;
+    this.configSnapshot = configSnapshot;
     this.startMs = startMs;
     this.startNs = startNs;
   }
@@ -166,7 +160,21 @@ public <T> QueryResponse<T> runSimple(
       final AuthorizationResult authorizationResult
   )
   {
-    initialize(query);
+    return runSimple(query, authenticationResult, authorizationResult, null);
+  }
+
+  /**
+   * As {@link #runSimple(Query, AuthenticationResult, AuthorizationResult)}, but takes the context keys the client
+   * actually set. See {@link #initialize(Query, Set)}.
+   */
+  public <T> QueryResponse<T> runSimple(
+      final Query<T> query,
+      final AuthenticationResult authenticationResult,
+      final AuthorizationResult authorizationResult,
+      @Nullable final Set<String> clientProvidedQueryContextKeys
+  )
+  {
+    initialize(query, clientProvidedQueryContextKeys);
 
     final Sequence<T> results;
 
@@ -212,61 +220,41 @@ public void after(final boolean isDone, final Throwable thrown)
    * @throws DruidException if the current state is not NEW, which indicates a bug
    */
   public void initialize(final Query<?> baseQuery)
+  {
+    initialize(baseQuery, null);
+  }
+
+  /**
+   * As {@link #initialize(Query)}, but takes the context keys the client actually set. Pass {@code null} to treat the
+   * whole context as client-set (native queries). The SQL layer merges static defaults into the context, so it must
+   * pass the real client-set keys so dynamic overrides can beat a merged-in default without overriding the client.
+   *
+   * @throws DruidException if the current state is not NEW, which indicates a bug
+   */
+  public void initialize(final Query<?> baseQuery, @Nullable final Set<String> clientProvidedQueryContextKeys)
   {
     transition(State.NEW, State.INITIALIZED);
 
-    userContextKeys = new HashSet<>(baseQuery.getContext().keySet());
+    final Map<String, Object> baseContext = baseQuery.getContext();
+    authorizationContextKeys = new HashSet<>(baseContext.keySet());
+
+    // Keys the client actually set (native queries pass null, so the whole context is client-set).
+    final Set<String> effectiveClientProvidedQueryContextKeys =
+        clientProvidedQueryContextKeys != null ? clientProvidedQueryContextKeys : baseContext.keySet();
+
     String queryId = baseQuery.getId();
     if (Strings.isNullOrEmpty(queryId)) {
       queryId = UUID.randomUUID().toString();
     }
 
-    // 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());
+    final Map<String, Object> finalContext =
+        configSnapshot.resolveContext(baseQuery, effectiveClientProvidedQueryContextKeys);
     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.
-   *
-   * For queries involving multiple datasources (e.g., joins or unions), the timeout from the first matching datasource is applied
-   * since getTableNames() returns a Set, the match order is non-deterministic.
-   */
-  private void applyPerDatasourcePerSegmentTimeout(
-      final Query<?> query,
-      final Map<String, Object> contextWithDefaults,
-      final String queryId
-  )
-  {
-    if (perSegmentTimeoutConfig.isEmpty()) {
-      return;
-    }
-
-    for (String tableName : query.getDataSource().getTableNames()) {
-      PerSegmentTimeoutConfig dsConfig = perSegmentTimeoutConfig.get(tableName);
-      if (dsConfig != null) {
-        if (dsConfig.isMonitorOnly()) {
-          log.debug(
-              "Per-segment timeout [%d ms] configured for datasource [%s] in monitorOnly mode (not enforced) for query [%s].",
-              dsConfig.getPerSegmentTimeoutMs(),
-              tableName,
-              queryId
-          );
-        } else {
-          contextWithDefaults.put(QueryContexts.PER_SEGMENT_TIMEOUT_KEY, dsConfig.getPerSegmentTimeoutMs());
-        }
-        return;
-      }
-    }
-  }
-
   /**
    * Returns {@link AuthorizationResult} based on {@code DRUID_AUTHENTICATION_RESULT} in the given request, base query
    * would be transformed with restrictions on the AuthorizationResult.
@@ -290,7 +278,7 @@ public AuthorizationResult authorize(HttpServletRequest req)
             AuthorizationUtils.DATASOURCE_READ_RA_GENERATOR
         ),
         Iterables.transform(
-            authConfig.contextKeysToAuthorize(userContextKeys),
+            authConfig.contextKeysToAuthorize(authorizationContextKeys),
             contextParam -> new ResourceAction(new Resource(contextParam, ResourceType.QUERY_CONTEXT), Action.WRITE)
         )
     );
@@ -328,7 +316,7 @@ public AuthorizationResult authorize(AuthenticationResult authenticationResult)
             AuthorizationUtils.DATASOURCE_READ_RA_GENERATOR
         ),
         Iterables.transform(
-            authConfig.contextKeysToAuthorize(userContextKeys),
+            authConfig.contextKeysToAuthorize(authorizationContextKeys),
             contextParam -> new ResourceAction(new Resource(contextParam, ResourceType.QUERY_CONTEXT), Action.WRITE)
         )
     );
@@ -364,7 +352,9 @@ private void preAuthorized(
    */
   private void checkQueryBlocklist()
   {
-    if (queryBlocklist == null || queryBlocklist.isEmpty()) {
+    final List<QueryBlocklistRule> queryBlocklist = configSnapshot.getQueryBlocklist();
+
+    if (queryBlocklist.isEmpty()) {
       return;
     }
 

server/src/main/java/org/apache/druid/server/QueryLifecycleFactory.java

@@ -28,15 +28,12 @@
 import org.apache.druid.query.QueryRunnerFactoryConglomerate;
 import org.apache.druid.query.QuerySegmentWalker;
 import org.apache.druid.query.policy.PolicyEnforcer;
-import org.apache.druid.server.broker.PerSegmentTimeoutConfig;
+import org.apache.druid.server.broker.QueryConfigSnapshot;
 import org.apache.druid.server.log.RequestLogger;
 import org.apache.druid.server.security.AuthConfig;
 import org.apache.druid.server.security.AuthorizerMapper;
 
 import javax.annotation.Nullable;
-import java.util.Collections;
-import java.util.List;
-import java.util.Map;
 
 @LazySingleton
 public class QueryLifecycleFactory
@@ -80,15 +77,11 @@ public QueryLifecycleFactory(
 
   public QueryLifecycle factorize()
   {
-    final List<QueryBlocklistRule> queryBlocklist;
-    final Map<String, PerSegmentTimeoutConfig> perSegmentTimeoutConfig;
-    if (brokerViewOfBrokerConfig != null && brokerViewOfBrokerConfig.getDynamicConfig() != null) {
-      queryBlocklist = brokerViewOfBrokerConfig.getDynamicConfig().getQueryBlocklist();
-      perSegmentTimeoutConfig = brokerViewOfBrokerConfig.getDynamicConfig().getPerSegmentTimeoutConfig();
-    } else {
-      queryBlocklist = Collections.emptyList();
-      perSegmentTimeoutConfig = Collections.emptyMap();
-    }
+    // Capture one snapshot per query so the default context, per-query overrides, and blocklist are consistent.
+    final QueryConfigSnapshot configSnapshot =
+        brokerViewOfBrokerConfig == null
+        ? new QueryConfigSnapshot(queryConfigProvider.getContext(), null)
+        : brokerViewOfBrokerConfig.snapshotForQuery();
 
     return new QueryLifecycle(
         conglomerate,
@@ -97,11 +90,9 @@ public QueryLifecycle factorize()
         emitter,
         requestLogger,
         authorizerMapper,
-        queryConfigProvider,
         authConfig,
         policyEnforcer,
-        queryBlocklist,
-        perSegmentTimeoutConfig,
+        configSnapshot,
         System.currentTimeMillis(),
         System.nanoTime()
     );

server/src/main/java/org/apache/druid/server/broker/BrokerDynamicConfig.java

@@ -22,7 +22,10 @@
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import org.apache.druid.common.config.Configs;
+import org.apache.druid.java.util.common.logger.Logger;
+import org.apache.druid.query.Query;
 import org.apache.druid.query.QueryContext;
+import org.apache.druid.query.QueryContexts;
 import org.apache.druid.server.QueryBlocklistRule;
 
 import javax.annotation.Nullable;
@@ -40,6 +43,8 @@
  */
 public class BrokerDynamicConfig
 {
+  private static final Logger log = new Logger(BrokerDynamicConfig.class);
+
   public static final String CONFIG_KEY = "broker.config";
 
   /**
@@ -91,6 +96,33 @@ public Map<String, PerSegmentTimeoutConfig> getPerSegmentTimeoutConfig()
     return perSegmentTimeoutConfig;
   }
 
+  /**
+   * Query-specific broker dynamic config query context overrides (e.g. per segment timeout).
+   */
+  public QueryContext getQuerySpecificContextOverrides(Query<?> query)
+  {
+    if (perSegmentTimeoutConfig.isEmpty()) {
+      return QueryContext.empty();
+    }
+
+    for (String tableName : query.getDataSource().getTableNames()) {
+      PerSegmentTimeoutConfig dsConfig = perSegmentTimeoutConfig.get(tableName);
+      if (dsConfig != null) {
+        if (dsConfig.isMonitorOnly()) {
+          log.debug(
+              "Per-segment timeout [%d ms] configured for datasource [%s] in monitorOnly mode (not enforced) for query [%s].",
+              dsConfig.getPerSegmentTimeoutMs(),
+              tableName,
+              query.getId()
+          );
+          return QueryContext.empty();
+        }
+        return QueryContext.of(Map.of(QueryContexts.PER_SEGMENT_TIMEOUT_KEY, dsConfig.getPerSegmentTimeoutMs()));
+      }
+    }
+    return QueryContext.empty();
+  }
+
   @Override
   public boolean equals(Object o)
   {