Allow coordinator to move segments cognizant of deploymentGroup

Author: jtuglu1

docs/api-reference/dynamic-configuration-api.md

@@ -110,7 +110,8 @@ Host: http://ROUTER_IP:ROUTER_PORT
     "cloneServers": {},
     "historicalTierAliases": {
         "hot": ["hot_1", "hot_2"]
-    }
+    },
+    "coordinatingVersions": []
 
 }
 ```

docs/configuration/index.md

@@ -813,6 +813,7 @@ The following table shows the dynamic configuration properties for the Coordinat
 |`turboLoadingNodes`| Experimental. List of Historical servers to place in turbo loading mode. These servers use a larger thread-pool to load segments faster but at the cost of query performance. For servers specified in `turboLoadingNodes`, `druid.coordinator.loadqueuepeon.http.batchSize` is ignored and the coordinator uses the value of the respective `numLoadingThreads` instead.<br/>Please use this config with caution. All servers should eventually be removed from this list once the segment loading on the respective historicals is finished. |none|
 |`cloneServers`| Experimental. Map from target Historical server to source Historical server which should be cloned by the target. The target Historical does not participate in regular segment assignment or balancing. Instead, the Coordinator mirrors any segment assignment made to the source Historical onto the target Historical, so that the target becomes an exact copy of the source. Segments on the target Historical do not count towards replica counts either. If the source disappears, the target remains in the last known state of the source server until removed from the configuration. <br/>Use this config with caution. All servers should eventually be removed from this list once the desired state on the respective Historicals is achieved. |none|
 |`historicalTierAliases`|Map from a virtual tier name to the set of real Historical tier names it expands to. When a load/drop rule references a virtual alias tier, the Coordinator replaces it with its real tiers — each receiving the full replica count independently. The alias key itself is never loaded to directly. For example, `{"hot": ["hot_1", "hot_2"]}` causes a rule of `{"hot": 2}` to load 2 replicas on each of `hot_1` and `hot_2`; `hot` receives no direct assignment. An alias value tier with no servers raises the normal invalid-tier alert. If a rule already specifies an explicit replica count for a tier that also appears as an alias value, the explicit count takes precedence. Duplicate tier names within a set are ignored. A virtual alias tier cannot also be a physical tier.|none|
+|`coordinatingVersions`|List of deployment groups for which the Coordinator enforces per-group Historical replication and handoff. When set, load/drop rules are applied independently to each listed group that has active servers in a tier, so each group receives the rule's required replicas. Servers outside the listed groups are not assigned new replicas and may have surplus replicas dropped. Empty disables per-group coordination.|none|
 
 ##### Smart segment loading
 

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.LinkedHashMap;
+import java.util.LinkedHashSet;
+import java.util.Map;
 import java.util.Set;
 
 /**
@@ -60,9 +63,27 @@ 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);
+        int remainingGroups = serversByGroup.size();
+        int remainingSegmentsToMove = maxSegmentsToMove;
+        for (final Set<ServerHolder> groupServers : serversByGroup.values()) {
+          final int groupMaxSegmentsToMove =
+              (remainingSegmentsToMove + remainingGroups - 1) / remainingGroups;
+          if (groupMaxSegmentsToMove > 0) {
+            new TierSegmentBalancer(tier, groupServers, groupMaxSegmentsToMove, params).run();
+            remainingSegmentsToMove -= groupMaxSegmentsToMove;
+          }
+          --remainingGroups;
+        }
+      }
+    });
 
     CoordinatorRunStats runStats = params.getCoordinatorStats();
     params.getBalancerStrategy()
@@ -123,4 +144,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 LinkedHashMap<>();
+    for (final ServerHolder server : servers) {
+      final String group = server.getServer().getMetadata().getDeploymentGroup();
+      byGroup.computeIfAbsent(group == null ? "" : group, g -> new LinkedHashSet<>()).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);
+  }
+}