This is an automated email from the ASF dual-hosted git repository.
kfaraz 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 0971e423682 perf: optimize coordinator HistoricalManagement duty
runtime (#19532)
0971e423682 is described below
commit 0971e423682fa683f500c2dd5233ce544734b553
Author: jtuglu1 <[email protected]>
AuthorDate: Thu Jul 16 06:45:07 2026 -0700
perf: optimize coordinator HistoricalManagement duty runtime (#19532)
Changes:
- `MarkOvershadowedSegmentsAsUnused` - build the set of datasources that
have
overshadowed segments first, then skip timeline construction for every
other datasource,
both when scanning servers and when checking zero-replica segments.
- `MarkEternityTombstonesAsUnused` - group overshadowed segments by
datasource up front,
then check each candidate tombstone only against its own datasource
segments, instead of
scanning the full overshadowed-segment list per candidate. Use cheaper
isOvershadowed check.
- `PrepareBalancerAndLoadQueues` - Use total counts already available in
ImmutableDruidDataSource
- `UpdateReplicationStatus` - Compute unavailable, under-replicated,
deep-storage only counts in a single pass
---
.../coordinator/ComputePlacementCostBenchmark.java | 141 ++++++++++++++++
.../MarkOvershadowedSegmentsAsUnusedBenchmark.java | 142 ++++++++++++++++
.../druid/client/ImmutableDruidDataSource.java | 6 +
.../druid/server/coordinator/DruidCluster.java | 2 +-
.../druid/server/coordinator/DruidCoordinator.java | 51 +++---
.../duty/MarkEternityTombstonesAsUnused.java | 94 +++++++----
.../duty/MarkOvershadowedSegmentsAsUnused.java | 32 +++-
.../duty/PrepareBalancerAndLoadQueues.java | 14 +-
.../coordinator/loading/SegmentReplicaCount.java | 10 ++
.../loading/SegmentReplicationStatus.java | 89 +++++++++-
.../coordinator/loading/SegmentStatsSnapshot.java | 78 +++++++++
.../loading/StrategicSegmentAssigner.java | 6 +-
.../loading/SegmentReplicationStatusTest.java | 187 +++++++++++++++++++++
13 files changed, 765 insertions(+), 87 deletions(-)
diff --git
a/benchmarks/src/test/java/org/apache/druid/server/coordinator/ComputePlacementCostBenchmark.java
b/benchmarks/src/test/java/org/apache/druid/server/coordinator/ComputePlacementCostBenchmark.java
new file mode 100644
index 00000000000..3ab9e0cbac2
--- /dev/null
+++
b/benchmarks/src/test/java/org/apache/druid/server/coordinator/ComputePlacementCostBenchmark.java
@@ -0,0 +1,141 @@
+/*
+ * 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.coordinator;
+
+import com.google.common.collect.ImmutableMap;
+import com.google.common.util.concurrent.ListeningExecutorService;
+import com.google.common.util.concurrent.MoreExecutors;
+import org.apache.druid.client.ImmutableDruidDataSource;
+import org.apache.druid.client.ImmutableDruidServer;
+import org.apache.druid.java.util.common.DateTimes;
+import org.apache.druid.java.util.common.Intervals;
+import org.apache.druid.server.coordination.DruidServerMetadata;
+import org.apache.druid.server.coordination.ServerType;
+import org.apache.druid.server.coordinator.balancer.CostBalancerStrategy;
+import org.apache.druid.server.coordinator.loading.TestLoadQueuePeon;
+import org.apache.druid.timeline.DataSegment;
+import org.apache.druid.timeline.SegmentId;
+import org.joda.time.Interval;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Level;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Param;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.TearDown;
+import org.openjdk.jmh.annotations.Warmup;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
+
+/**
+ * Benchmarks {@link CostBalancerStrategy#computePlacementCost}, the
per-server placement-cost computation the balancer
+ * invokes for every candidate server when placing a segment. The cost of a
single call scales with the number of
+ * interval buckets the server's segments occupy, so {@code historyDays} (the
span of daily intervals held by the
+ * server) is the primary parameter to vary.
+ */
+@State(Scope.Benchmark)
+@BenchmarkMode(Mode.AverageTime)
+@OutputTimeUnit(TimeUnit.MICROSECONDS)
+@Warmup(iterations = 3, time = 2)
+@Measurement(iterations = 5, time = 2)
+@Fork(1)
+public class ComputePlacementCostBenchmark
+{
+ private static final long DAY_MILLIS = TimeUnit.DAYS.toMillis(1);
+ private static final long T0 =
DateTimes.of("2026-01-01T00:00:00Z").getMillis();
+ private static final String DATASOURCE = "ds0";
+
+ /** Span of contiguous daily intervals held by the server. */
+ @Param({"180", "730", "3650"})
+ private int historyDays;
+
+ private ListeningExecutorService exec;
+ private ExposedCostBalancerStrategy strategy;
+ private ServerHolder server;
+ private DataSegment proposalSegment;
+
+ @Setup(Level.Trial)
+ public void setup()
+ {
+ exec = MoreExecutors.newDirectExecutorService();
+ strategy = new ExposedCostBalancerStrategy(exec);
+
+ final List<DataSegment> segments = new ArrayList<>(historyDays);
+ for (int day = 0; day < historyDays; day++) {
+ final long start = T0 - (long) (day + 1) * DAY_MILLIS;
+ segments.add(createSegment(Intervals.utc(start, start + DAY_MILLIS)));
+ }
+
+ final ImmutableDruidServer immutableServer = new ImmutableDruidServer(
+ new DruidServerMetadata("server", "host", null, 1L << 40, null,
ServerType.HISTORICAL, "_default_tier", 0),
+ 0L,
+ ImmutableMap.of(DATASOURCE, new ImmutableDruidDataSource(DATASOURCE,
Collections.emptyMap(), segments)),
+ segments.size()
+ );
+ server = new ServerHolder(immutableServer, new TestLoadQueuePeon());
+
+ proposalSegment = createSegment(Intervals.utc(T0 - DAY_MILLIS, T0));
+ }
+
+ @TearDown(Level.Trial)
+ public void tearDown()
+ {
+ exec.shutdownNow();
+ }
+
+ @Benchmark
+ public double computePlacementCost()
+ {
+ return strategy.computePlacementCost(proposalSegment, server);
+ }
+
+ private static DataSegment createSegment(Interval interval)
+ {
+ return DataSegment.builder(SegmentId.of(DATASOURCE, interval, "v1", 0))
+ .size(1)
+ .build();
+ }
+
+ /**
+ * Exposes the protected {@link CostBalancerStrategy#computePlacementCost}
so the benchmark exercises the production
+ * implementation directly.
+ */
+ private static class ExposedCostBalancerStrategy extends CostBalancerStrategy
+ {
+ ExposedCostBalancerStrategy(ListeningExecutorService exec)
+ {
+ super(exec);
+ }
+
+ @Override
+ public double computePlacementCost(DataSegment proposalSegment,
ServerHolder server)
+ {
+ return super.computePlacementCost(proposalSegment, server);
+ }
+ }
+}
diff --git
a/benchmarks/src/test/java/org/apache/druid/server/coordinator/MarkOvershadowedSegmentsAsUnusedBenchmark.java
b/benchmarks/src/test/java/org/apache/druid/server/coordinator/MarkOvershadowedSegmentsAsUnusedBenchmark.java
new file mode 100644
index 00000000000..07fab6bd3ba
--- /dev/null
+++
b/benchmarks/src/test/java/org/apache/druid/server/coordinator/MarkOvershadowedSegmentsAsUnusedBenchmark.java
@@ -0,0 +1,142 @@
+/*
+ * 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.coordinator;
+
+import com.google.common.collect.ImmutableMap;
+import org.apache.druid.client.DataSourcesSnapshot;
+import org.apache.druid.client.ImmutableDruidDataSource;
+import org.apache.druid.client.ImmutableDruidServer;
+import org.apache.druid.java.util.common.DateTimes;
+import org.apache.druid.server.coordination.DruidServerMetadata;
+import org.apache.druid.server.coordination.ServerType;
+import org.apache.druid.server.coordinator.balancer.RandomBalancerStrategy;
+import
org.apache.druid.server.coordinator.duty.MarkOvershadowedSegmentsAsUnused;
+import org.apache.druid.server.coordinator.loading.SegmentLoadQueueManager;
+import org.apache.druid.server.coordinator.loading.TestLoadQueuePeon;
+import org.apache.druid.timeline.DataSegment;
+import org.apache.druid.timeline.SegmentId;
+import org.joda.time.DateTime;
+import org.joda.time.Interval;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Level;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Param;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.Warmup;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
+
+/**
+ * Benchmarks the {@code MarkOvershadowedSegmentsAsUnused} coordinator duty's
{@code run} method against a cluster where
+ * most datasources have no overshadowed segments (as in production). The duty
builds segment timelines from the served
+ * segments only for the datasources that actually have overshadowed segments,
so the cost is dominated by how many
+ * datasources are relevant relative to the total served.
+ */
+@State(Scope.Benchmark)
+@BenchmarkMode(Mode.AverageTime)
+@OutputTimeUnit(TimeUnit.MILLISECONDS)
+@Warmup(iterations = 3, time = 2)
+@Measurement(iterations = 5, time = 2)
+@Fork(1)
+public class MarkOvershadowedSegmentsAsUnusedBenchmark
+{
+ private static final DateTime START = DateTimes.of("2024-01-01");
+
+ /** Total datasources served by the cluster. */
+ @Param({"1000"})
+ private int numDatasources;
+
+ /** Datasources that have overshadowed segments (an older version shadowed
by a newer one). */
+ @Param({"10", "100"})
+ private int relevantDatasources;
+
+ @Param({"4"})
+ private int intervalsPerDatasource;
+
+ private MarkOvershadowedSegmentsAsUnused duty;
+ private DruidCoordinatorRuntimeParams params;
+
+ @Setup(Level.Trial)
+ public void setup()
+ {
+ final List<DataSegment> usedSegments = new ArrayList<>();
+ final ImmutableMap.Builder<String, ImmutableDruidDataSource> dataSources =
ImmutableMap.builder();
+
+ for (int d = 0; d < numDatasources; d++) {
+ final String datasource = "ds_" + d;
+ final boolean hasOvershadowed = d < relevantDatasources;
+ final List<DataSegment> dsSegments = new ArrayList<>();
+ for (int i = 0; i < intervalsPerDatasource; i++) {
+ final Interval interval = new Interval(START.plusDays(i),
START.plusDays(i + 1));
+ dsSegments.add(segment(datasource, interval, "v1"));
+ if (hasOvershadowed) {
+ // A newer version overshadows the v1 segment for the same interval.
+ dsSegments.add(segment(datasource, interval, "v2"));
+ }
+ }
+ usedSegments.addAll(dsSegments);
+ dataSources.put(datasource, new ImmutableDruidDataSource(datasource,
ImmutableMap.of(), dsSegments));
+ }
+
+ final ImmutableDruidServer server = new ImmutableDruidServer(
+ new DruidServerMetadata("server", "host", null, 1L << 40, null,
ServerType.HISTORICAL, "_default_tier", 0),
+ 0L,
+ dataSources.build(),
+ usedSegments.size()
+ );
+ final DruidCluster cluster = DruidCluster
+ .builder()
+ .add(new ServerHolder(server, new TestLoadQueuePeon()))
+ .build();
+
+ params = DruidCoordinatorRuntimeParams
+ .builder()
+
.withDataSourcesSnapshot(DataSourcesSnapshot.fromUsedSegments(usedSegments))
+ .withDruidCluster(cluster)
+
.withDynamicConfigs(CoordinatorDynamicConfig.builder().withMarkSegmentAsUnusedDelayMillis(0).build())
+ .withBalancerStrategy(new RandomBalancerStrategy())
+ .withSegmentAssignerUsing(new SegmentLoadQueueManager(null, null))
+ .build();
+
+ // A no-op delete handler so the benchmark measures the timeline build +
overshadow check, not the metadata write.
+ duty = new MarkOvershadowedSegmentsAsUnused((datasource, segmentIds) ->
segmentIds.size());
+ }
+
+ @Benchmark
+ public DruidCoordinatorRuntimeParams run()
+ {
+ return duty.run(params);
+ }
+
+ private static DataSegment segment(String datasource, Interval interval,
String version)
+ {
+ return DataSegment.builder(SegmentId.of(datasource, interval, version, 0))
+ .size(1)
+ .build();
+ }
+}
diff --git
a/server/src/main/java/org/apache/druid/client/ImmutableDruidDataSource.java
b/server/src/main/java/org/apache/druid/client/ImmutableDruidDataSource.java
index 5120681fa64..e855ca8f1fc 100644
--- a/server/src/main/java/org/apache/druid/client/ImmutableDruidDataSource.java
+++ b/server/src/main/java/org/apache/druid/client/ImmutableDruidDataSource.java
@@ -107,6 +107,12 @@ public class ImmutableDruidDataSource
return Collections.unmodifiableCollection(idToSegments.values());
}
+ @JsonIgnore
+ public int getNumSegments()
+ {
+ return idToSegments.size();
+ }
+
@JsonIgnore
@Nullable
public DataSegment getSegment(SegmentId segmentId)
diff --git
a/server/src/main/java/org/apache/druid/server/coordinator/DruidCluster.java
b/server/src/main/java/org/apache/druid/server/coordinator/DruidCluster.java
index 5e3bcfd31de..2743cff0f4d 100644
--- a/server/src/main/java/org/apache/druid/server/coordinator/DruidCluster.java
+++ b/server/src/main/java/org/apache/druid/server/coordinator/DruidCluster.java
@@ -101,7 +101,7 @@ public class DruidCluster
return brokers;
}
- public Iterable<String> getTierNames()
+ public Set<String> getTierNames()
{
return historicals.keySet();
}
diff --git
a/server/src/main/java/org/apache/druid/server/coordinator/DruidCoordinator.java
b/server/src/main/java/org/apache/druid/server/coordinator/DruidCoordinator.java
index 41237e35299..dd12caf904e 100644
---
a/server/src/main/java/org/apache/druid/server/coordinator/DruidCoordinator.java
+++
b/server/src/main/java/org/apache/druid/server/coordinator/DruidCoordinator.java
@@ -86,6 +86,7 @@ import
org.apache.druid.server.coordinator.loading.LoadQueueTaskMaster;
import org.apache.druid.server.coordinator.loading.SegmentLoadQueueManager;
import org.apache.druid.server.coordinator.loading.SegmentReplicaCount;
import org.apache.druid.server.coordinator.loading.SegmentReplicationStatus;
+import org.apache.druid.server.coordinator.loading.SegmentStatsSnapshot;
import org.apache.druid.server.coordinator.stats.CoordinatorRunStats;
import org.apache.druid.server.coordinator.stats.CoordinatorStat;
import org.apache.druid.server.coordinator.stats.Dimension;
@@ -247,12 +248,10 @@ public class DruidCoordinator
final Iterable<DataSegment> dataSegments =
metadataManager.iterateAllUsedSegments();
for (DataSegment segment : dataSegments) {
- SegmentReplicaCount replicaCount =
segmentReplicationStatus.getReplicaCountsInCluster(segment.getId());
- if (replicaCount != null && (replicaCount.totalLoaded() > 0 ||
replicaCount.required() == 0)) {
- datasourceToUnavailableSegments.addTo(segment.getDataSource(), 0);
- } else {
- datasourceToUnavailableSegments.addTo(segment.getDataSource(), 1);
- }
+ datasourceToUnavailableSegments.addTo(
+ segment.getDataSource(),
+ segmentReplicationStatus.isUnavailable(segment.getId()) ? 1 : 0
+ );
}
return datasourceToUnavailableSegments;
@@ -268,8 +267,7 @@ public class DruidCoordinator
final Iterable<DataSegment> dataSegments =
metadataManager.iterateAllUsedSegments();
for (DataSegment segment : dataSegments) {
- SegmentReplicaCount replicaCount =
segmentReplicationStatus.getReplicaCountsInCluster(segment.getId());
- if (replicaCount != null && replicaCount.totalLoaded() == 0 &&
replicaCount.required() == 0) {
+ if (segmentReplicationStatus.isDeepStorageOnly(segment.getId())) {
datasourceToDeepStorageOnlySegments.addTo(segment.getDataSource(), 1);
}
}
@@ -784,9 +782,11 @@ public class DruidCoordinator
{
broadcastSegments = params.getBroadcastSegments();
segmentReplicationStatus = params.getSegmentReplicationStatus();
+
if (coordinatorSegmentMetadataCache != null) {
coordinatorSegmentMetadataCache.updateSegmentReplicationStatus(segmentReplicationStatus);
}
+
return params;
}
}
@@ -801,25 +801,28 @@ public class DruidCoordinator
{
// Collect stats for unavailable and under-replicated segments
final CoordinatorRunStats stats = params.getCoordinatorStats();
- getDatasourceToUnavailableSegmentCount().forEach(
- (dataSource, numUnavailable) -> stats.add(
- Stats.Segments.UNAVAILABLE,
- RowKey.of(Dimension.DATASOURCE, dataSource),
- numUnavailable
- )
+
+ // Compute unavailable, under-replicated, deep-storage only counts in a
single pass
+ final SegmentStatsSnapshot snapshot =
+ segmentReplicationStatus == null
+ ? SegmentStatsSnapshot.empty()
+ :
segmentReplicationStatus.computeSegmentStats(metadataManager.iterateAllUsedSegments(),
true);
+
+ snapshot.getDatasourceToUnavailableCount().forEach(
+ (datasource, count) ->
+ stats.add(Stats.Segments.UNAVAILABLE,
RowKey.of(Dimension.DATASOURCE, datasource), count)
);
- getTierToDatasourceToUnderReplicatedCount(false).forEach(
- (tier, countsPerDatasource) -> countsPerDatasource.forEach(
- (dataSource, underReplicatedCount) ->
- stats.addToSegmentStat(Stats.Segments.UNDER_REPLICATED,
tier, dataSource, underReplicatedCount)
+
+ snapshot.getTierToDatasourceToUnderReplicatedCount().forEach(
+ (tier, datasourceToCount) -> datasourceToCount.forEach(
+ (datasource, count) ->
+ stats.addToSegmentStat(Stats.Segments.UNDER_REPLICATED,
tier, datasource, count)
)
);
- getDatasourceToDeepStorageQueryOnlySegmentCount().forEach(
- (dataSource, numDeepStorageOnly) -> stats.add(
- Stats.Segments.DEEP_STORAGE_ONLY,
- RowKey.of(Dimension.DATASOURCE, dataSource),
- numDeepStorageOnly
- )
+
+ snapshot.getDatasourceToDeepStorageOnlyCount().forEach(
+ (datasource, count) ->
+ stats.add(Stats.Segments.DEEP_STORAGE_ONLY,
RowKey.of(Dimension.DATASOURCE, datasource), count)
);
return params;
diff --git
a/server/src/main/java/org/apache/druid/server/coordinator/duty/MarkEternityTombstonesAsUnused.java
b/server/src/main/java/org/apache/druid/server/coordinator/duty/MarkEternityTombstonesAsUnused.java
index f332f11020c..64c70c3a826 100644
---
a/server/src/main/java/org/apache/druid/server/coordinator/duty/MarkEternityTombstonesAsUnused.java
+++
b/server/src/main/java/org/apache/druid/server/coordinator/duty/MarkEternityTombstonesAsUnused.java
@@ -19,10 +19,9 @@
package org.apache.druid.server.coordinator.duty;
-import com.google.common.base.Optional;
import org.apache.druid.client.DataSourcesSnapshot;
+import org.apache.druid.client.ImmutableDruidDataSource;
import org.apache.druid.java.util.common.DateTimes;
-import org.apache.druid.java.util.common.Intervals;
import org.apache.druid.java.util.common.granularity.Granularities;
import org.apache.druid.java.util.common.logger.Logger;
import org.apache.druid.server.coordinator.DruidCoordinatorRuntimeParams;
@@ -31,7 +30,6 @@ import org.apache.druid.server.coordinator.stats.Dimension;
import org.apache.druid.server.coordinator.stats.RowKey;
import org.apache.druid.server.coordinator.stats.Stats;
import org.apache.druid.timeline.DataSegment;
-import org.apache.druid.timeline.Partitions;
import org.apache.druid.timeline.SegmentId;
import org.apache.druid.timeline.SegmentTimeline;
import org.apache.druid.timeline.partition.TombstoneShardSpec;
@@ -82,9 +80,9 @@ public class MarkEternityTombstonesAsUnused implements
CoordinatorDuty
@Override
public DruidCoordinatorRuntimeParams run(final DruidCoordinatorRuntimeParams
params)
{
- DataSourcesSnapshot dataSourcesSnapshot = params.getDataSourcesSnapshot();
+ final DataSourcesSnapshot dataSourcesSnapshot =
params.getDataSourcesSnapshot();
- final Map<String, Set<SegmentId>>
datasourceToNonOvershadowedEternityTombstones =
+ final Map<String, Set<SegmentId>>
datasourceToNonOvershadowedEternityTombstones =
determineNonOvershadowedEternityTombstones(
dataSourcesSnapshot
);
@@ -132,43 +130,69 @@ public class MarkEternityTombstonesAsUnused implements
CoordinatorDuty
private Map<String, Set<SegmentId>>
determineNonOvershadowedEternityTombstones(final DataSourcesSnapshot
dataSourcesSnapshot)
{
final Map<String, Set<SegmentId>>
datasourceToNonOvershadowedEternityTombstones = new HashMap<>();
+ final Map<String, Set<DataSegment>> overshadowedSegmentsByDatasource = new
HashMap<>();
- dataSourcesSnapshot.getDataSourcesMap().keySet().forEach((datasource) -> {
+ for (final DataSegment overshadowedSegment :
dataSourcesSnapshot.getOvershadowedSegments()) {
+ overshadowedSegmentsByDatasource
+ .computeIfAbsent(overshadowedSegment.getDataSource(), ds -> new
HashSet<>())
+ .add(overshadowedSegment);
+ }
+
+ for (final ImmutableDruidDataSource dataSource :
dataSourcesSnapshot.getDataSourcesWithAllUsedSegments()) {
+ final String datasourceName = dataSource.getName();
+ // The timeline is guaranteed to be present since it is derived from the
same set of
+ // datasources that we are iterating over.
final SegmentTimeline usedSegmentsTimeline
- =
dataSourcesSnapshot.getUsedSegmentsTimelinesPerDataSource().get(datasource);
-
- final Optional<Set<DataSegment>> usedNonOvershadowedSegments =
- Optional.fromNullable(usedSegmentsTimeline)
- .transform(timeline ->
timeline.findNonOvershadowedObjectsInInterval(
- Intervals.ETERNITY,
- Partitions.ONLY_COMPLETE
- ));
-
- if (usedNonOvershadowedSegments.isPresent()) {
- usedNonOvershadowedSegments.get().forEach(candidateSegment -> {
- if (isNewGenerationEternityTombstone(candidateSegment)) {
- boolean overlaps =
dataSourcesSnapshot.getOvershadowedSegments().stream()
- .filter(overshadowedSegment
->
-
candidateSegment.getDataSource()
-
.equals(overshadowedSegment.getDataSource()))
- .anyMatch(
- overshadowedSegment ->
-
candidateSegment.getInterval()
-
.overlaps(overshadowedSegment.getInterval())
- );
- if (!overlaps) {
- datasourceToNonOvershadowedEternityTombstones
- .computeIfAbsent(datasource, ds -> new HashSet<>())
- .add(candidateSegment.getId());
- }
- }
- });
+ =
dataSourcesSnapshot.getUsedSegmentsTimelinesPerDataSource().get(datasourceName);
+
+ for (final DataSegment candidateSegment : dataSource.getSegments()) {
+ if (shouldMarkAsUnused(candidateSegment,
overshadowedSegmentsByDatasource, usedSegmentsTimeline)) {
+ datasourceToNonOvershadowedEternityTombstones
+ .computeIfAbsent(datasourceName, ds -> new HashSet<>())
+ .add(candidateSegment.getId());
+ }
}
- });
+ }
return datasourceToNonOvershadowedEternityTombstones;
}
+ /**
+ * A candidate segment should be marked as unused if it is a new-generation
eternity tombstone that is neither
+ * overshadowed in the used-segments timeline nor overlapping any
overshadowed segment in its datasource.
+ */
+ private boolean shouldMarkAsUnused(
+ final DataSegment candidateSegment,
+ final Map<String, Set<DataSegment>> overshadowedSegmentsByDatasource,
+ final SegmentTimeline usedSegmentsTimeline
+ )
+ {
+ return isNewGenerationEternityTombstone(candidateSegment)
+ && !usedSegmentsTimeline.isOvershadowed(candidateSegment)
+ && !overlapsAnyOvershadowedSegment(
+ candidateSegment,
+
overshadowedSegmentsByDatasource.get(candidateSegment.getDataSource())
+ );
+ }
+
+ private boolean overlapsAnyOvershadowedSegment(
+ final DataSegment candidateSegment,
+ final Set<DataSegment> overshadowedSegments
+ )
+ {
+ if (overshadowedSegments == null) {
+ return false;
+ }
+
+ for (final DataSegment overshadowedSegment : overshadowedSegments) {
+ if
(candidateSegment.getInterval().overlaps(overshadowedSegment.getInterval())) {
+ return true;
+ }
+ }
+
+ return false;
+ }
+
private boolean isNewGenerationEternityTombstone(final DataSegment segment)
{
return segment.isTombstone() &&
segment.getShardSpec().getNumCorePartitions() == 0 && (
diff --git
a/server/src/main/java/org/apache/druid/server/coordinator/duty/MarkOvershadowedSegmentsAsUnused.java
b/server/src/main/java/org/apache/druid/server/coordinator/duty/MarkOvershadowedSegmentsAsUnused.java
index dba11d2c290..7772245ca5a 100644
---
a/server/src/main/java/org/apache/druid/server/coordinator/duty/MarkOvershadowedSegmentsAsUnused.java
+++
b/server/src/main/java/org/apache/druid/server/coordinator/duty/MarkOvershadowedSegmentsAsUnused.java
@@ -39,6 +39,7 @@ import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
+import java.util.stream.Collectors;
/**
* Marks a segment as unused if it is overshadowed by:
@@ -80,23 +81,33 @@ public class MarkOvershadowedSegmentsAsUnused implements
CoordinatorDuty
return params;
}
+ // Identify the datasources that actually have overshadowed segments to
check.
+ // Timelines only need to be built for these datasources.
+ final Set<String> eligibleDatasources = allOvershadowedSegments
+ .stream()
+ .map(DataSegment::getDataSource)
+ .collect(Collectors.toSet());
+
final DruidCluster cluster = params.getDruidCluster();
final Map<String, SegmentTimeline> timelines = new HashMap<>();
cluster.getManagedHistoricals().values().forEach(
historicals -> historicals.forEach(
- historical -> addSegmentsFromServer(historical, timelines)
+ historical -> addSegmentsFromServer(historical, timelines,
eligibleDatasources)
)
);
cluster.getBrokers().forEach(
- broker -> addSegmentsFromServer(broker, timelines)
+ broker -> addSegmentsFromServer(broker, timelines, eligibleDatasources)
);
// Include all segments that require zero replicas to be loaded
params.getSegmentAssigner().getSegmentsWithZeroRequiredReplicas().forEach(
- (datasource, segments) -> timelines
- .computeIfAbsent(datasource, ds -> new SegmentTimeline())
- .addSegments(segments.iterator())
+ (datasource, segments) -> {
+ if (eligibleDatasources.contains(datasource)) {
+ timelines.computeIfAbsent(datasource, ds -> new SegmentTimeline())
+ .addSegments(segments.iterator());
+ }
+ }
);
// Do not include segments served by ingestion services such as tasks or
indexers,
@@ -132,12 +143,17 @@ public class MarkOvershadowedSegmentsAsUnused implements
CoordinatorDuty
private void addSegmentsFromServer(
ServerHolder serverHolder,
- Map<String, SegmentTimeline> timelines
+ Map<String, SegmentTimeline> timelines,
+ Set<String> eligibleDatasources
)
{
- ImmutableDruidServer server = serverHolder.getServer();
+ final ImmutableDruidServer server = serverHolder.getServer();
+
+ for (final ImmutableDruidDataSource dataSource : server.getDataSources()) {
+ if (!eligibleDatasources.contains(dataSource.getName())) {
+ continue;
+ }
- for (ImmutableDruidDataSource dataSource : server.getDataSources()) {
timelines
.computeIfAbsent(dataSource.getName(), dsName -> new
SegmentTimeline())
.addSegments(dataSource.getSegments().iterator());
diff --git
a/server/src/main/java/org/apache/druid/server/coordinator/duty/PrepareBalancerAndLoadQueues.java
b/server/src/main/java/org/apache/druid/server/coordinator/duty/PrepareBalancerAndLoadQueues.java
index ec0ddcaf0d0..0ad2c1cb199 100644
---
a/server/src/main/java/org/apache/druid/server/coordinator/duty/PrepareBalancerAndLoadQueues.java
+++
b/server/src/main/java/org/apache/druid/server/coordinator/duty/PrepareBalancerAndLoadQueues.java
@@ -20,6 +20,7 @@
package org.apache.druid.server.coordinator.duty;
import org.apache.druid.client.DruidServer;
+import org.apache.druid.client.ImmutableDruidDataSource;
import org.apache.druid.client.ImmutableDruidServer;
import org.apache.druid.client.ServerInventoryView;
import org.apache.druid.java.util.common.logger.Logger;
@@ -36,7 +37,6 @@ import
org.apache.druid.server.coordinator.stats.CoordinatorRunStats;
import org.apache.druid.server.coordinator.stats.Dimension;
import org.apache.druid.server.coordinator.stats.RowKey;
import org.apache.druid.server.coordinator.stats.Stats;
-import org.apache.druid.timeline.DataSegment;
import java.util.HashSet;
import java.util.List;
@@ -203,14 +203,14 @@ public class PrepareBalancerAndLoadQueues implements
CoordinatorDuty
private void collectUsedSegmentStats(DruidCoordinatorRuntimeParams params,
CoordinatorRunStats stats)
{
- params.getUsedSegmentsTimelinesPerDataSource().forEach((dataSource,
timeline) -> {
- long totalSizeOfUsedSegments = timeline.iterateAllObjects().stream()
-
.mapToLong(DataSegment::getSize).sum();
+ for (final ImmutableDruidDataSource dataSource :
params.getDataSourcesSnapshot().getDataSourcesWithAllUsedSegments()) {
+ final long totalSizeOfUsedSegments = dataSource.getTotalSizeOfSegments();
+ final int numUsedSegments = dataSource.getNumSegments();
- RowKey datasourceKey = RowKey.of(Dimension.DATASOURCE, dataSource);
+ final RowKey datasourceKey = RowKey.of(Dimension.DATASOURCE,
dataSource.getName());
stats.add(Stats.Segments.USED_BYTES, datasourceKey,
totalSizeOfUsedSegments);
- stats.add(Stats.Segments.USED, datasourceKey, timeline.getNumObjects());
- });
+ stats.add(Stats.Segments.USED, datasourceKey, numUsedSegments);
+ }
}
private void collectDebugStats(SegmentLoadingConfig config,
CoordinatorRunStats stats)
diff --git
a/server/src/main/java/org/apache/druid/server/coordinator/loading/SegmentReplicaCount.java
b/server/src/main/java/org/apache/druid/server/coordinator/loading/SegmentReplicaCount.java
index 1f9650a59f0..8f0bf11b7d4 100644
---
a/server/src/main/java/org/apache/druid/server/coordinator/loading/SegmentReplicaCount.java
+++
b/server/src/main/java/org/apache/druid/server/coordinator/loading/SegmentReplicaCount.java
@@ -163,6 +163,16 @@ public class SegmentReplicaCount
return Math.max(requiredAndLoadable() - totalLoaded(), 0);
}
+ /**
+ * Returns a new, independent instance with the same counts as this one.
+ */
+ SegmentReplicaCount copy()
+ {
+ final SegmentReplicaCount copy = new SegmentReplicaCount();
+ copy.accumulate(this);
+ return copy;
+ }
+
/**
* Accumulates counts from the given {@code SegmentReplicaCount} into this
instance.
*/
diff --git
a/server/src/main/java/org/apache/druid/server/coordinator/loading/SegmentReplicationStatus.java
b/server/src/main/java/org/apache/druid/server/coordinator/loading/SegmentReplicationStatus.java
index 7121642f25e..ec64b00241d 100644
---
a/server/src/main/java/org/apache/druid/server/coordinator/loading/SegmentReplicationStatus.java
+++
b/server/src/main/java/org/apache/druid/server/coordinator/loading/SegmentReplicationStatus.java
@@ -19,12 +19,14 @@
package org.apache.druid.server.coordinator.loading;
-import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Maps;
+import it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap;
import it.unimi.dsi.fastutil.objects.Object2LongMap;
import it.unimi.dsi.fastutil.objects.Object2LongOpenHashMap;
import org.apache.druid.timeline.DataSegment;
import org.apache.druid.timeline.SegmentId;
+import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
@@ -40,15 +42,25 @@ public class SegmentReplicationStatus
public SegmentReplicationStatus(Map<SegmentId, Map<String,
SegmentReplicaCount>> replicaCountsInTier)
{
- this.replicaCountsInTier = ImmutableMap.copyOf(replicaCountsInTier);
-
- final Map<SegmentId, SegmentReplicaCount> totalReplicaCounts = new
HashMap<>();
- replicaCountsInTier.forEach((segmentId, tierToReplicaCount) -> {
+ // Make a deep copy of replicaCountsInTier as it may be mutated further
+ // in the same coordinator cycle (e.g. by BalanceSegments)
+ final Map<SegmentId, Map<String, SegmentReplicaCount>>
replicaCountsInTierCopy
+ = Maps.newHashMapWithExpectedSize(replicaCountsInTier.size());
+ final Map<SegmentId, SegmentReplicaCount> totalReplicaCounts
+ = Maps.newHashMapWithExpectedSize(replicaCountsInTier.size());
+ for (Map.Entry<SegmentId, Map<String, SegmentReplicaCount>> entry :
replicaCountsInTier.entrySet()) {
+ final Map<String, SegmentReplicaCount> tierCopy =
Maps.newHashMapWithExpectedSize(entry.getValue().size());
final SegmentReplicaCount total = new SegmentReplicaCount();
- tierToReplicaCount.values().forEach(total::accumulate);
- totalReplicaCounts.put(segmentId, total);
- });
- this.totalReplicaCounts = ImmutableMap.copyOf(totalReplicaCounts);
+ for (Map.Entry<String, SegmentReplicaCount> tierEntry :
entry.getValue().entrySet()) {
+ final SegmentReplicaCount countCopy = tierEntry.getValue().copy();
+ tierCopy.put(tierEntry.getKey(), countCopy);
+ total.accumulate(countCopy);
+ }
+ replicaCountsInTierCopy.put(entry.getKey(),
Collections.unmodifiableMap(tierCopy));
+ totalReplicaCounts.put(entry.getKey(), total);
+ }
+ this.replicaCountsInTier =
Collections.unmodifiableMap(replicaCountsInTierCopy);
+ this.totalReplicaCounts = totalReplicaCounts;
}
public SegmentReplicaCount getReplicaCountsInCluster(SegmentId segmentId)
@@ -56,6 +68,25 @@ public class SegmentReplicationStatus
return totalReplicaCounts.get(segmentId);
}
+ /**
+ * A segment is unavailable if it has no loaded replicas while still
requiring at least one.
+ */
+ public boolean isUnavailable(SegmentId segmentId)
+ {
+ final SegmentReplicaCount totalCount = totalReplicaCounts.get(segmentId);
+ return totalCount == null || (totalCount.totalLoaded() <= 0 &&
totalCount.required() != 0);
+ }
+
+ /**
+ * A segment is deep-storage only if it is not loaded anywhere and requires
zero replicas,
+ * so it is served exclusively from deep storage.
+ */
+ public boolean isDeepStorageOnly(SegmentId segmentId)
+ {
+ final SegmentReplicaCount totalCount = totalReplicaCounts.get(segmentId);
+ return totalCount != null && totalCount.totalLoaded() == 0 &&
totalCount.required() == 0;
+ }
+
public Map<String, Object2LongMap<String>>
getTierToDatasourceToUnderReplicated(
Iterable<DataSegment> usedSegments,
boolean ignoreMissingServers
@@ -81,4 +112,44 @@ public class SegmentReplicationStatus
return tierToUnderReplicated;
}
+
+ /**
+ * Computes unavailable, under-replicated and deep-storage-only segment
counts in a single
+ * pass over {@code usedSegments}, instead of three independent full
iterations. Produces
+ * results identical to calling {@link #getReplicaCountsInCluster} and
+ * {@link #getTierToDatasourceToUnderReplicated} independently for each
segment.
+ *
+ * @param ignoreMissingServers same semantics as in {@link
#getTierToDatasourceToUnderReplicated}.
+ */
+ public SegmentStatsSnapshot computeSegmentStats(Iterable<DataSegment>
usedSegments, boolean ignoreMissingServers)
+ {
+ final Object2IntOpenHashMap<String> datasourceToUnavailable = new
Object2IntOpenHashMap<>();
+ final Object2IntOpenHashMap<String> datasourceToDeepStorageOnly = new
Object2IntOpenHashMap<>();
+ final Map<String, Object2LongMap<String>> tierToUnderReplicated = new
HashMap<>();
+
+ for (DataSegment segment : usedSegments) {
+ final SegmentId segmentId = segment.getId();
+ final String datasource = segment.getDataSource();
+
+ // addTo with 0 ensures the datasource is present in the map even when
it has no unavailable segments.
+ datasourceToUnavailable.addTo(datasource, isUnavailable(segmentId) ? 1 :
0);
+ if (isDeepStorageOnly(segmentId)) {
+ datasourceToDeepStorageOnly.addTo(datasource, 1);
+ }
+
+ final Map<String, SegmentReplicaCount> tierToReplicaCount =
replicaCountsInTier.get(segmentId);
+ if (tierToReplicaCount != null) {
+ tierToReplicaCount.forEach((tier, counts) -> {
+ final int underReplicated = ignoreMissingServers ? counts.missing()
: counts.missingAndLoadable();
+ if (underReplicated >= 0) {
+ Object2LongOpenHashMap<String> datasourceToUnderReplicated =
(Object2LongOpenHashMap<String>)
+ tierToUnderReplicated.computeIfAbsent(tier, ds -> new
Object2LongOpenHashMap<>());
+ datasourceToUnderReplicated.addTo(datasource, underReplicated);
+ }
+ });
+ }
+ }
+
+ return new SegmentStatsSnapshot(datasourceToUnavailable,
tierToUnderReplicated, datasourceToDeepStorageOnly);
+ }
}
diff --git
a/server/src/main/java/org/apache/druid/server/coordinator/loading/SegmentStatsSnapshot.java
b/server/src/main/java/org/apache/druid/server/coordinator/loading/SegmentStatsSnapshot.java
new file mode 100644
index 00000000000..3a7fb163fb2
--- /dev/null
+++
b/server/src/main/java/org/apache/druid/server/coordinator/loading/SegmentStatsSnapshot.java
@@ -0,0 +1,78 @@
+/*
+ * 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.coordinator.loading;
+
+import it.unimi.dsi.fastutil.objects.Object2IntMap;
+import it.unimi.dsi.fastutil.objects.Object2IntMaps;
+import it.unimi.dsi.fastutil.objects.Object2LongMap;
+
+import java.util.Collections;
+import java.util.Map;
+
+/**
+ * Holder for the three segment-stat views computed together by
+ * {@link SegmentReplicationStatus#computeSegmentStats}.
+ */
+public class SegmentStatsSnapshot
+{
+ private static final SegmentStatsSnapshot EMPTY = new SegmentStatsSnapshot(
+ Object2IntMaps.emptyMap(),
+ Collections.emptyMap(),
+ Object2IntMaps.emptyMap()
+ );
+
+ private final Object2IntMap<String> datasourceToUnavailableCount;
+ private final Map<String, Object2LongMap<String>>
tierToDatasourceToUnderReplicatedCount;
+ private final Object2IntMap<String> datasourceToDeepStorageOnlyCount;
+
+ SegmentStatsSnapshot(
+ Object2IntMap<String> datasourceToUnavailableCount,
+ Map<String, Object2LongMap<String>>
tierToDatasourceToUnderReplicatedCount,
+ Object2IntMap<String> datasourceToDeepStorageOnlyCount
+ )
+ {
+ this.datasourceToUnavailableCount = datasourceToUnavailableCount;
+ this.tierToDatasourceToUnderReplicatedCount =
tierToDatasourceToUnderReplicatedCount;
+ this.datasourceToDeepStorageOnlyCount = datasourceToDeepStorageOnlyCount;
+ }
+
+ /**
+ * Returns an immutable, empty snapshot, used when no replication status is
available.
+ */
+ public static SegmentStatsSnapshot empty()
+ {
+ return EMPTY;
+ }
+
+ public Object2IntMap<String> getDatasourceToUnavailableCount()
+ {
+ return datasourceToUnavailableCount;
+ }
+
+ public Map<String, Object2LongMap<String>>
getTierToDatasourceToUnderReplicatedCount()
+ {
+ return tierToDatasourceToUnderReplicatedCount;
+ }
+
+ public Object2IntMap<String> getDatasourceToDeepStorageOnlyCount()
+ {
+ return datasourceToDeepStorageOnlyCount;
+ }
+}
diff --git
a/server/src/main/java/org/apache/druid/server/coordinator/loading/StrategicSegmentAssigner.java
b/server/src/main/java/org/apache/druid/server/coordinator/loading/StrategicSegmentAssigner.java
index fc6936e4df1..f65ed4c46f4 100644
---
a/server/src/main/java/org/apache/druid/server/coordinator/loading/StrategicSegmentAssigner.java
+++
b/server/src/main/java/org/apache/druid/server/coordinator/loading/StrategicSegmentAssigner.java
@@ -19,7 +19,6 @@
package org.apache.druid.server.coordinator.loading;
-import com.google.common.collect.Sets;
import it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap;
import org.apache.druid.client.DruidServer;
import org.apache.druid.server.coordinator.DruidCluster;
@@ -64,6 +63,8 @@ public class StrategicSegmentAssigner implements
SegmentActionHandler
private final RoundRobinServerSelector serverSelector;
private final BalancerStrategy strategy;
+ private final Set<String> allTiersInCluster;
+
private final boolean useRoundRobinAssignment;
private final Map<String, Set<String>> historicalTierAliases;
private final Map<String, String> tierToAliasName;
@@ -86,6 +87,7 @@ public class StrategicSegmentAssigner implements
SegmentActionHandler
this.cluster = cluster;
this.strategy = strategy;
this.loadQueueManager = loadQueueManager;
+ this.allTiersInCluster = Set.copyOf(cluster.getTierNames());
this.replicaCountMap = SegmentReplicaCountMap.create(cluster);
this.replicationThrottler = createReplicationThrottler(cluster,
loadingConfig);
this.useRoundRobinAssignment =
loadingConfig.isUseRoundRobinSegmentAssignment();
@@ -234,7 +236,6 @@ public class StrategicSegmentAssigner implements
SegmentActionHandler
public void replicateSegment(DataSegment segment, Map<String, Integer>
tierToReplicaCount)
{
final Map<String, Integer> effectiveTierToReplicaCount =
expandWithAliases(tierToReplicaCount);
- final Set<String> allTiersInCluster =
Sets.newHashSet(cluster.getTierNames());
if (effectiveTierToReplicaCount.isEmpty()) {
// Track the counts for a segment even if it requires 0 replicas on all
tiers
@@ -290,7 +291,6 @@ public class StrategicSegmentAssigner implements
SegmentActionHandler
)
{
final Map<String, Integer> effectiveTierToReplicaCount =
expandWithAliases(tierToReplicaCount);
- final Set<String> allTiersInCluster =
Sets.newHashSet(cluster.getTierNames());
if (effectiveTierToReplicaCount.isEmpty()) {
replicaCountMap.computeIfAbsent(segment.getId(),
DruidServer.DEFAULT_TIER);
diff --git
a/server/src/test/java/org/apache/druid/server/coordinator/loading/SegmentReplicationStatusTest.java
b/server/src/test/java/org/apache/druid/server/coordinator/loading/SegmentReplicationStatusTest.java
new file mode 100644
index 00000000000..5338b2f0ef9
--- /dev/null
+++
b/server/src/test/java/org/apache/druid/server/coordinator/loading/SegmentReplicationStatusTest.java
@@ -0,0 +1,187 @@
+/*
+ * 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.coordinator.loading;
+
+import it.unimi.dsi.fastutil.objects.Object2IntMap;
+import it.unimi.dsi.fastutil.objects.Object2LongMap;
+import org.apache.druid.java.util.common.granularity.Granularities;
+import org.apache.druid.server.coordinator.CreateDataSegments;
+import org.apache.druid.timeline.DataSegment;
+import org.apache.druid.timeline.SegmentId;
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Verifies that {@link SegmentReplicationStatus#computeSegmentStats} produces
results
+ * identical to computing unavailable, under-replicated and deep-storage-only
counts
+ * independently via {@link
SegmentReplicationStatus#getReplicaCountsInCluster} and
+ * {@link SegmentReplicationStatus#getTierToDatasourceToUnderReplicated}.
+ */
+public class SegmentReplicationStatusTest
+{
+ private static final String TIER_1 = "tier1";
+ private static final String TIER_2 = "tier2";
+
+ private final List<DataSegment> segments = CreateDataSegments
+ .ofDatasource("wiki")
+ .forIntervals(5, Granularities.DAY)
+ .eachOfSize(100);
+
+ @Test
+ public void testComputeSegmentStatsMatchesIndependentComputation()
+ {
+ final Map<SegmentId, Map<String, SegmentReplicaCount>> replicaCountsInTier
= new HashMap<>();
+
+ // Segment 0: fully loaded, single tier, no under-replication.
+ replicaCountsInTier.put(segments.get(0).getId(), Map.of(TIER_1, countOf(2,
2, 2)));
+
+ // Segment 1: unavailable (required > 0, nothing loaded).
+ replicaCountsInTier.put(segments.get(1).getId(), Map.of(TIER_1, countOf(2,
2, 0)));
+
+ // Segment 2: deep-storage-only (required == 0, nothing loaded).
+ replicaCountsInTier.put(segments.get(2).getId(), Map.of(TIER_1, countOf(0,
0, 0)));
+
+ // Segment 3: under-replicated in one tier, satisfied in another.
+ final Map<String, SegmentReplicaCount> tierMap3 = new HashMap<>();
+ tierMap3.put(TIER_1, countOf(2, 2, 1));
+ tierMap3.put(TIER_2, countOf(1, 1, 1));
+ replicaCountsInTier.put(segments.get(3).getId(), tierMap3);
+
+ // Segment 4: present in the used-segment list but absent from the
replication map
+ // (e.g. metadata race) -- must be treated as unavailable, not throw.
+ // Intentionally not added to replicaCountsInTier.
+
+ final SegmentReplicationStatus status = new
SegmentReplicationStatus(replicaCountsInTier);
+
+ for (boolean ignoreMissingServers : new boolean[]{true, false}) {
+ final SegmentStatsSnapshot snapshot =
+ status.computeSegmentStats(segments, ignoreMissingServers);
+
+ final Object2IntMap<String> expectedUnavailable =
computeExpectedUnavailable(status, segments);
+ final Object2IntMap<String> expectedDeepStorageOnly =
computeExpectedDeepStorageOnly(status, segments);
+ final Map<String, Object2LongMap<String>> expectedUnderReplicated =
+ status.getTierToDatasourceToUnderReplicated(segments,
ignoreMissingServers);
+
+ Assert.assertEquals(expectedUnavailable,
snapshot.getDatasourceToUnavailableCount());
+ Assert.assertEquals(expectedDeepStorageOnly,
snapshot.getDatasourceToDeepStorageOnlyCount());
+ Assert.assertEquals(expectedUnderReplicated,
snapshot.getTierToDatasourceToUnderReplicatedCount());
+ }
+ }
+
+ @Test
+ public void testIgnoreMissingServersUsesMissingNotMissingAndLoadable()
+ {
+ final Map<SegmentId, Map<String, SegmentReplicaCount>> replicaCountsInTier
= new HashMap<>();
+
+ // required=3, requiredAndLoadable=1 (only 1 loadable server), loaded=0.
+ // missing() = 3, missingAndLoadable() = 1.
+ final SegmentReplicaCount count = new SegmentReplicaCount();
+ count.setRequired(3, 1);
+ replicaCountsInTier.put(segments.get(0).getId(), Map.of(TIER_1, count));
+
+ final SegmentReplicationStatus status = new
SegmentReplicationStatus(replicaCountsInTier);
+ final List<DataSegment> singleSegment = List.of(segments.get(0));
+
+ final SegmentStatsSnapshot ignoreMissing =
+ status.computeSegmentStats(singleSegment, true);
+ Assert.assertEquals(
+ 3L,
+
ignoreMissing.getTierToDatasourceToUnderReplicatedCount().get(TIER_1).getLong("wiki")
+ );
+
+ final SegmentStatsSnapshot respectMissing =
+ status.computeSegmentStats(singleSegment, false);
+ Assert.assertEquals(
+ 1L,
+
respectMissing.getTierToDatasourceToUnderReplicatedCount().get(TIER_1).getLong("wiki")
+ );
+ }
+
+ @Test
+ public void testConstructorSnapshotsInputAndIsUnaffectedByLaterMutation()
+ {
+ final Map<SegmentId, Map<String, SegmentReplicaCount>> replicaCountsInTier
= new HashMap<>();
+ final SegmentReplicaCount count = countOf(2, 2, 1);
+ final Map<String, SegmentReplicaCount> tierMap = new HashMap<>();
+ tierMap.put(TIER_1, count);
+ replicaCountsInTier.put(segments.get(0).getId(), tierMap);
+
+ final SegmentReplicationStatus status = new
SegmentReplicationStatus(replicaCountsInTier);
+ final SegmentReplicaCount totalBefore =
status.getReplicaCountsInCluster(segments.get(0).getId());
+ Assert.assertEquals(1, totalBefore.totalLoaded());
+
+ // Mutate the caller's inputs after construction, as BalanceSegments does
later in the same
+ // coordinator cycle: add a brand new segment key and mutate an
already-tracked count in place.
+ count.incrementLoaded();
+ tierMap.put(TIER_2, countOf(1, 1, 1));
+ replicaCountsInTier.put(segments.get(1).getId(), Map.of(TIER_1, countOf(1,
1, 1)));
+
+ Assert.assertEquals(1,
status.getReplicaCountsInCluster(segments.get(0).getId()).totalLoaded());
+
Assert.assertNull(status.getReplicaCountsInCluster(segments.get(1).getId()));
+ }
+
+ private static SegmentReplicaCount countOf(int required, int
requiredAndLoadable, int loaded)
+ {
+ final SegmentReplicaCount count = new SegmentReplicaCount();
+ count.setRequired(required, requiredAndLoadable);
+ for (int i = 0; i < loaded; i++) {
+ count.incrementLoaded();
+ }
+ return count;
+ }
+
+ private static Object2IntMap<String> computeExpectedUnavailable(
+ SegmentReplicationStatus status,
+ List<DataSegment> segments
+ )
+ {
+ final it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap<String> result =
+ new it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap<>();
+ for (DataSegment segment : segments) {
+ final SegmentReplicaCount rc =
status.getReplicaCountsInCluster(segment.getId());
+ if (rc != null && (rc.totalLoaded() > 0 || rc.required() == 0)) {
+ result.addTo(segment.getDataSource(), 0);
+ } else {
+ result.addTo(segment.getDataSource(), 1);
+ }
+ }
+ return result;
+ }
+
+ private static Object2IntMap<String> computeExpectedDeepStorageOnly(
+ SegmentReplicationStatus status,
+ List<DataSegment> segments
+ )
+ {
+ final it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap<String> result =
+ new it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap<>();
+ for (DataSegment segment : segments) {
+ final SegmentReplicaCount rc =
status.getReplicaCountsInCluster(segment.getId());
+ if (rc != null && rc.totalLoaded() == 0 && rc.required() == 0) {
+ result.addTo(segment.getDataSource(), 1);
+ }
+ }
+ return result;
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
