This is an automated email from the ASF dual-hosted git repository.
jtuglu1 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/druid.git
The following commit(s) were added to refs/heads/master by this push:
new f813306c603 feat: Weighted Query Laning Strategy (#19225)
f813306c603 is described below
commit f813306c60357b32f25f16e659644e8f2f3f52af
Author: mshahid6 <[email protected]>
AuthorDate: Mon Jul 6 11:30:08 2026 -0700
feat: Weighted Query Laning Strategy (#19225)
Adds a new WeightedQueryLaningStrategy (strategy: "weighted") that scores
queries by how many configured thresholds they breach and assigns them to
graduated lanes with different capacity limits.
---
docs/configuration/index.md | 32 ++
.../apache/druid/server/QueryLaningStrategy.java | 4 +-
.../org/apache/druid/server/QueryScheduler.java | 10 +-
.../scheduling/ManualQueryLaningStrategy.java | 2 +-
.../scheduling/WeightedQueryLaningStrategy.java | 323 +++++++++++++++++
.../apache/druid/server/QuerySchedulerTest.java | 100 ++++++
.../WeightedQueryLaningStrategyTest.java | 389 +++++++++++++++++++++
7 files changed, 856 insertions(+), 4 deletions(-)
diff --git a/docs/configuration/index.md b/docs/configuration/index.md
index 92c2fc9103b..a89619b381a 100644
--- a/docs/configuration/index.md
+++ b/docs/configuration/index.md
@@ -1740,6 +1740,38 @@ This laning strategy is best suited for cases where one
or more external applica
|`druid.query.scheduler.laning.lanes.{name}`|Maximum percent or exact limit of
queries that can concurrently run in the defined lanes. Any number of lanes may
be defined like this. The lane names 'total' and 'default' are reserved for
internal use.|No default, must define at least one lane with a limit above 0.
If `druid.query.scheduler.laning.isLimitPercent` is set to `true`, values must
be integers in the range of 1 to 100.|
|`druid.query.scheduler.laning.isLimitPercent`|If set to `true`, the values
set for `druid.query.scheduler.laning.lanes` will be treated as a percent of
the smaller number of `druid.server.http.numThreads` or
`druid.query.scheduler.numThreads`. Note that in this mode, these lane values
across lanes are _not_ required to add up to, and can exceed, 100%.|`false`|
+###### 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-prior [...]
+
+If a lane is specified in the [query
context](../querying/query-context-reference.md) `lane` parameter, this will
override the computed lane.
+
+This strategy can be enabled by setting
`druid.query.scheduler.laning.strategy=weighted`.
+
+|Property|Description|Default|
+|--------|-----------|-------|
+|`druid.query.scheduler.laning.periodThreshold`|ISO 8601 period. A query is
charged 1 if any of its intervals starts before `now - periodThreshold` (that
is, it reads data older than this).|null (not evaluated)|
+|`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.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.sc [...]
+
+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%:
+
+```json
+{
+ "strategy": "weighted",
+ "periodThreshold": "P1M",
+ "durationThreshold": "P1D",
+ "segmentCountThreshold": 1000,
+ "segmentRangeThreshold": "P180D",
+ "lanes": {
+ "low": { "minCost": 1, "maxPercent": 30 },
+ "very-low": { "minCost": 3, "maxPercent": 10 }
+ }
+}
+```
+
##### Server configuration
Druid uses Jetty to serve HTTP requests. Each query being processed consumes a
single thread from `druid.server.http.numThreads`, so consider defining
`druid.query.scheduler.numThreads` to a lower value in order to reserve HTTP
threads for responding to health checks, lookup loading, and other non-query,
(in most cases) comparatively very short-lived, HTTP requests.
diff --git
a/server/src/main/java/org/apache/druid/server/QueryLaningStrategy.java
b/server/src/main/java/org/apache/druid/server/QueryLaningStrategy.java
index d601384c665..2e95f9d1b3a 100644
--- a/server/src/main/java/org/apache/druid/server/QueryLaningStrategy.java
+++ b/server/src/main/java/org/apache/druid/server/QueryLaningStrategy.java
@@ -28,6 +28,7 @@ import org.apache.druid.query.QueryPlus;
import org.apache.druid.server.scheduling.HiLoQueryLaningStrategy;
import org.apache.druid.server.scheduling.ManualQueryLaningStrategy;
import org.apache.druid.server.scheduling.NoQueryLaningStrategy;
+import org.apache.druid.server.scheduling.WeightedQueryLaningStrategy;
import java.util.Optional;
import java.util.Set;
@@ -37,7 +38,8 @@ import java.util.Set;
@JsonSubTypes(value = {
@JsonSubTypes.Type(name = "none", value = NoQueryLaningStrategy.class),
@JsonSubTypes.Type(name = "hilo", value = HiLoQueryLaningStrategy.class),
- @JsonSubTypes.Type(name = "manual", value =
ManualQueryLaningStrategy.class)
+ @JsonSubTypes.Type(name = "manual", value =
ManualQueryLaningStrategy.class),
+ @JsonSubTypes.Type(name = "weighted", value =
WeightedQueryLaningStrategy.class)
})
public interface QueryLaningStrategy
{
diff --git a/server/src/main/java/org/apache/druid/server/QueryScheduler.java
b/server/src/main/java/org/apache/druid/server/QueryScheduler.java
index 489af47e869..ad43d514890 100644
--- a/server/src/main/java/org/apache/druid/server/QueryScheduler.java
+++ b/server/src/main/java/org/apache/druid/server/QueryScheduler.java
@@ -66,6 +66,12 @@ public class QueryScheduler implements QueryWatcher
private static final Logger LOGGER = new Logger(QueryScheduler.class);
public static final int UNAVAILABLE = -1;
public static final String TOTAL = "total";
+ /**
+ * Lane name used for queries that are not assigned to any explicit lane
(interactive pool). It has special meaning
+ * for the resilience4j bulkhead used to enforce lane limits, and is also
emitted as the {@code lane} metric dimension
+ * for un-laned queries. It is therefore reserved and cannot be used as a
configured lane name.
+ */
+ public static final String DEFAULT = "default";
private final int totalCapacity;
private final QueryPrioritizationStrategy prioritizationStrategy;
private final QueryLaningStrategy laningStrategy;
@@ -165,12 +171,12 @@ public class QueryScheduler implements QueryWatcher
Optional<String> lane =
laningStrategy.computeLane(queryPlus.withQuery(query), segments);
LOGGER.debug(
"[%s] lane assigned to [%s] query with [%,d] priority",
- lane.orElse("default"),
+ lane.orElse(DEFAULT),
query.getType(),
priority.orElse(0)
);
final ServiceMetricEvent.Builder builderUsr =
ServiceMetricEvent.builder().setFeed("metrics")
-
.setDimension("lane", lane.orElse("default"))
+
.setDimension("lane", lane.orElse(DEFAULT))
.setDimension("dataSource", query.getDataSource().getTableNames())
.setDimension("type", query.getType());
emitter.emit(builderUsr.setMetric("query/priority",
priority.orElse(Integer.valueOf(0))));
diff --git
a/server/src/main/java/org/apache/druid/server/scheduling/ManualQueryLaningStrategy.java
b/server/src/main/java/org/apache/druid/server/scheduling/ManualQueryLaningStrategy.java
index 0f17adfae78..146a07e5038 100644
---
a/server/src/main/java/org/apache/druid/server/scheduling/ManualQueryLaningStrategy.java
+++
b/server/src/main/java/org/apache/druid/server/scheduling/ManualQueryLaningStrategy.java
@@ -64,7 +64,7 @@ public class ManualQueryLaningStrategy implements
QueryLaningStrategy
// 'default' has special meaning for resilience4j bulkhead used by query
scheduler, this restriction
// can potentially be relaxed if we ever change enforcement mechanism
Preconditions.checkArgument(
- lanes.keySet().stream().noneMatch("default"::equals),
+ lanes.keySet().stream().noneMatch(QueryScheduler.DEFAULT::equals),
"Lane cannot be named 'default'"
);
}
diff --git
a/server/src/main/java/org/apache/druid/server/scheduling/WeightedQueryLaningStrategy.java
b/server/src/main/java/org/apache/druid/server/scheduling/WeightedQueryLaningStrategy.java
new file mode 100644
index 00000000000..b3f075981c7
--- /dev/null
+++
b/server/src/main/java/org/apache/druid/server/scheduling/WeightedQueryLaningStrategy.java
@@ -0,0 +1,323 @@
+/*
+ * 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.druid.server.scheduling;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.google.common.base.Preconditions;
+import it.unimi.dsi.fastutil.objects.Object2IntArrayMap;
+import it.unimi.dsi.fastutil.objects.Object2IntMap;
+import org.apache.druid.client.SegmentServerSelector;
+import org.apache.druid.java.util.common.DateTimes;
+import org.apache.druid.query.Query;
+import org.apache.druid.query.QueryPlus;
+import org.apache.druid.server.QueryLaningStrategy;
+import org.apache.druid.server.QueryScheduler;
+import org.joda.time.DateTime;
+import org.joda.time.Duration;
+import org.joda.time.Interval;
+import org.joda.time.Period;
+import org.joda.time.base.AbstractInterval;
+
+import javax.annotation.Nullable;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Set;
+
+/**
+ * Query laning strategy that assigns a cost to queries based on how many
configured thresholds
+ * they breach, then assigns them to the most restrictive matching lane. This
provides more nuanced
+ * lane assignment than {@link HiLoQueryLaningStrategy}, which uses a binary
high/low split.
+ *
+ * <p>Configuration example:
+ * <pre>{@code
+ * {
+ * "strategy": "weighted",
+ * "periodThreshold": "P1M",
+ * "segmentCountThreshold": 1000,
+ * "lanes": {
+ * "low": { "minCost": 1, "maxPercent": 30 },
+ * "very-low": { "minCost": 3, "maxPercent": 10 }
+ * }
+ * }
+ * }</pre>
+ */
+public class WeightedQueryLaningStrategy implements QueryLaningStrategy
+{
+ @JsonProperty
+ @Nullable
+ private final Integer segmentCountThreshold;
+ // Stored as the original config strings so that Jackson round-trip serde is
consistent:
+ // the @JsonCreator constructor accepts String parameters, so serialization
must also emit
+ // strings. Storing Period/Duration objects directly with @JsonProperty
would cause Jackson
+ // to serialize them as complex JSON objects that the constructor cannot
deserialize back.
+ @JsonProperty("periodThreshold")
+ @Nullable
+ private final String periodThresholdString;
+ @JsonProperty("durationThreshold")
+ @Nullable
+ private final String durationThresholdString;
+ @JsonProperty("segmentRangeThreshold")
+ @Nullable
+ private final String segmentRangeThresholdString;
+
+ // Parsed from the string fields above at construction time; not serialized.
+ @Nullable
+ private final Period periodThreshold;
+ @Nullable
+ private final Duration durationThreshold;
+ @Nullable
+ private final Duration segmentRangeThreshold;
+
+ @JsonProperty
+ private final Map<String, LaneConfig> lanes;
+
+ @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
+ )
+ {
+ final Period parsedPeriod;
+ if (periodThresholdString == null) {
+ parsedPeriod = null;
+ } else {
+ try {
+ parsedPeriod = new Period(periodThresholdString);
+ }
+ catch (IllegalArgumentException e) {
+ throw new IllegalArgumentException(
+ "periodThreshold is not a valid ISO 8601 period, got [" +
periodThresholdString + "]"
+ );
+ }
+ }
+ final Duration parsedDuration;
+ if (durationThresholdString == null) {
+ parsedDuration = null;
+ } else {
+ try {
+ parsedDuration = new
Period(durationThresholdString).toStandardDuration();
+ }
+ catch (IllegalArgumentException | UnsupportedOperationException e) {
+ throw new IllegalArgumentException(
+ "durationThreshold is not a valid ISO 8601 duration, "
+ + "got [" + durationThresholdString + "]"
+ );
+ }
+ }
+ final Duration parsedSegmentRange;
+ if (segmentRangeThresholdString == null) {
+ parsedSegmentRange = null;
+ } else {
+ try {
+ parsedSegmentRange = new
Period(segmentRangeThresholdString).toStandardDuration();
+ }
+ catch (IllegalArgumentException | UnsupportedOperationException e) {
+ throw new IllegalArgumentException(
+ "segmentRangeThreshold is not a valid ISO 8601 duration, "
+ + "got [" + segmentRangeThresholdString + "]"
+ );
+ }
+ }
+
+ Preconditions.checkArgument(
+ segmentCountThreshold != null || parsedPeriod != null ||
parsedDuration != null || parsedSegmentRange != null,
+ "At least one of periodThreshold, durationThreshold,
segmentCountThreshold, or segmentRangeThreshold must be set"
+ );
+ Preconditions.checkArgument(
+ segmentCountThreshold == null || segmentCountThreshold > 0,
+ "segmentCountThreshold must be > 0, got [%s]", segmentCountThreshold
+ );
+ if (parsedDuration != null) {
+ Preconditions.checkArgument(parsedDuration.getMillis() > 0,
"durationThreshold must be positive, got [%s]", durationThresholdString);
+ }
+ if (parsedSegmentRange != null) {
+ Preconditions.checkArgument(parsedSegmentRange.getMillis() > 0,
"segmentRangeThreshold must be positive, got [%s]",
segmentRangeThresholdString);
+ }
+ if (parsedPeriod != null) {
+ DateTime now = DateTimes.nowUtc();
+ Preconditions.checkArgument(
+ now.minus(parsedPeriod.toDurationFrom(now)).isBefore(now),
+ "periodThreshold must be positive, got [%s]", periodThresholdString
+ );
+ }
+ Preconditions.checkArgument(
+ lanes != null && !lanes.isEmpty(),
+ "At least one lane must be defined"
+ );
+ Preconditions.checkArgument(
+ !lanes.containsKey(QueryScheduler.TOTAL),
+ "Lane cannot be named '%s'", QueryScheduler.TOTAL
+ );
+ Preconditions.checkArgument(
+ !lanes.containsKey(QueryScheduler.DEFAULT),
+ "Lane cannot be named '%s'", QueryScheduler.DEFAULT
+ );
+ long distinctCosts =
lanes.values().stream().mapToInt(LaneConfig::getMinCost).distinct().count();
+ Preconditions.checkArgument(
+ distinctCosts == lanes.size(),
+ "Each lane must have a unique minCost so that lane selection is
deterministic "
+ + "(a query is assigned to the lane with the highest minCost it meets;
equal costs "
+ + "produce non-deterministic results). Found duplicate minCost values
in lanes: [%s]",
+ lanes
+ );
+
+ this.segmentCountThreshold = segmentCountThreshold;
+ this.periodThresholdString = periodThresholdString;
+ this.durationThresholdString = durationThresholdString;
+ this.segmentRangeThresholdString = segmentRangeThresholdString;
+ this.periodThreshold = parsedPeriod;
+ this.durationThreshold = parsedDuration;
+ this.segmentRangeThreshold = parsedSegmentRange;
+ this.lanes = lanes;
+ }
+
+ @Override
+ public Object2IntMap<String> getLaneLimits(int totalLimit)
+ {
+ Object2IntMap<String> limits = new Object2IntArrayMap<>(lanes.size());
+ for (Map.Entry<String, LaneConfig> entry : lanes.entrySet()) {
+ limits.put(entry.getKey(), computeLimitFromPercent(totalLimit,
entry.getValue().maxPercent));
+ }
+ return limits;
+ }
+
+ @Override
+ public <T> Optional<String> computeLane(QueryPlus<T> query,
Set<SegmentServerSelector> segments)
+ {
+ final String existingLane = query.getQuery().context().getLane();
+ if (existingLane != null) {
+ return Optional.of(existingLane);
+ }
+
+ int cost = computeCost(query.getQuery(), segments);
+ if (cost == 0) {
+ return Optional.empty();
+ }
+
+ // Find the lane with the highest minCost that this query meets
+ String highestLane = null;
+ int highestMinCost = 0;
+ for (Map.Entry<String, LaneConfig> entry : lanes.entrySet()) {
+ int minCost = entry.getValue().minCost;
+ if (cost >= minCost && minCost > highestMinCost) {
+ highestLane = entry.getKey();
+ highestMinCost = minCost;
+ }
+ }
+ return Optional.ofNullable(highestLane);
+ }
+
+ private <T> int computeCost(Query<T> query, Set<SegmentServerSelector>
segments)
+ {
+ int cost = 0;
+
+ if (periodThreshold != null) {
+ final DateTime now = DateTimes.nowUtc();
+ final DateTime cutoff = now.minus(periodThreshold.toDurationFrom(now));
+ // Query intervals are condensed and sorted ascending by start (see
JodaUtils.condenseIntervals, applied by
+ // every QuerySegmentSpec), so the earliest start is the first interval.
Checking only it avoids scanning a
+ // query with hundreds of intervals.
+ final List<Interval> intervals = query.getIntervals();
+ if (!intervals.isEmpty() &&
intervals.get(0).getStart().isBefore(cutoff)) {
+ cost++;
+ }
+ }
+
+ if (durationThreshold != null &&
query.getDuration().isLongerThan(durationThreshold)) {
+ cost++;
+ }
+
+ if (segmentCountThreshold != null && segments.size() >
segmentCountThreshold) {
+ cost++;
+ }
+
+ if (segmentRangeThreshold != null) {
+ long segmentRangeMs = segments.stream()
+ .filter(s -> s.getSegmentDescriptor() !=
null)
+ .map(s ->
s.getSegmentDescriptor().getInterval())
+ .distinct()
+
.mapToLong(AbstractInterval::toDurationMillis)
+ .sum();
+ if (segmentRangeMs > segmentRangeThreshold.getMillis()) {
+ cost++;
+ }
+ }
+
+ return cost;
+ }
+
+ public static class LaneConfig
+ {
+ private final int minCost;
+ private final int maxPercent;
+
+ @JsonCreator
+ public LaneConfig(
+ @JsonProperty("minCost") int minCost,
+ @JsonProperty("maxPercent") int maxPercent
+ )
+ {
+ Preconditions.checkArgument(minCost > 0, "minCost must be > 0, got
[%s]", minCost);
+ Preconditions.checkArgument(
+ maxPercent > 0 && maxPercent <= 100,
+ "maxPercent must be in the range 1 to 100, got [%s]", maxPercent
+ );
+ this.minCost = minCost;
+ this.maxPercent = maxPercent;
+ }
+
+ @JsonProperty
+ public int getMinCost()
+ {
+ return minCost;
+ }
+
+ @JsonProperty
+ public int getMaxPercent()
+ {
+ return maxPercent;
+ }
+
+ @Override
+ public boolean equals(Object o)
+ {
+ if (this == o) {
+ return true;
+ }
+ if (o == null || getClass() != o.getClass()) {
+ return false;
+ }
+ LaneConfig that = (LaneConfig) o;
+ return minCost == that.minCost && maxPercent == that.maxPercent;
+ }
+
+ @Override
+ public int hashCode()
+ {
+ return Objects.hash(minCost, maxPercent);
+ }
+ }
+}
diff --git
a/server/src/test/java/org/apache/druid/server/QuerySchedulerTest.java
b/server/src/test/java/org/apache/druid/server/QuerySchedulerTest.java
index ed602678a82..9d584dccc7c 100644
--- a/server/src/test/java/org/apache/druid/server/QuerySchedulerTest.java
+++ b/server/src/test/java/org/apache/druid/server/QuerySchedulerTest.java
@@ -68,6 +68,7 @@ import org.apache.druid.server.initialization.ServerConfig;
import org.apache.druid.server.scheduling.HiLoQueryLaningStrategy;
import org.apache.druid.server.scheduling.ManualQueryPrioritizationStrategy;
import org.apache.druid.server.scheduling.NoQueryLaningStrategy;
+import org.apache.druid.server.scheduling.WeightedQueryLaningStrategy;
import org.easymock.EasyMock;
import org.hamcrest.text.StringContainsInOrder;
import org.junit.After;
@@ -79,7 +80,9 @@ import java.io.IOException;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
+import java.util.Map;
import java.util.Properties;
+import java.util.Set;
import java.util.UUID;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Future;
@@ -606,6 +609,103 @@ public class QuerySchedulerTest
Assert.assertEquals(QueryScheduler.UNAVAILABLE,
scheduler.getLaneAvailableCapacity("non-existent"));
}
+ @Test
+ public void testConfigWeighted()
+ {
+ final Injector injector = createInjector();
+ final String propertyPrefix = "druid.query.scheduler";
+ final JsonConfigProvider<QuerySchedulerProvider> provider =
JsonConfigProvider.of(
+ propertyPrefix,
+ QuerySchedulerProvider.class
+ );
+ final Properties properties = new Properties();
+ properties.setProperty(propertyPrefix + ".numThreads", "10");
+ properties.setProperty(propertyPrefix + ".laning.strategy", "weighted");
+ properties.setProperty(propertyPrefix + ".laning.segmentCountThreshold",
"1");
+ properties.setProperty(propertyPrefix + ".laning.lanes", "{\"low\":
{\"minCost\": 1, \"maxPercent\": 30}}");
+ provider.inject(properties, injector.getInstance(JsonConfigurator.class));
+ final QueryScheduler scheduler = provider.get().get();
+ Assert.assertEquals(10, scheduler.getTotalAvailableCapacity());
+ Assert.assertEquals(3, scheduler.getLaneAvailableCapacity("low"));
+ Assert.assertEquals(QueryScheduler.UNAVAILABLE,
scheduler.getLaneAvailableCapacity("non-existent"));
+ }
+
+ @Test
+ public void testWeightedLaneAssignment()
+ {
+ ObservableQueryScheduler weightedScheduler = new ObservableQueryScheduler(
+ TEST_HI_CAPACITY,
+ ManualQueryPrioritizationStrategy.INSTANCE,
+ new WeightedQueryLaningStrategy(
+ null,
+ null,
+ 1,
+ null,
+ Map.of("low", new WeightedQueryLaningStrategy.LaneConfig(1, 40))
+ ),
+ SERVER_CONFIG_WITH_TOTAL
+ );
+
+ // Query with 2 segments exceeds segmentCountThreshold=1 → "low" lane
+ Query<?> query = weightedScheduler.prioritizeAndLaneQuery(
+ QueryPlus.wrap(makeDefaultQuery()),
+ Set.of(
+ EasyMock.createMock(SegmentServerSelector.class),
+ EasyMock.createMock(SegmentServerSelector.class)
+ )
+ );
+ Assert.assertEquals("low", query.context().getLane());
+
+ // Query with 0 segments → no lane
+ Query<?> noLaneQuery = weightedScheduler.prioritizeAndLaneQuery(
+ QueryPlus.wrap(makeDefaultQuery()),
+ Set.of()
+ );
+ Assert.assertNull(noLaneQuery.context().getLane());
+ }
+
+ @Test
+ public void testWeightedFailsWhenOutOfLaneCapacity()
+ {
+ ObservableQueryScheduler weightedScheduler = new ObservableQueryScheduler(
+ TEST_HI_CAPACITY,
+ ManualQueryPrioritizationStrategy.INSTANCE,
+ new WeightedQueryLaningStrategy(
+ null,
+ null,
+ 1,
+ null,
+ Map.of("low", new WeightedQueryLaningStrategy.LaneConfig(1, 40))
+ ),
+ SERVER_CONFIG_WITH_TOTAL
+ );
+
+ Set<SegmentServerSelector> manySegments = Set.of(
+ EasyMock.createMock(SegmentServerSelector.class),
+ EasyMock.createMock(SegmentServerSelector.class)
+ );
+
+ // Fill the low lane (capacity = ceil(5 * 40/100) = 2)
+ Query<?> q1 =
weightedScheduler.prioritizeAndLaneQuery(QueryPlus.wrap(makeDefaultQuery()),
manySegments);
+ Yielders.each(weightedScheduler.run(q1, Sequences.empty()));
+ Assert.assertEquals(1, weightedScheduler.getLaneAvailableCapacity("low"));
+
+ Query<?> q2 =
weightedScheduler.prioritizeAndLaneQuery(QueryPlus.wrap(makeDefaultQuery()),
manySegments);
+ Yielders.each(weightedScheduler.run(q2, Sequences.empty()));
+ Assert.assertEquals(0, weightedScheduler.getLaneAvailableCapacity("low"));
+
+ // Third should fail with 429
+ Assert.assertThrows(
+ QueryCapacityExceededException.class,
+ () -> Yielders.each(
+ weightedScheduler.run(
+
weightedScheduler.prioritizeAndLaneQuery(QueryPlus.wrap(makeDefaultQuery()),
manySegments),
+ Sequences.empty()
+ )
+ )
+ );
+ }
+
private void maybeDelayNextIteration(int i) throws InterruptedException
{
if (i > 0 && i % 10 == 0) {
diff --git
a/server/src/test/java/org/apache/druid/server/scheduling/WeightedQueryLaningStrategyTest.java
b/server/src/test/java/org/apache/druid/server/scheduling/WeightedQueryLaningStrategyTest.java
new file mode 100644
index 00000000000..65eeb50d4da
--- /dev/null
+++
b/server/src/test/java/org/apache/druid/server/scheduling/WeightedQueryLaningStrategyTest.java
@@ -0,0 +1,389 @@
+/*
+ * 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.druid.server.scheduling;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import it.unimi.dsi.fastutil.objects.Object2IntMap;
+import org.apache.druid.client.SegmentServerSelector;
+import org.apache.druid.java.util.common.Intervals;
+import org.apache.druid.java.util.common.granularity.Granularities;
+import org.apache.druid.query.Druids;
+import org.apache.druid.query.QueryContexts;
+import org.apache.druid.query.QueryPlus;
+import org.apache.druid.query.SegmentDescriptor;
+import org.apache.druid.query.aggregation.CountAggregatorFactory;
+import org.apache.druid.query.timeseries.TimeseriesQuery;
+import org.apache.druid.segment.TestHelper;
+import org.apache.druid.server.QueryLaningStrategy;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+
+public class WeightedQueryLaningStrategyTest
+{
+ private static final Map<String, WeightedQueryLaningStrategy.LaneConfig>
TWO_LANES = Map.of(
+ "low", new WeightedQueryLaningStrategy.LaneConfig(1, 30),
+ "very-low", new WeightedQueryLaningStrategy.LaneConfig(3, 10)
+ );
+
+ private Druids.TimeseriesQueryBuilder queryBuilder;
+
+ @Before
+ public void setup()
+ {
+ queryBuilder = Druids.newTimeseriesQueryBuilder()
+ .dataSource("test")
+
.intervals(List.of(Intervals.of("2020-01-01/2020-01-02")))
+ .granularity(Granularities.DAY)
+ .aggregators(new CountAggregatorFactory("count"));
+ }
+
+ @Test
+ public void testGetLaneLimits()
+ {
+ WeightedQueryLaningStrategy strategy = newStrategy(null, null, 10, null);
+
+ Object2IntMap<String> limits = strategy.getLaneLimits(100);
+ Assert.assertEquals(2, limits.size());
+ Assert.assertEquals(30, limits.getInt("low"));
+ Assert.assertEquals(10, limits.getInt("very-low"));
+ }
+
+ @Test
+ public void testComputeLane_noViolations()
+ {
+ WeightedQueryLaningStrategy strategy = newStrategy(null, null, 10000,
null);
+ TimeseriesQuery query = queryBuilder.build();
+ Optional<String> lane = strategy.computeLane(QueryPlus.wrap(query),
Set.of());
+ Assert.assertFalse(lane.isPresent());
+ }
+
+ @Test
+ public void testComputeLane_oneViolation_segmentCount()
+ {
+ // segmentCountThreshold=1, query has 5 segments → cost=1 → "low"
+ WeightedQueryLaningStrategy strategy = newStrategy(null, null, 1, null);
+ TimeseriesQuery query = queryBuilder.build();
+ Set<SegmentServerSelector> segments = makeSegments(5);
+ Optional<String> lane = strategy.computeLane(QueryPlus.wrap(query),
segments);
+ Assert.assertTrue(lane.isPresent());
+ Assert.assertEquals("low", lane.get());
+ }
+
+ @Test
+ public void testComputeLane_twoViolations_matchesLowerMinCost()
+ {
+ // segmentCountThreshold=1 + durationThreshold=PT1S → cost=2
+ // Matches "low" (minCost=1) but NOT "very-low" (minCost=3)
+ WeightedQueryLaningStrategy strategy = new WeightedQueryLaningStrategy(
+ null,
+ "PT1S",
+ 1,
+ null,
+ TWO_LANES
+ );
+ TimeseriesQuery query = queryBuilder.build();
+ Set<SegmentServerSelector> segments = makeSegments(5);
+ Optional<String> lane = strategy.computeLane(QueryPlus.wrap(query),
segments);
+ Assert.assertTrue(lane.isPresent());
+ Assert.assertEquals("low", lane.get());
+ }
+
+ @Test
+ public void testComputeLane_allViolations_mostRestrictiveLane()
+ {
+ // All 4 thresholds set very low → cost=4 → meets "very-low" (minCost=3)
+ WeightedQueryLaningStrategy strategy = new WeightedQueryLaningStrategy(
+ "PT1S",
+ "PT1S",
+ 1,
+ "PT1S",
+ TWO_LANES
+ );
+ TimeseriesQuery query = queryBuilder.build();
+ Set<SegmentServerSelector> segments = makeSegments(5);
+ Optional<String> lane = strategy.computeLane(QueryPlus.wrap(query),
segments);
+ Assert.assertTrue(lane.isPresent());
+ Assert.assertEquals("very-low", lane.get());
+ }
+
+ @Test
+ public void testComputeLane_existingLane_preserved()
+ {
+ WeightedQueryLaningStrategy strategy = newStrategy(null, null, 10000,
null);
+ TimeseriesQuery query = queryBuilder
+ .context(Map.of(QueryContexts.LANE_KEY, "custom"))
+ .build();
+ Optional<String> lane = strategy.computeLane(QueryPlus.wrap(query),
Set.of());
+ Assert.assertTrue(lane.isPresent());
+ Assert.assertEquals("custom", lane.get());
+ }
+
+ @Test
+ public void testComputeLane_segmentRangeWithDifferentIntervals()
+ {
+ // Segments spanning different intervals to validate segmentRange summing
+ WeightedQueryLaningStrategy strategy = newStrategy(null, null, null,
"PT1S");
+ TimeseriesQuery query = queryBuilder.build();
+ Set<SegmentServerSelector> segments = new HashSet<>();
+ segments.add(new SegmentServerSelector(
+ new SegmentDescriptor(Intervals.of("2020-01-01/2020-01-02"), "v1", 0)
+ ));
+ segments.add(new SegmentServerSelector(
+ new SegmentDescriptor(Intervals.of("2020-01-02/2020-01-03"), "v1", 1)
+ ));
+ // Total range = 1 day + 1 day = 2 days > 1 second → cost=1 → "low"
+ Optional<String> lane = strategy.computeLane(QueryPlus.wrap(query),
segments);
+ Assert.assertTrue(lane.isPresent());
+ Assert.assertEquals("low", lane.get());
+ }
+
+ @Test
+ public void testComputeLane_periodThreshold_usesEarliestInterval()
+ {
+ // periodThreshold=P1D → cutoff is ~1 day ago. Provide multiple intervals
in unsorted order where only the
+ // earliest (old) interval is before the cutoff. getIntervals() is
condensed/sorted ascending by start, so the
+ // period check inspects the earliest interval and must charge cost even
though later intervals are recent.
+ WeightedQueryLaningStrategy strategy = newStrategy("P1D", null, null,
null);
+ TimeseriesQuery query = queryBuilder
+ .intervals(List.of(
+ Intervals.of("2038-01-01/2038-01-02"),
+ Intervals.of("2000-01-01/2000-01-02")
+ ))
+ .build();
+ Optional<String> lane = strategy.computeLane(QueryPlus.wrap(query),
Set.of());
+ Assert.assertTrue(lane.isPresent());
+ Assert.assertEquals("low", lane.get());
+ }
+
+ @Test
+ public void testValidation_noThresholds()
+ {
+ Assert.assertThrows(
+ IllegalArgumentException.class,
+ () -> newStrategy(null, null, null, null)
+ );
+ }
+
+ @Test
+ public void testValidation_noLanes()
+ {
+ Assert.assertThrows(
+ IllegalArgumentException.class,
+ () -> new WeightedQueryLaningStrategy(
+ null,
+ null,
+ 10,
+ null,
+ Map.of()
+ )
+ );
+ }
+
+ @Test
+ public void testValidation_nullLanes()
+ {
+ Assert.assertThrows(
+ IllegalArgumentException.class,
+ () -> new WeightedQueryLaningStrategy(
+ null,
+ null,
+ 10,
+ null,
+ null
+ )
+ );
+ }
+
+ @Test
+ public void testValidation_reservedLaneNameTotal()
+ {
+ Assert.assertThrows(
+ IllegalArgumentException.class,
+ () -> new WeightedQueryLaningStrategy(
+ null,
+ null,
+ 10,
+ null,
+ Map.of("total", new WeightedQueryLaningStrategy.LaneConfig(1, 30))
+ )
+ );
+ }
+
+ @Test
+ public void testValidation_reservedLaneNameDefault()
+ {
+ Assert.assertThrows(
+ IllegalArgumentException.class,
+ () -> new WeightedQueryLaningStrategy(
+ null,
+ null,
+ 10,
+ null,
+ Map.of("default", new WeightedQueryLaningStrategy.LaneConfig(1,
30))
+ )
+ );
+ }
+
+ @Test
+ public void testLaneConfig_invalidMinCost()
+ {
+ Assert.assertThrows(
+ IllegalArgumentException.class,
+ () -> new WeightedQueryLaningStrategy.LaneConfig(0, 30)
+ );
+ }
+
+ @Test
+ public void testLaneConfig_invalidMaxPercent()
+ {
+ Assert.assertThrows(
+ IllegalArgumentException.class,
+ () -> new WeightedQueryLaningStrategy.LaneConfig(1, 0)
+ );
+ }
+
+ @Test
+ public void testSerde() throws Exception
+ {
+ ObjectMapper mapper = TestHelper.makeJsonMapper();
+ String json = "{\n"
+ + " \"strategy\": \"weighted\",\n"
+ + " \"segmentCountThreshold\": 1000,\n"
+ + " \"durationThreshold\": \"P1D\",\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);
+
+ Object2IntMap<String> limits = deserialized.getLaneLimits(100);
+ Assert.assertEquals(30, limits.getInt("low"));
+ Assert.assertEquals(10, limits.getInt("very-low"));
+ }
+
+ private static WeightedQueryLaningStrategy newStrategy(
+ String periodThreshold,
+ String durationThreshold,
+ Integer segmentCountThreshold,
+ String segmentRangeThreshold
+ )
+ {
+ return new WeightedQueryLaningStrategy(
+ periodThreshold,
+ durationThreshold,
+ segmentCountThreshold,
+ segmentRangeThreshold,
+ TWO_LANES
+ );
+ }
+
+ @Test
+ public void testValidation_segmentCountThresholdZero()
+ {
+ Assert.assertThrows(
+ IllegalArgumentException.class,
+ () -> new WeightedQueryLaningStrategy(null, null, 0, null, TWO_LANES)
+ );
+ }
+
+ @Test
+ public void testValidation_segmentCountThresholdNegative()
+ {
+ Assert.assertThrows(
+ IllegalArgumentException.class,
+ () -> new WeightedQueryLaningStrategy(null, null, -1, null, TWO_LANES)
+ );
+ }
+
+ @Test
+ public void testValidation_durationThresholdZero()
+ {
+ Assert.assertThrows(
+ IllegalArgumentException.class,
+ () -> new WeightedQueryLaningStrategy(null, "PT0S", null, null,
TWO_LANES)
+ );
+ }
+
+ @Test
+ public void testValidation_durationThresholdNegative()
+ {
+ Assert.assertThrows(
+ IllegalArgumentException.class,
+ () -> new WeightedQueryLaningStrategy(null, "-PT1S", null, null,
TWO_LANES)
+ );
+ }
+
+ @Test
+ public void testValidation_segmentRangeThresholdZero()
+ {
+ Assert.assertThrows(
+ IllegalArgumentException.class,
+ () -> new WeightedQueryLaningStrategy(null, null, null, "PT0S",
TWO_LANES)
+ );
+ }
+
+ @Test
+ public void testValidation_segmentRangeThresholdNegative()
+ {
+ Assert.assertThrows(
+ IllegalArgumentException.class,
+ () -> new WeightedQueryLaningStrategy(null, null, null, "-PT1S",
TWO_LANES)
+ );
+ }
+
+ @Test
+ public void testValidation_periodThresholdZero()
+ {
+ Assert.assertThrows(
+ IllegalArgumentException.class,
+ () -> new WeightedQueryLaningStrategy("PT0S", null, null, null,
TWO_LANES)
+ );
+ }
+
+ @Test
+ public void testValidation_periodThresholdNegative()
+ {
+ Assert.assertThrows(
+ IllegalArgumentException.class,
+ () -> new WeightedQueryLaningStrategy("-PT1S", null, null, null,
TWO_LANES)
+ );
+ }
+
+ private static Set<SegmentServerSelector> makeSegments(int count)
+ {
+ Set<SegmentServerSelector> segments = new HashSet<>();
+ for (int i = 0; i < count; i++) {
+ segments.add(new SegmentServerSelector(
+ new SegmentDescriptor(Intervals.of("2020-01-01/2020-01-02"), "v1", i)
+ ));
+ }
+ return segments;
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
