Allow for setting the total shards per node in the Allocate ILM action (#76134)

This adds a new optional field to the allocate ILM action called "total_shards_per_node". If present, the
value of this field is set as the value of "index.routing.allocation.total_shards_per_node" before the allocation
takes place.
Relates to #44070

docs/reference/ilm/actions/ilm-allocate.asciidoc

@@ -32,6 +32,11 @@ see <<shard-allocation-filtering>>.
 (Optional, integer)
 Number of replicas to assign to the index.
 
+`total_shards_per_node`::
+(Optional, integer)
+The maximum number of shards for the index on a single {es} node. A value of `-1` is
+interpreted as unlimited. See <<allocation-total-shards, total shards>>.
+
 `include`::
 (Optional, object)
 Assigns an index to nodes that have at least _one_ of the specified custom attributes.
@@ -48,7 +53,8 @@ Assigns an index to nodes that have _all_ of the specified custom attributes.
 ==== Example
 
 The allocate action in the following policy changes the index's number of replicas to `2`.
-The index allocation rules are not changed.
+No more than 200 shards for the index will be placed on any single node. Otherwise the index
+allocation rules are not changed.
 
 [source,console]
 --------------------------------------------------
@@ -59,7 +65,8 @@ PUT _ilm/policy/my_policy
       "warm": {
         "actions": {
           "allocate" : {
-            "number_of_replicas" : 2
+            "number_of_replicas" : 2,
+            "total_shards_per_node" : 200
           }
         }
       }

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ilm/AllocateAction.java

@@ -6,8 +6,10 @@
  */
 package org.elasticsearch.xpack.core.ilm;
 
+import org.elasticsearch.Version;
 import org.elasticsearch.client.Client;
 import org.elasticsearch.cluster.metadata.IndexMetadata;
+import org.elasticsearch.cluster.routing.allocation.decider.ShardsLimitAllocationDecider;
 import org.elasticsearch.common.xcontent.ParseField;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.io.stream.StreamInput;
@@ -29,22 +31,26 @@ public class AllocateAction implements LifecycleAction {
 
     public static final String NAME = "allocate";
     public static final ParseField NUMBER_OF_REPLICAS_FIELD = new ParseField("number_of_replicas");
+    public static final ParseField TOTAL_SHARDS_PER_NODE_FIELD = new ParseField("total_shards_per_node");
     public static final ParseField INCLUDE_FIELD = new ParseField("include");
     public static final ParseField EXCLUDE_FIELD = new ParseField("exclude");
     public static final ParseField REQUIRE_FIELD = new ParseField("require");
 
     @SuppressWarnings("unchecked")
     private static final ConstructingObjectParser<AllocateAction, Void> PARSER = new ConstructingObjectParser<>(NAME,
-            a -> new AllocateAction((Integer) a[0], (Map<String, String>) a[1], (Map<String, String>) a[2], (Map<String, String>) a[3]));
+            a -> new AllocateAction((Integer) a[0], (Integer) a[1], (Map<String, String>) a[2], (Map<String, String>) a[3],
+                (Map<String, String>) a[4]));
 
     static {
         PARSER.declareInt(ConstructingObjectParser.optionalConstructorArg(), NUMBER_OF_REPLICAS_FIELD);
+        PARSER.declareInt(ConstructingObjectParser.optionalConstructorArg(), TOTAL_SHARDS_PER_NODE_FIELD);
         PARSER.declareObject(ConstructingObjectParser.optionalConstructorArg(), (p, c) -> p.mapStrings(), INCLUDE_FIELD);
         PARSER.declareObject(ConstructingObjectParser.optionalConstructorArg(), (p, c) -> p.mapStrings(), EXCLUDE_FIELD);
         PARSER.declareObject(ConstructingObjectParser.optionalConstructorArg(), (p, c) -> p.mapStrings(), REQUIRE_FIELD);
     }
 
     private final Integer numberOfReplicas;
+    private final Integer totalShardsPerNode;
     private final Map<String, String> include;
     private final Map<String, String> exclude;
     private final Map<String, String> require;
@@ -53,7 +59,8 @@ public class AllocateAction implements LifecycleAction {
         return PARSER.apply(parser, null);
     }
 
-    public AllocateAction(Integer numberOfReplicas, Map<String, String> include, Map<String, String> exclude, Map<String, String> require) {
+    public AllocateAction(Integer numberOfReplicas, Integer totalShardsPerNode, Map<String, String> include, Map<String, String> exclude,
+                          Map<String, String> require) {
         if (include == null) {
             this.include = Collections.emptyMap();
         } else {
@@ -78,18 +85,27 @@ public class AllocateAction implements LifecycleAction {
             throw new IllegalArgumentException("[" + NUMBER_OF_REPLICAS_FIELD.getPreferredName() + "] must be >= 0");
         }
         this.numberOfReplicas = numberOfReplicas;
+        if (totalShardsPerNode != null && totalShardsPerNode < -1) {
+            throw new IllegalArgumentException("[" + TOTAL_SHARDS_PER_NODE_FIELD.getPreferredName() + "] must be >= -1");
+        }
+        this.totalShardsPerNode = totalShardsPerNode;
     }
 
     @SuppressWarnings("unchecked")
     public AllocateAction(StreamInput in) throws IOException {
-        this(in.readOptionalVInt(), (Map<String, String>) in.readGenericValue(), (Map<String, String>) in.readGenericValue(),
-                (Map<String, String>) in.readGenericValue());
+        this(in.readOptionalVInt(), in.getVersion().onOrAfter(Version.V_8_0_0) ? in.readOptionalInt() : null,
+            (Map<String, String>) in.readGenericValue(), (Map<String, String>) in.readGenericValue(),
+            (Map<String, String>) in.readGenericValue());
     }
 
     public Integer getNumberOfReplicas() {
         return numberOfReplicas;
     }
 
+    public Integer getTotalShardsPerNode() {
+        return totalShardsPerNode;
+    }
+
     public Map<String, String> getInclude() {
         return include;
     }
@@ -105,6 +121,9 @@ public class AllocateAction implements LifecycleAction {
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         out.writeOptionalVInt(numberOfReplicas);
+        if (out.getVersion().onOrAfter(Version.V_8_0_0)) {
+            out.writeOptionalInt(totalShardsPerNode);
+        }
         out.writeGenericValue(include);
         out.writeGenericValue(exclude);
         out.writeGenericValue(require);
@@ -121,6 +140,9 @@ public class AllocateAction implements LifecycleAction {
         if (numberOfReplicas != null) {
             builder.field(NUMBER_OF_REPLICAS_FIELD.getPreferredName(), numberOfReplicas);
         }
+        if (totalShardsPerNode != null) {
+            builder.field(TOTAL_SHARDS_PER_NODE_FIELD.getPreferredName(), totalShardsPerNode);
+        }
         builder.field(INCLUDE_FIELD.getPreferredName(), include);
         builder.field(EXCLUDE_FIELD.getPreferredName(), exclude);
         builder.field(REQUIRE_FIELD.getPreferredName(), require);
@@ -145,6 +167,9 @@ public class AllocateAction implements LifecycleAction {
         include.forEach((key, value) -> newSettings.put(IndexMetadata.INDEX_ROUTING_INCLUDE_GROUP_SETTING.getKey() + key, value));
         exclude.forEach((key, value) -> newSettings.put(IndexMetadata.INDEX_ROUTING_EXCLUDE_GROUP_SETTING.getKey() + key, value));
         require.forEach((key, value) -> newSettings.put(IndexMetadata.INDEX_ROUTING_REQUIRE_GROUP_SETTING.getKey() + key, value));
+        if (totalShardsPerNode != null) {
+            newSettings.put(ShardsLimitAllocationDecider.INDEX_TOTAL_SHARDS_PER_NODE_SETTING.getKey(), totalShardsPerNode);
+        }
         UpdateSettingsStep allocateStep = new UpdateSettingsStep(allocateKey, allocationRoutedKey, client, newSettings.build());
         AllocationRoutedStep routedCheckStep = new AllocationRoutedStep(allocationRoutedKey, nextStepKey);
         return Arrays.asList(allocateStep, routedCheckStep);
@@ -152,7 +177,7 @@ public class AllocateAction implements LifecycleAction {
 
     @Override
     public int hashCode() {
-        return Objects.hash(numberOfReplicas, include, exclude, require);
+        return Objects.hash(numberOfReplicas, totalShardsPerNode, include, exclude, require);
     }
 
     @Override
@@ -165,6 +190,7 @@ public class AllocateAction implements LifecycleAction {
         }
         AllocateAction other = (AllocateAction) obj;
         return Objects.equals(numberOfReplicas, other.numberOfReplicas) &&
+            Objects.equals(totalShardsPerNode, other.totalShardsPerNode) &&
             Objects.equals(include, other.include) &&
             Objects.equals(exclude, other.exclude) &&
             Objects.equals(require, other.require);

x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ilm/AllocateActionTests.java

@@ -8,6 +8,7 @@ package org.elasticsearch.xpack.core.ilm;
 
 import org.elasticsearch.cluster.metadata.IndexMetadata;
 import org.elasticsearch.cluster.node.DiscoveryNodeRole;
+import org.elasticsearch.cluster.routing.allocation.decider.ShardsLimitAllocationDecider;
 import org.elasticsearch.common.io.stream.Writeable.Reader;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.xcontent.XContentParser;
@@ -18,6 +19,7 @@ import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
+import static org.elasticsearch.cluster.routing.allocation.decider.ShardsLimitAllocationDecider.INDEX_TOTAL_SHARDS_PER_NODE_SETTING;
 import static org.hamcrest.Matchers.equalTo;
 
 public class AllocateActionTests extends AbstractActionTestCase<AllocateAction> {
@@ -55,7 +57,8 @@ public class AllocateActionTests extends AbstractActionTestCase<AllocateAction>
             requires = randomBoolean() ? null : Collections.emptyMap();
         }
         Integer numberOfReplicas = randomBoolean() ? null : randomIntBetween(0, 10);
-        return new AllocateAction(numberOfReplicas, includes, excludes, requires);
+        Integer totalShardsPerNode = randomBoolean() ? null : randomIntBetween(-1, 300);
+        return new AllocateAction(numberOfReplicas, totalShardsPerNode, includes, excludes, requires);
     }
 
 
@@ -70,6 +73,7 @@ public class AllocateActionTests extends AbstractActionTestCase<AllocateAction>
         Map<String, String> exclude = instance.getExclude();
         Map<String, String> require = instance.getRequire();
         Integer numberOfReplicas = instance.getNumberOfReplicas();
+        Integer totalShardsPerNode = instance.getTotalShardsPerNode();
         switch (randomIntBetween(0, 3)) {
         case 0:
             include = new HashMap<>(include);
@@ -89,7 +93,7 @@ public class AllocateActionTests extends AbstractActionTestCase<AllocateAction>
         default:
             throw new AssertionError("Illegal randomisation branch");
         }
-        return new AllocateAction(numberOfReplicas, include, exclude, require);
+        return new AllocateAction(numberOfReplicas, totalShardsPerNode, include, exclude, require);
     }
 
     public void testAllMapsNullOrEmpty() {
@@ -97,7 +101,7 @@ public class AllocateActionTests extends AbstractActionTestCase<AllocateAction>
         Map<String, String> exclude = randomBoolean() ? null : Collections.emptyMap();
         Map<String, String> require = randomBoolean() ? null : Collections.emptyMap();
         IllegalArgumentException exception = expectThrows(IllegalArgumentException.class,
-                () -> new AllocateAction(null, include, exclude, require));
+                () -> new AllocateAction(null, null, include, exclude, require));
         assertEquals("At least one of " + AllocateAction.INCLUDE_FIELD.getPreferredName() + ", "
                 + AllocateAction.EXCLUDE_FIELD.getPreferredName() + " or " + AllocateAction.REQUIRE_FIELD.getPreferredName()
                 + "must contain attributes for action " + AllocateAction.NAME, exception.getMessage());
@@ -108,10 +112,19 @@ public class AllocateActionTests extends AbstractActionTestCase<AllocateAction>
         Map<String, String> exclude = randomBoolean() ? null : Collections.emptyMap();
         Map<String, String> require = randomBoolean() ? null : Collections.emptyMap();
         IllegalArgumentException exception = expectThrows(IllegalArgumentException.class,
-            () -> new AllocateAction(randomIntBetween(-1000, -1), include, exclude, require));
+            () -> new AllocateAction(randomIntBetween(-1000, -1), randomIntBetween(0, 300), include, exclude, require));
         assertEquals("[" + AllocateAction.NUMBER_OF_REPLICAS_FIELD.getPreferredName() + "] must be >= 0", exception.getMessage());
     }
 
+    public void testInvalidTotalShardsPerNode() {
+        Map<String, String> include = randomAllocationRoutingMap(1, 5);
+        Map<String, String> exclude = randomBoolean() ? null : Collections.emptyMap();
+        Map<String, String> require = randomBoolean() ? null : Collections.emptyMap();
+        IllegalArgumentException exception = expectThrows(IllegalArgumentException.class,
+            () -> new AllocateAction(randomIntBetween(0, 300), randomIntBetween(-1000, -2), include, exclude, require));
+        assertEquals("[" + AllocateAction.TOTAL_SHARDS_PER_NODE_FIELD.getPreferredName() + "] must be >= -1", exception.getMessage());
+    }
+
     public static Map<String, String> randomAllocationRoutingMap(int minEntries, int maxEntries) {
         Map<String, String> map = new HashMap<>();
         int numIncludes = randomIntBetween(minEntries, maxEntries);
@@ -146,10 +159,32 @@ public class AllocateActionTests extends AbstractActionTestCase<AllocateAction>
             (key, value) -> expectedSettings.put(IndexMetadata.INDEX_ROUTING_EXCLUDE_GROUP_SETTING.getKey() + key, value));
         action.getRequire().forEach(
             (key, value) -> expectedSettings.put(IndexMetadata.INDEX_ROUTING_REQUIRE_GROUP_SETTING.getKey() + key, value));
+        if (action.getTotalShardsPerNode() != null) {
+            expectedSettings.put(ShardsLimitAllocationDecider.INDEX_TOTAL_SHARDS_PER_NODE_SETTING.getKey(), action.getTotalShardsPerNode());
+        }
+
         assertThat(firstStep.getSettings(), equalTo(expectedSettings.build()));
         AllocationRoutedStep secondStep = (AllocationRoutedStep) steps.get(1);
         assertEquals(expectedSecondStepKey, secondStep.getKey());
         assertEquals(nextStepKey, secondStep.getNextStepKey());
     }
 
+    public void testTotalNumberOfShards() throws Exception {
+        Integer totalShardsPerNode = randomIntBetween(-1, 1000);
+        Integer numberOfReplicas = randomIntBetween(0, 4);
+        AllocateAction action = new AllocateAction(numberOfReplicas, totalShardsPerNode, null, null, null);
+        String phase = randomAlphaOfLengthBetween(1, 10);
+        StepKey nextStepKey = new StepKey(randomAlphaOfLengthBetween(1, 10), randomAlphaOfLengthBetween(1, 10),
+            randomAlphaOfLengthBetween(1, 10));
+        List<Step> steps = action.toSteps(null, phase, nextStepKey);
+        UpdateSettingsStep firstStep = (UpdateSettingsStep) steps.get(0);
+        assertEquals(totalShardsPerNode, firstStep.getSettings().getAsInt(INDEX_TOTAL_SHARDS_PER_NODE_SETTING.getKey(), null));
+
+        totalShardsPerNode = null;
+        action = new AllocateAction(numberOfReplicas, totalShardsPerNode, null, null, null);
+        steps = action.toSteps(null, phase, nextStepKey);
+        firstStep = (UpdateSettingsStep) steps.get(0);
+        assertEquals(null, firstStep.getSettings().get(INDEX_TOTAL_SHARDS_PER_NODE_SETTING.getKey()));
+    }
+
 }

x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ilm/PhaseCacheManagementTests.java

@@ -236,7 +236,7 @@ public class PhaseCacheManagementTests extends ESTestCase {
 
         Map<String, LifecycleAction> actions = new HashMap<>();
         actions.put("forcemerge", new ForceMergeAction(5, null));
-        actions.put("allocate", new AllocateAction(1, null, null, null));
+        actions.put("allocate", new AllocateAction(1, 20, null, null, null));
         PhaseExecutionInfo pei = new PhaseExecutionInfo("policy", new Phase("wonky", TimeValue.ZERO, actions), 1, 1);
         String phaseDef = Strings.toString(pei);
         logger.info("--> phaseDef: {}", phaseDef);

x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ilm/TimeseriesLifecycleTypeTests.java

@@ -55,7 +55,7 @@ import static org.hamcrest.Matchers.notNullValue;
 public class TimeseriesLifecycleTypeTests extends ESTestCase {
 
     private static final AllocateAction TEST_ALLOCATE_ACTION =
-        new AllocateAction(2, Collections.singletonMap("node", "node1"),null, null);
+        new AllocateAction(2, 20, Collections.singletonMap("node", "node1"),null, null);
     private static final DeleteAction TEST_DELETE_ACTION = new DeleteAction();
     private static final WaitForSnapshotAction TEST_WAIT_FOR_SNAPSHOT_ACTION = new WaitForSnapshotAction("policy");
     private static final ForceMergeAction TEST_FORCE_MERGE_ACTION = new ForceMergeAction(1, null);
@@ -634,7 +634,7 @@ public class TimeseriesLifecycleTypeTests extends ESTestCase {
         {
             // the allocate action only specifies the number of replicas
             Map<String, LifecycleAction> actions = new HashMap<>();
-            actions.put(TEST_ALLOCATE_ACTION.getWriteableName(), new AllocateAction(2, null, null, null));
+            actions.put(TEST_ALLOCATE_ACTION.getWriteableName(), new AllocateAction(2, 20, null, null, null));
             Phase phase = new Phase(WARM_PHASE, TimeValue.ZERO, actions);
             assertThat(TimeseriesLifecycleType.shouldInjectMigrateStepForPhase(phase), is(true));
         }
@@ -862,7 +862,8 @@ public class TimeseriesLifecycleTypeTests extends ESTestCase {
         return Arrays.asList(availableActionNames).stream().map(n -> {
             switch (n) {
             case AllocateAction.NAME:
-                return new AllocateAction(null, Collections.singletonMap("foo", "bar"), Collections.emptyMap(), Collections.emptyMap());
+                return new AllocateAction(null, null, Collections.singletonMap("foo", "bar"), Collections.emptyMap(),
+                    Collections.emptyMap());
             case DeleteAction.NAME:
                 return new DeleteAction();
             case ForceMergeAction.NAME:

x-pack/plugin/ilm/qa/multi-node/src/javaRestTest/java/org/elasticsearch/xpack/MigrateToDataTiersIT.java

@@ -100,11 +100,11 @@ public class MigrateToDataTiersIT extends ESRestTestCase {
         Map<String, LifecycleAction> warmActions = new HashMap<>();
         warmActions.put(SetPriorityAction.NAME, new SetPriorityAction(50));
         warmActions.put(ForceMergeAction.NAME, new ForceMergeAction(1, null));
-        warmActions.put(AllocateAction.NAME, new AllocateAction(null, singletonMap("data", "warm"), null, null));
+        warmActions.put(AllocateAction.NAME, new AllocateAction(null, null, singletonMap("data", "warm"), null, null));
         warmActions.put(ShrinkAction.NAME, new ShrinkAction(1, null));
         Map<String, LifecycleAction> coldActions = new HashMap<>();
         coldActions.put(SetPriorityAction.NAME, new SetPriorityAction(0));
-        coldActions.put(AllocateAction.NAME, new AllocateAction(0, null, null, singletonMap("data", "cold")));
+        coldActions.put(AllocateAction.NAME, new AllocateAction(0, null, null, null, singletonMap("data", "cold")));
 
         createPolicy(client(), policy,
             new Phase("hot", TimeValue.ZERO, hotActions),
@@ -221,11 +221,11 @@ public class MigrateToDataTiersIT extends ESRestTestCase {
         Map<String, LifecycleAction> warmActions = new HashMap<>();
         warmActions.put(SetPriorityAction.NAME, new SetPriorityAction(50));
         warmActions.put(ForceMergeAction.NAME, new ForceMergeAction(1, null));
-        warmActions.put(AllocateAction.NAME, new AllocateAction(null, singletonMap("data", "warm"), null, null));
+        warmActions.put(AllocateAction.NAME, new AllocateAction(null, null, singletonMap("data", "warm"), null, null));
         warmActions.put(ShrinkAction.NAME, new ShrinkAction(1, null));
         Map<String, LifecycleAction> coldActions = new HashMap<>();
         coldActions.put(SetPriorityAction.NAME, new SetPriorityAction(0));
-        coldActions.put(AllocateAction.NAME, new AllocateAction(0, null, null, singletonMap("data", "cold")));
+        coldActions.put(AllocateAction.NAME, new AllocateAction(0, null, null, null, singletonMap("data", "cold")));
 
         createPolicy(client(), policy,
             new Phase("hot", TimeValue.ZERO, hotActions),

x-pack/plugin/ilm/qa/multi-node/src/javaRestTest/java/org/elasticsearch/xpack/TimeSeriesRestDriver.java

@@ -176,12 +176,14 @@ public final class TimeSeriesRestDriver {
         Map<String, LifecycleAction> warmActions = new HashMap<>();
         warmActions.put(SetPriorityAction.NAME, new SetPriorityAction(50));
         warmActions.put(ForceMergeAction.NAME, new ForceMergeAction(1, null));
-        warmActions.put(AllocateAction.NAME, new AllocateAction(1, singletonMap("_name", "javaRestTest-0,javaRestTest-1,javaRestTest-2," +
+        warmActions.put(AllocateAction.NAME, new AllocateAction(1, null, singletonMap("_name", "javaRestTest-0,javaRestTest-1," +
+            "javaRestTest-2," +
             "javaRestTest-3"), null, null));
         warmActions.put(ShrinkAction.NAME, new ShrinkAction(1, null));
         Map<String, LifecycleAction> coldActions = new HashMap<>();
         coldActions.put(SetPriorityAction.NAME, new SetPriorityAction(0));
-        coldActions.put(AllocateAction.NAME, new AllocateAction(0, singletonMap("_name", "javaRestTest-0,javaRestTest-1,javaRestTest-2," +
+        coldActions.put(AllocateAction.NAME, new AllocateAction(0, null, singletonMap("_name", "javaRestTest-0,javaRestTest-1," +
+            "javaRestTest-2," +
             "javaRestTest-3"), null, null));
         Map<String, Phase> phases = new HashMap<>();
         phases.put("hot", new Phase("hot", hotTime, hotActions));

x-pack/plugin/ilm/qa/multi-node/src/javaRestTest/java/org/elasticsearch/xpack/ilm/ChangePolicyforIndexIT.java

@@ -64,13 +64,13 @@ public class ChangePolicyforIndexIT extends ESRestTestCase {
         Map<String, Phase> phases1 = new HashMap<>();
         phases1.put("hot", new Phase("hot", TimeValue.ZERO, singletonMap(RolloverAction.NAME, new RolloverAction(null, null, null, 1L))));
         phases1.put("warm", new Phase("warm", TimeValue.ZERO,
-                singletonMap(AllocateAction.NAME, new AllocateAction(1, singletonMap("_name", "foobarbaz"), null, null))));
+                singletonMap(AllocateAction.NAME, new AllocateAction(1, null, singletonMap("_name", "foobarbaz"), null, null))));
         LifecyclePolicy lifecyclePolicy1 = new LifecyclePolicy("policy_1", phases1);
         Map<String, Phase> phases2 = new HashMap<>();
         phases2.put("hot", new Phase("hot", TimeValue.ZERO, singletonMap(RolloverAction.NAME, new RolloverAction(null, null, null, 1000L))));
         phases2.put("warm", new Phase("warm", TimeValue.ZERO,
                 singletonMap(AllocateAction.NAME,
-                    new AllocateAction(1, singletonMap("_name", "javaRestTest-0,javaRestTest-1,javaRestTest-2,javaRestTest-3"),
+                    new AllocateAction(1, null, singletonMap("_name", "javaRestTest-0,javaRestTest-1,javaRestTest-2,javaRestTest-3"),
                         null, null))));
         LifecyclePolicy lifecyclePolicy2 = new LifecyclePolicy("policy_2", phases2);
         // PUT policy_1 and policy_2

x-pack/plugin/ilm/qa/multi-node/src/javaRestTest/java/org/elasticsearch/xpack/ilm/TimeSeriesLifecycleActionsIT.java

@@ -168,7 +168,7 @@ public class TimeSeriesLifecycleActionsIT extends ESRestTestCase {
         createIndexWithSettings(client(), index, alias, Settings.builder().put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 2)
             .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 0));
         String allocateNodeName = "javaRestTest-0,javaRestTest-1,javaRestTest-2,javaRestTest-3";
-        AllocateAction allocateAction = new AllocateAction(null, singletonMap("_name", allocateNodeName), null, null);
+        AllocateAction allocateAction = new AllocateAction(null, null, singletonMap("_name", allocateNodeName), null, null);
         String endPhase = randomFrom("warm", "cold");
         createNewSingletonPolicy(client(), policy, endPhase, allocateAction);
         updatePolicy(client(), index, policy);
@@ -183,7 +183,7 @@ public class TimeSeriesLifecycleActionsIT extends ESRestTestCase {
         int finalNumReplicas = (numReplicas + 1) % 2;
         createIndexWithSettings(client(), index, alias, Settings.builder().put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, numShards)
             .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, numReplicas));
-        AllocateAction allocateAction = new AllocateAction(finalNumReplicas, null, null, null);
+        AllocateAction allocateAction = new AllocateAction(finalNumReplicas, null, null, null, null);
         String endPhase = randomFrom("warm", "cold");
         createNewSingletonPolicy(client(), policy, endPhase, allocateAction);
         updatePolicy(client(), index, policy);

x-pack/plugin/ilm/src/main/java/org/elasticsearch/xpack/cluster/metadata/MetadataMigrateToDataTiersRoutingService.java

@@ -323,7 +323,8 @@ public final class MetadataMigrateToDataTiersRoutingService {
                 if (allocateAction.getNumberOfReplicas() != null) {
                     // keep the number of replicas configuration
                     AllocateAction updatedAllocateAction =
-                        new AllocateAction(allocateAction.getNumberOfReplicas(), null, null, null);
+                        new AllocateAction(allocateAction.getNumberOfReplicas(), allocateAction.getTotalShardsPerNode(),
+                            null, null, null);
                     actionMap.put(allocateAction.getWriteableName(), updatedAllocateAction);
                     logger.debug("ILM policy [{}], phase [{}]: updated the allocate action to [{}]", lifecyclePolicy.getName(),
                         phase.getName(), allocateAction);

x-pack/plugin/ilm/src/test/java/org/elasticsearch/xpack/cluster/metadata/MetadataMigrateToDataTiersRoutingServiceTests.java

@@ -90,8 +90,8 @@ public class MetadataMigrateToDataTiersRoutingServiceTests extends ESTestCase {
 
     public void testMigrateIlmPolicyForIndexWithoutILMMetadata() {
         ShrinkAction shrinkAction = new ShrinkAction(2, null);
-        AllocateAction warmAllocateAction = new AllocateAction(null, Map.of("data", "warm"), null, Map.of("rack", "rack1"));
-        AllocateAction coldAllocateAction = new AllocateAction(0, null, null, Map.of("data", "cold"));
+        AllocateAction warmAllocateAction = new AllocateAction(null, null, Map.of("data", "warm"), null, Map.of("rack", "rack1"));
+        AllocateAction coldAllocateAction = new AllocateAction(0, null, null, null, Map.of("data", "cold"));
         SetPriorityAction warmSetPriority = new SetPriorityAction(100);
         LifecyclePolicyMetadata policyMetadata = getWarmColdPolicyMeta(warmSetPriority, shrinkAction, warmAllocateAction,
             coldAllocateAction);
@@ -125,7 +125,7 @@ public class MetadataMigrateToDataTiersRoutingServiceTests extends ESTestCase {
 
     public void testMigrateIlmPolicyFOrPhaseWithDeactivatedMigrateAction() {
         ShrinkAction shrinkAction = new ShrinkAction(2, null);
-        AllocateAction warmAllocateAction = new AllocateAction(null, Map.of("data", "warm"), null, Map.of("rack", "rack1"));
+        AllocateAction warmAllocateAction = new AllocateAction(null, null, Map.of("data", "warm"), null, Map.of("rack", "rack1"));
         MigrateAction deactivatedMigrateAction = new MigrateAction(false);
 
         LifecyclePolicy policy = new LifecyclePolicy(lifecycleName,
@@ -160,8 +160,8 @@ public class MetadataMigrateToDataTiersRoutingServiceTests extends ESTestCase {
     @SuppressWarnings("unchecked")
     public void testMigrateIlmPolicyRefreshesCachedPhase() {
         ShrinkAction shrinkAction = new ShrinkAction(2, null);
-        AllocateAction warmAllocateAction = new AllocateAction(null, Map.of("data", "warm"), null, Map.of("rack", "rack1"));
-        AllocateAction coldAllocateAction = new AllocateAction(0, null, null, Map.of("data", "cold"));
+        AllocateAction warmAllocateAction = new AllocateAction(null, null, Map.of("data", "warm"), null, Map.of("rack", "rack1"));
+        AllocateAction coldAllocateAction = new AllocateAction(0, null, null, null, Map.of("data", "cold"));
         SetPriorityAction warmSetPriority = new SetPriorityAction(100);
         LifecyclePolicyMetadata policyMetadata = getWarmColdPolicyMeta(warmSetPriority, shrinkAction, warmAllocateAction,
             coldAllocateAction);
@@ -337,11 +337,12 @@ public class MetadataMigrateToDataTiersRoutingServiceTests extends ESTestCase {
     }
 
     public void testAllocateActionDefinesRoutingRules() {
-        assertThat(allocateActionDefinesRoutingRules("data", new AllocateAction(null, Map.of("data", "cold"), null, null)), is(true));
-        assertThat(allocateActionDefinesRoutingRules("data", new AllocateAction(null, null, Map.of("data", "cold"), null)), is(true));
-        assertThat(allocateActionDefinesRoutingRules("data", new AllocateAction(null, Map.of("another_attribute", "rack1"), null,
+        assertThat(allocateActionDefinesRoutingRules("data", new AllocateAction(null, null, Map.of("data", "cold"), null, null)), is(true));
+        assertThat(allocateActionDefinesRoutingRules("data", new AllocateAction(null, null, null, Map.of("data", "cold"), null)), is(true));
+        assertThat(allocateActionDefinesRoutingRules("data", new AllocateAction(null, null, Map.of("another_attribute", "rack1"), null,
             Map.of("data", "cold"))), is(true));
-        assertThat(allocateActionDefinesRoutingRules("data", new AllocateAction(null, null, null, Map.of("another_attribute", "cold"))),
+        assertThat(allocateActionDefinesRoutingRules("data", new AllocateAction(null, null, null, null, Map.of("another_attribute",
+            "cold"))),
             is(false));
         assertThat(allocateActionDefinesRoutingRules("data", null), is(false));
     }
@@ -547,8 +548,9 @@ public class MetadataMigrateToDataTiersRoutingServiceTests extends ESTestCase {
     }
 
     public void testMigrateToDataTiersRouting() {
-        AllocateAction allocateActionWithDataAttribute = new AllocateAction(null, Map.of("data", "warm"), null, Map.of("rack", "rack1"));
-        AllocateAction allocateActionWithOtherAttribute = new AllocateAction(0, null, null, Map.of("other", "cold"));
+        AllocateAction allocateActionWithDataAttribute = new AllocateAction(null, null, Map.of("data", "warm"), null, Map.of("rack",
+            "rack1"));
+        AllocateAction allocateActionWithOtherAttribute = new AllocateAction(0, null, null, null, Map.of("other", "cold"));
 
         LifecyclePolicy policyToMigrate = new LifecyclePolicy(lifecycleName,
             Map.of("warm",