Weighted laning: configurable per-threshold weights

docs/configuration/index.md

@@ -1742,7 +1742,7 @@ This laning strategy is best suited for cases where one or more external applica
 
 ###### Weighted laning strategy
 
-This laning strategy assigns a cost to each query based on how many of a configurable set of thresholds it breaches, then routes the query to the most restrictive lane whose `minCost` the query meets. It is a more granular alternative to the 'High/Low' strategy: rather than a binary split, a query that breaches one threshold can be assigned to a different lane than one that breaches several. The thresholds are the same ones used by the [threshold prioritization strategy](#threshold-prioritization-strategy) (data age, query interval duration, segment count, and total segment range). Each threshold a query breaches adds 1 to its cost. A query with cost 0 is not assigned a lane and runs in the interactive (default) pool.
+This laning strategy assigns a cost to each query based on how many of a configurable set of thresholds it breaches, then routes the query to the most restrictive lane whose `minCost` the query meets. It is a more granular alternative to the 'High/Low' strategy: rather than a binary split, a query that breaches one threshold can be assigned to a different lane than one that breaches several. The thresholds are the same ones used by the [threshold prioritization strategy](#threshold-prioritization-strategy) (data age, query interval duration, segment count, and total segment range). Each threshold a query breaches adds that threshold's weight to its cost — by default 1, or a larger value if you set the corresponding per-threshold weight (`periodWeight`, `durationWeight`, `segmentCountWeight`, `segmentRangeWeight`), which lets a breach on one dimension count for more than another. A query with cost 0 is not assigned a lane and runs in the interactive (default) pool.
 
 If a lane is specified in the [query context](../querying/query-context-reference.md) `lane` parameter, this will override the computed lane.
 
@@ -1754,6 +1754,10 @@ This strategy can be enabled by setting `druid.query.scheduler.laning.strategy=w
 |`druid.query.scheduler.laning.durationThreshold`|ISO 8601 duration (must not contain month or year components). A query is charged 1 if its total interval duration exceeds this.|null (not evaluated)|
 |`druid.query.scheduler.laning.segmentCountThreshold`|A query is charged 1 if the number of segments it involves exceeds this. Must be greater than 0.|null (not evaluated)|
 |`druid.query.scheduler.laning.segmentRangeThreshold`|ISO 8601 duration (must not contain month or year components). A query is charged 1 if the summed time range of its distinct segments exceeds this.|null (not evaluated)|
+|`druid.query.scheduler.laning.periodWeight`|Optional weight added to a query's cost when it breaches `periodThreshold`. Must be an integer >= 1, and `periodThreshold` must be set.|1|
+|`druid.query.scheduler.laning.durationWeight`|Optional weight added to a query's cost when it breaches `durationThreshold`. Must be an integer >= 1, and `durationThreshold` must be set.|1|
+|`druid.query.scheduler.laning.segmentCountWeight`|Optional weight added to a query's cost when it breaches `segmentCountThreshold`. Must be an integer >= 1, and `segmentCountThreshold` must be set.|1|
+|`druid.query.scheduler.laning.segmentRangeWeight`|Optional weight added to a query's cost when it breaches `segmentRangeThreshold`. Must be an integer >= 1, and `segmentRangeThreshold` must be set.|1|
 |`druid.query.scheduler.laning.lanes`|A map of lane name to its configuration. At least one lane must be defined. Each lane's configuration has two fields: `minCost`, the minimum query cost required to be assigned to the lane (a query is assigned to the lane with the highest `minCost` it meets; must be greater than 0 and unique across lanes so lane selection is deterministic), and `maxPercent`, the maximum percent of the smaller number of `druid.server.http.numThreads` or `druid.query.scheduler.numThreads` that queries in the lane may use concurrently (an integer in the range 1 to 100). The lane names 'total' and 'default' are reserved for internal use.|No default, must define at least one lane|
 
 At least one of `periodThreshold`, `durationThreshold`, `segmentCountThreshold`, or `segmentRangeThreshold` must be set. For example, the following configuration routes queries breaching 1 or 2 thresholds to a `low` lane capped at 30% capacity, and queries breaching 3 or 4 thresholds to a `very-low` lane capped at 10%:

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

@@ -179,6 +179,12 @@ public <T> Query<T> prioritizeAndLaneQuery(QueryPlus<T> queryPlus, Set<SegmentSe
                                                                     .setDimension("lane", lane.orElse(DEFAULT))
                                                                     .setDimension("dataSource", query.getDataSource().getTableNames())
                                                                     .setDimension("type", query.getType());
+    // "id" (the query id) lets query/priority be attributed to an individual query rather than only aggregated at
+    // datasource/type granularity. High cardinality, so downstream pipelines must not forward it to low-cardinality
+    // stores (e.g. Atlas). Guarded because some (typically internal) queries carry no id and setDimension rejects null.
+    if (query.getId() != null) {
+      builderUsr.setDimension("id", query.getId());
+    }
     emitter.emit(builderUsr.setMetric("query/priority", priority.orElse(Integer.valueOf(0))));
     return lane.map(query::withLane).orElse(query);
   }

server/src/main/java/org/apache/druid/server/scheduling/WeightedQueryLaningStrategy.java

@@ -21,6 +21,7 @@
 
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import it.unimi.dsi.fastutil.objects.Object2IntArrayMap;
 import it.unimi.dsi.fastutil.objects.Object2IntMap;
@@ -91,13 +92,32 @@ public class WeightedQueryLaningStrategy implements QueryLaningStrategy
   @JsonProperty
   private final Map<String, LaneConfig> lanes;
 
+  // Optional per-threshold weight: how much a breach of that threshold adds to the query's cost. Null means the
+  // default weight of 1 (the original flat scoring). Lets a breach on one dimension count for more than another.
+  @JsonProperty("periodWeight")
+  @Nullable
+  private final Integer periodWeight;
+  @JsonProperty("durationWeight")
+  @Nullable
+  private final Integer durationWeight;
+  @JsonProperty("segmentCountWeight")
+  @Nullable
+  private final Integer segmentCountWeight;
+  @JsonProperty("segmentRangeWeight")
+  @Nullable
+  private final Integer segmentRangeWeight;
+
   @JsonCreator
   public WeightedQueryLaningStrategy(
       @JsonProperty("periodThreshold") @Nullable String periodThresholdString,
       @JsonProperty("durationThreshold") @Nullable String durationThresholdString,
       @JsonProperty("segmentCountThreshold") @Nullable Integer segmentCountThreshold,
       @JsonProperty("segmentRangeThreshold") @Nullable String segmentRangeThresholdString,
-      @JsonProperty("lanes") Map<String, LaneConfig> lanes
+      @JsonProperty("lanes") Map<String, LaneConfig> lanes,
+      @JsonProperty("periodWeight") @Nullable Integer periodWeight,
+      @JsonProperty("durationWeight") @Nullable Integer durationWeight,
+      @JsonProperty("segmentCountWeight") @Nullable Integer segmentCountWeight,
+      @JsonProperty("segmentRangeWeight") @Nullable Integer segmentRangeWeight
   )
   {
     final Period parsedPeriod;
@@ -183,6 +203,12 @@ public WeightedQueryLaningStrategy(
         + "produce non-deterministic results). Found duplicate minCost values in lanes: [%s]",
         lanes
     );
+    // Each weight, if set, must be >= 1 and must correspond to a threshold that is actually configured
+    // (otherwise the weight is silently inert, which is almost certainly a misconfiguration).
+    validateWeight("periodWeight", periodWeight, parsedPeriod != null, "periodThreshold");
+    validateWeight("durationWeight", durationWeight, parsedDuration != null, "durationThreshold");
+    validateWeight("segmentCountWeight", segmentCountWeight, segmentCountThreshold != null, "segmentCountThreshold");
+    validateWeight("segmentRangeWeight", segmentRangeWeight, parsedSegmentRange != null, "segmentRangeThreshold");
 
     this.segmentCountThreshold = segmentCountThreshold;
     this.periodThresholdString = periodThresholdString;
@@ -192,6 +218,47 @@ public WeightedQueryLaningStrategy(
     this.durationThreshold = parsedDuration;
     this.segmentRangeThreshold = parsedSegmentRange;
     this.lanes = lanes;
+    this.periodWeight = periodWeight;
+    this.durationWeight = durationWeight;
+    this.segmentCountWeight = segmentCountWeight;
+    this.segmentRangeWeight = segmentRangeWeight;
+  }
+
+  /**
+   * Convenience constructor without per-threshold weights (every breach adds 1). Retained so existing callers
+   * and tests need not pass the optional weight parameters.
+   */
+  @VisibleForTesting
+  public WeightedQueryLaningStrategy(
+      @Nullable String periodThresholdString,
+      @Nullable String durationThresholdString,
+      @Nullable Integer segmentCountThreshold,
+      @Nullable String segmentRangeThresholdString,
+      Map<String, LaneConfig> lanes
+  )
+  {
+    this(periodThresholdString, durationThresholdString, segmentCountThreshold, segmentRangeThresholdString, lanes,
+         null, null, null, null);
+  }
+
+  private static void validateWeight(String name, @Nullable Integer weight, boolean thresholdSet, String thresholdName)
+  {
+    if (weight == null) {
+      return;
+    }
+    Preconditions.checkArgument(weight >= 1, "%s must be >= 1, got [%s]", name, weight);
+    Preconditions.checkArgument(
+        thresholdSet,
+        "%s is set but %s is not configured, so the weight would have no effect", name, thresholdName
+    );
+  }
+
+  /**
+   * Weight (cost contribution) for a breached threshold; defaults to 1 when not configured.
+   */
+  private static int weightOrDefault(@Nullable Integer weight)
+  {
+    return weight == null ? 1 : weight;
   }
 
   @Override
@@ -242,16 +309,16 @@ private <T> int computeCost(Query<T> query, Set<SegmentServerSelector> segments)
       // query with hundreds of intervals.
       final List<Interval> intervals = query.getIntervals();
       if (!intervals.isEmpty() && intervals.get(0).getStart().isBefore(cutoff)) {
-        cost++;
+        cost += weightOrDefault(periodWeight);
       }
     }
 
     if (durationThreshold != null && query.getDuration().isLongerThan(durationThreshold)) {
-      cost++;
+      cost += weightOrDefault(durationWeight);
     }
 
     if (segmentCountThreshold != null && segments.size() > segmentCountThreshold) {
-      cost++;
+      cost += weightOrDefault(segmentCountWeight);
     }
 
     if (segmentRangeThreshold != null) {
@@ -262,7 +329,7 @@ private <T> int computeCost(Query<T> query, Set<SegmentServerSelector> segments)
                                     .mapToLong(AbstractInterval::toDurationMillis)
                                     .sum();
       if (segmentRangeMs > segmentRangeThreshold.getMillis()) {
-        cost++;
+        cost += weightOrDefault(segmentRangeWeight);
       }
     }
 

server/src/test/java/org/apache/druid/server/scheduling/WeightedQueryLaningStrategyTest.java

@@ -376,6 +376,82 @@ public void testValidation_periodThresholdNegative()
     );
   }
 
+  @Test
+  public void testComputeLane_segmentCountWeight_bumpsCost()
+  {
+    // A single segmentCount breach normally scores 1 -> "low"; weighting it 3 makes cost=3 -> "very-low".
+    WeightedQueryLaningStrategy strategy =
+        new WeightedQueryLaningStrategy(null, null, 1, null, TWO_LANES, null, null, 3, null);
+    TimeseriesQuery query = queryBuilder.build();
+    Optional<String> lane = strategy.computeLane(QueryPlus.wrap(query), makeSegments(5));
+    Assert.assertTrue(lane.isPresent());
+    Assert.assertEquals("very-low", lane.get());
+  }
+
+  @Test
+  public void testComputeLane_defaultWeightUnchanged()
+  {
+    // Same single segmentCount breach without weights stays "low" (weight defaults to 1).
+    WeightedQueryLaningStrategy strategy = new WeightedQueryLaningStrategy(null, null, 1, null, TWO_LANES);
+    Optional<String> lane = strategy.computeLane(QueryPlus.wrap(queryBuilder.build()), makeSegments(5));
+    Assert.assertTrue(lane.isPresent());
+    Assert.assertEquals("low", lane.get());
+  }
+
+  @Test
+  public void testComputeLane_weightsSumAcrossThresholds()
+  {
+    // Breach durationThreshold (weight 2) + segmentCountThreshold (weight 2) => cost 4 => "very-low".
+    WeightedQueryLaningStrategy strategy =
+        new WeightedQueryLaningStrategy(null, "PT1S", 1, null, TWO_LANES, null, 2, 2, null);
+    Optional<String> lane = strategy.computeLane(QueryPlus.wrap(queryBuilder.build()), makeSegments(5));
+    Assert.assertTrue(lane.isPresent());
+    Assert.assertEquals("very-low", lane.get());
+  }
+
+  @Test
+  public void testSerde_withWeights() throws Exception
+  {
+    ObjectMapper mapper = TestHelper.makeJsonMapper();
+    String json = "{\n"
+                  + "  \"strategy\": \"weighted\",\n"
+                  + "  \"segmentCountThreshold\": 1,\n"
+                  + "  \"segmentCountWeight\": 3,\n"
+                  + "  \"lanes\": {\n"
+                  + "    \"low\": { \"minCost\": 1, \"maxPercent\": 30 },\n"
+                  + "    \"very-low\": { \"minCost\": 3, \"maxPercent\": 10 }\n"
+                  + "  }\n"
+                  + "}";
+    QueryLaningStrategy deserialized = mapper.readValue(json, QueryLaningStrategy.class);
+    Assert.assertTrue(deserialized instanceof WeightedQueryLaningStrategy);
+    // weight 3 applied to the single segmentCount breach -> cost 3 -> "very-low"
+    Optional<String> lane = deserialized.computeLane(QueryPlus.wrap(queryBuilder.build()), makeSegments(5));
+    Assert.assertEquals("very-low", lane.orElse(null));
+  }
+
+  @Test
+  public void testValidation_weightBelowOne()
+  {
+    Assert.assertThrows(
+        IllegalArgumentException.class,
+        () -> new WeightedQueryLaningStrategy(null, null, 1, null, TWO_LANES, null, null, 0, null)
+    );
+    Assert.assertThrows(
+        IllegalArgumentException.class,
+        () -> new WeightedQueryLaningStrategy(null, null, 1, null, TWO_LANES, null, null, -1, null)
+    );
+  }
+
+  @Test
+  public void testValidation_weightForUnsetThreshold()
+  {
+    // periodWeight set but no periodThreshold configured -> reject (weight would be inert).
+    Assert.assertThrows(
+        IllegalArgumentException.class,
+        () -> new WeightedQueryLaningStrategy(null, null, 1, null, TWO_LANES, 5, null, null, null)
+    );
+  }
+
   private static Set<SegmentServerSelector> makeSegments(int count)
   {
     Set<SegmentServerSelector> segments = new HashSet<>();