Allow coordinator to move segments cognizant of deploymentGroup

Author: jtuglu1

server/src/main/java/org/apache/druid/server/coordinator/CoordinatorDynamicConfig.java

@@ -84,6 +84,15 @@ public class CoordinatorDynamicConfig
    */
   private final Map<String, Set<String>> historicalTierAliases;
 
+  /**
+   * Deployment groups the coordinator enforces per-group replication and handoff for.
+   * When non-empty, the coordinator ensures {@code requiredReplicas} per listed group rather than
+   * per tier, and the handoff endpoint requires at least one server in each listed group (that has
+   * online servers) to serve the segment before declaring handoff complete.
+   * An empty set disables this behavior and restores default tier-wide replication.
+   */
+  private final Set<String> coordinatingVersions;
+
   /**
    * Stale pending segments belonging to the data sources in this list are not killed by {@code
    * KillStalePendingSegments}. In other words, segments in these data sources are "protected".
@@ -135,7 +144,8 @@ public CoordinatorDynamicConfig(
       @JsonProperty("debugDimensions") @Nullable Map<String, String> debugDimensions,
       @JsonProperty("turboLoadingNodes") @Nullable Set<String> turboLoadingNodes,
       @JsonProperty("cloneServers") @Nullable Map<String, String> cloneServers,
-      @JsonProperty("historicalTierAliases") @Nullable Map<String, Set<String>> historicalTierAliases
+      @JsonProperty("historicalTierAliases") @Nullable Map<String, Set<String>> historicalTierAliases,
+      @JsonProperty("coordinatingVersions") @Nullable Set<String> coordinatingVersions
   )
   {
     this.markSegmentAsUnusedDelayMillis =
@@ -183,6 +193,7 @@ public CoordinatorDynamicConfig(
     this.cloneServers = Configs.valueOrDefault(cloneServers, Map.of());
 
     this.historicalTierAliases = Configs.valueOrDefault(historicalTierAliases, Map.of());
+    this.coordinatingVersions = Configs.valueOrDefault(coordinatingVersions, Set.of());
     final Set<String> aliasKeys = this.historicalTierAliases.keySet();
     for (Set<String> mappedTiers : this.historicalTierAliases.values()) {
       if (!Sets.intersection(mappedTiers, aliasKeys).isEmpty()) {
@@ -364,6 +375,12 @@ public Map<String, Set<String>> getHistoricalTierAliases()
     return historicalTierAliases;
   }
 
+  @JsonProperty
+  public Set<String> getCoordinatingVersions()
+  {
+    return coordinatingVersions;
+  }
+
   /**
    * List of servers to put in turbo-loading mode. These servers will use a larger thread pool to load
    * segments. This causes decreases the average time taken to load segments. However, this also means less resources
@@ -398,6 +415,7 @@ public String toString()
            ", turboLoadingNodes=" + turboLoadingNodes +
            ", cloneServers=" + cloneServers +
            ", historicalTierAliases=" + historicalTierAliases +
+           ", coordinatingVersions=" + coordinatingVersions +
            '}';
   }
 
@@ -435,7 +453,8 @@ public boolean equals(Object o)
            && Objects.equals(turboLoadingNodes, that.turboLoadingNodes)
            && Objects.equals(debugDimensions, that.debugDimensions)
            && Objects.equals(cloneServers, that.cloneServers)
-           && Objects.equals(historicalTierAliases, that.historicalTierAliases);
+           && Objects.equals(historicalTierAliases, that.historicalTierAliases)
+           && Objects.equals(coordinatingVersions, that.coordinatingVersions);
   }
 
   @Override
@@ -460,7 +479,8 @@ public int hashCode()
         debugDimensions,
         turboLoadingNodes,
         cloneServers,
-        historicalTierAliases
+        historicalTierAliases,
+        coordinatingVersions
     );
   }
 
@@ -518,6 +538,7 @@ public static class Builder
     private Set<String> turboLoadingNodes;
     private Map<String, String> cloneServers;
     private Map<String, Set<String>> historicalTierAliases;
+    private Set<String> coordinatingVersions;
 
     public Builder()
     {
@@ -543,7 +564,8 @@ public Builder(
         @JsonProperty("debugDimensions") @Nullable Map<String, String> debugDimensions,
         @JsonProperty("turboLoadingNodes") @Nullable Set<String> turboLoadingNodes,
         @JsonProperty("cloneServers") @Nullable Map<String, String> cloneServers,
-        @JsonProperty("historicalTierAliases") @Nullable Map<String, Set<String>> historicalTierAliases
+        @JsonProperty("historicalTierAliases") @Nullable Map<String, Set<String>> historicalTierAliases,
+        @JsonProperty("coordinatingVersions") @Nullable Set<String> coordinatingVersions
     )
     {
       this.markSegmentAsUnusedDelayMillis = markSegmentAsUnusedDelayMillis;
@@ -565,6 +587,7 @@ public Builder(
       this.turboLoadingNodes = turboLoadingNodes;
       this.cloneServers = cloneServers;
       this.historicalTierAliases = historicalTierAliases;
+      this.coordinatingVersions = coordinatingVersions;
     }
 
     public Builder withMarkSegmentAsUnusedDelayMillis(long leadingTimeMillis)
@@ -669,6 +692,12 @@ public Builder withHistoricalTierAliases(Map<String, Set<String>> historicalTier
       return this;
     }
 
+    public Builder withCoordinatingVersions(Set<String> coordinatingVersions)
+    {
+      this.coordinatingVersions = coordinatingVersions;
+      return this;
+    }
+
     /**
      * Builds a CoordinatoryDynamicConfig using either the configured values, or
      * the default value if not configured.
@@ -697,7 +726,8 @@ public CoordinatorDynamicConfig build()
           debugDimensions,
           turboLoadingNodes,
           cloneServers,
-          historicalTierAliases
+          historicalTierAliases,
+          coordinatingVersions
       );
     }
 
@@ -730,7 +760,8 @@ public CoordinatorDynamicConfig build(CoordinatorDynamicConfig defaults)
           valueOrDefault(debugDimensions, defaults.getDebugDimensions()),
           valueOrDefault(turboLoadingNodes, defaults.getTurboLoadingNodes()),
           valueOrDefault(cloneServers, defaults.getCloneServers()),
-          valueOrDefault(historicalTierAliases, defaults.getHistoricalTierAliases())
+          valueOrDefault(historicalTierAliases, defaults.getHistoricalTierAliases()),
+          valueOrDefault(coordinatingVersions, defaults.getCoordinatingVersions())
       );
     }
   }

server/src/main/java/org/apache/druid/server/coordinator/DruidCluster.java

@@ -33,6 +33,7 @@
 import java.util.Map;
 import java.util.NavigableSet;
 import java.util.Set;
+import java.util.TreeSet;
 import java.util.stream.Collectors;
 
 /**
@@ -46,6 +47,11 @@ public class DruidCluster
   private final Set<ServerHolder> realtimes;
   private final Map<String, NavigableSet<ServerHolder>> historicals;
   private final Map<String, NavigableSet<ServerHolder>> managedHistoricals;
+  /**
+   * Managed historicals indexed by tier and then by deployment group. Only servers with a non-null
+   * deploymentGroup appear here. Used to support per-group replication and handoff enforcement.
+   */
+  private final Map<String, Map<String, NavigableSet<ServerHolder>>> managedHistoricalsByTierAndGroup;
   private final Set<ServerHolder> brokers;
   private final List<ServerHolder> allManagedServers;
 
@@ -70,6 +76,7 @@ private DruidCluster(
           return CollectionUtils.newTreeSet(Comparator.naturalOrder(), managedServers);
         }
     );
+    this.managedHistoricalsByTierAndGroup = initManagedHistoricalsByTierAndGroup();
     this.brokers = Collections.unmodifiableSet(brokers);
     this.allManagedServers = initAllManagedServers();
   }
@@ -111,6 +118,73 @@ public NavigableSet<ServerHolder> getManagedHistoricalsByTier(String tier)
     return managedHistoricals.get(tier);
   }
 
+  /**
+   * Returns the distinct non-null deployment groups present among managed historicals in the given tier.
+   */
+  public Set<String> getDeploymentGroupsForTier(String tier)
+  {
+    final Map<String, NavigableSet<ServerHolder>> groupMap = managedHistoricalsByTierAndGroup.get(tier);
+    return groupMap == null ? Collections.emptySet() : groupMap.keySet();
+  }
+
+  /**
+   * Returns managed historicals in the given tier that belong to the given deployment group.
+   * Returns an empty set if no servers for that (tier, group) pair exist.
+   */
+  public NavigableSet<ServerHolder> getManagedHistoricalsByTierAndGroup(String tier, String group)
+  {
+    final Map<String, NavigableSet<ServerHolder>> groupMap = managedHistoricalsByTierAndGroup.get(tier);
+    if (groupMap == null) {
+      return Collections.emptyNavigableSet();
+    }
+    final NavigableSet<ServerHolder> servers = groupMap.get(group);
+    return servers == null ? Collections.emptyNavigableSet() : servers;
+  }
+
+  /**
+   * Returns managed historicals in the given tier whose {@code deploymentGroup} is null or is not
+   * present in {@code coordinatingVersions}. These are the "uncoordinated" servers whose replica
+   * counts roll up to the tier-wide {@link ReplicaCountKey}; they must still be visited by drop /
+   * cancellation passes when the rest of the tier is operating in per-group mode.
+   */
+  public NavigableSet<ServerHolder> getUncoordinatedManagedHistoricalsByTier(
+      String tier,
+      Set<String> coordinatingVersions
+  )
+  {
+    final NavigableSet<ServerHolder> all = managedHistoricals.get(tier);
+    if (all == null || all.isEmpty()) {
+      return Collections.emptyNavigableSet();
+    }
+    if (coordinatingVersions == null || coordinatingVersions.isEmpty()) {
+      return all;
+    }
+    final NavigableSet<ServerHolder> filtered = new TreeSet<>(Comparator.naturalOrder());
+    for (ServerHolder server : all) {
+      final String group = server.getServer().getMetadata().getDeploymentGroup();
+      if (group == null || !coordinatingVersions.contains(group)) {
+        filtered.add(server);
+      }
+    }
+    return filtered;
+  }
+
+  private Map<String, Map<String, NavigableSet<ServerHolder>>> initManagedHistoricalsByTierAndGroup()
+  {
+    final Map<String, Map<String, NavigableSet<ServerHolder>>> result = new HashMap<>();
+    managedHistoricals.forEach((tier, servers) -> {
+      for (ServerHolder server : servers) {
+        final String group = server.getServer().getMetadata().getDeploymentGroup();
+        if (group != null) {
+          result.computeIfAbsent(tier, t -> new HashMap<>())
+                .computeIfAbsent(group, g -> new TreeSet<>(Comparator.naturalOrder()))
+                .add(server);
+        }
+      }
+    });
+    return Collections.unmodifiableMap(result);
+  }
+
   public List<ServerHolder> getAllManagedServers()
   {
     return allManagedServers;

server/src/main/java/org/apache/druid/server/coordinator/duty/BalanceSegments.java

@@ -31,6 +31,9 @@
 import org.apache.druid.server.coordinator.stats.Stats;
 import org.joda.time.Duration;
 
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
 import java.util.Set;
 
 /**
@@ -60,9 +63,20 @@ public DruidCoordinatorRuntimeParams run(DruidCoordinatorRuntimeParams params)
       return params;
     }
 
-    params.getDruidCluster().getManagedHistoricals().forEach(
-        (tier, servers) -> new TierSegmentBalancer(tier, servers, maxSegmentsToMove, params).run()
-    );
+    final Set<String> coordinatingVersions = params.getCoordinatorDynamicConfig().getCoordinatingVersions();
+    params.getDruidCluster().getManagedHistoricals().forEach((tier, servers) -> {
+      if (coordinatingVersions.isEmpty()) {
+        new TierSegmentBalancer(tier, servers, maxSegmentsToMove, params).run();
+      } else {
+        // Partition tier servers by deployment group so segments never move across groups.
+        // Servers with no deploymentGroup form their own partition keyed under the empty string.
+        final Map<String, Set<ServerHolder>> serversByGroup = partitionByDeploymentGroup(servers);
+        final int perGroupMax = Math.max(1, maxSegmentsToMove / serversByGroup.size());
+        serversByGroup.forEach(
+            (group, groupServers) -> new TierSegmentBalancer(tier, groupServers, perGroupMax, params).run()
+        );
+      }
+    });
 
     CoordinatorRunStats runStats = params.getCoordinatorStats();
     params.getBalancerStrategy()
@@ -123,4 +137,13 @@ private Pair<Integer, Integer> getNumHistoricalsAndSegments(DruidCluster cluster
     return Pair.of(numHistoricals, numSegments);
   }
 
+  private static Map<String, Set<ServerHolder>> partitionByDeploymentGroup(Set<ServerHolder> servers)
+  {
+    final Map<String, Set<ServerHolder>> byGroup = new HashMap<>();
+    for (ServerHolder server : servers) {
+      final String group = server.getServer().getMetadata().getDeploymentGroup();
+      byGroup.computeIfAbsent(group == null ? "" : group, g -> new HashSet<>()).add(server);
+    }
+    return byGroup;
+  }
 }

server/src/main/java/org/apache/druid/server/coordinator/loading/ReplicaCountKey.java

@@ -0,0 +1,48 @@
+/*
+ * 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 javax.annotation.Nullable;
+import java.util.Set;
+
+/**
+ * Map key used by {@link SegmentReplicaCountMap}. A null {@code group} represents tier-wide
+ * tracking; a non-null {@code group} represents a specific deployment group within the tier
+ * (used when the coordinator is enforcing per-group replication via {@code coordinatingVersions}).
+ */
+public record ReplicaCountKey(String tier, @Nullable String group)
+{
+  public static ReplicaCountKey forTier(String tier)
+  {
+    return new ReplicaCountKey(tier, null);
+  }
+
+  /**
+   * Returns a (tier, group) key when {@code group} is non-null and present in
+   * {@code coordinatingVersions}; otherwise a plain tier-wide key.
+   */
+  public static ReplicaCountKey from(String tier, @Nullable String group, Set<String> coordinatingVersions)
+  {
+    if (group != null && coordinatingVersions.contains(group)) {
+      return new ReplicaCountKey(tier, group);
+    }
+    return new ReplicaCountKey(tier, null);
+  }
+}

server/src/main/java/org/apache/druid/server/coordinator/loading/RoundRobinServerSelector.java

@@ -48,12 +48,25 @@
 public class RoundRobinServerSelector
 {
   private final Map<String, CircularServerList> tierToServers = new HashMap<>();
+  private final Map<String, Map<String, CircularServerList>> tierGroupToServers = new HashMap<>();
 
-  public RoundRobinServerSelector(DruidCluster cluster)
+  public RoundRobinServerSelector(DruidCluster cluster, Set<String> coordinatingVersions)
   {
     cluster.getManagedHistoricals().forEach(
         (tier, servers) -> tierToServers.put(tier, new CircularServerList(servers))
     );
+
+    if (!coordinatingVersions.isEmpty()) {
+      cluster.getManagedHistoricals().keySet().forEach(tier -> {
+        for (String group : coordinatingVersions) {
+          final var groupServers = cluster.getManagedHistoricalsByTierAndGroup(tier, group);
+          if (!groupServers.isEmpty()) {
+            tierGroupToServers.computeIfAbsent(tier, t -> new HashMap<>())
+                              .put(group, new CircularServerList(groupServers));
+          }
+        }
+      });
+    }
   }
 
   /**
@@ -70,6 +83,20 @@ public Iterator<ServerHolder> getServersInTierToLoadSegment(String tier, DataSeg
     return new EligibleServerIterator(segment, iterator);
   }
 
+  /**
+   * Returns an iterator over servers in the given tier and deployment group that are eligible to
+   * load the given segment.
+   */
+  public Iterator<ServerHolder> getServersInTierAndGroupToLoadSegment(String tier, String group, DataSegment segment)
+  {
+    final Map<String, CircularServerList> groupMap = tierGroupToServers.get(tier);
+    if (groupMap == null) {
+      return Collections.emptyIterator();
+    }
+    final CircularServerList list = groupMap.get(group);
+    return list == null ? Collections.emptyIterator() : new EligibleServerIterator(segment, list);
+  }
+
   /**
    * Iterator over servers in a tier that are eligible to load a given segment.
    */