Autoscaling policy roles specification (#64222)

Add a roles specification to autoscaling policies. This is used to map
the policy to a set of nodes governed by the policy. The list of roles
is mandatory when adding a policy, optional on updates.

This commit also removes the outer level "policy" element from autoscaling
policy PUT and GET requests.

docs/reference/autoscaling/apis/delete-autoscaling-policy.asciidoc

@@ -15,10 +15,9 @@ Delete autoscaling policy.
 --------------------------------------------------
 PUT /_autoscaling/policy/my_autoscaling_policy
 {
-  "policy": {
-    "deciders": {
-      "fixed": {
-      }
+  "roles": [],
+  "deciders": {
+    "fixed": {
     }
   }
 }

docs/reference/autoscaling/apis/get-autoscaling-policy.asciidoc

@@ -15,10 +15,9 @@ Get autoscaling policy.
 --------------------------------------------------
 PUT /_autoscaling/policy/my_autoscaling_policy
 {
-  "policy": {
-    "deciders": {
-      "fixed": {
-      }
+  "roles" : [],
+  "deciders": {
+    "fixed": {
     }
   }
 }
@@ -70,9 +69,9 @@ The API returns the following result:
 [source,console-result]
 --------------------------------------------------
 {
-  "policy": {
-     "deciders": <deciders>
-  }
+   "roles": <roles>,
+   "deciders": <deciders>
 }
 --------------------------------------------------
-// TEST[s/<deciders>/$body.policy.deciders/]
+// TEST[s/<roles>/$body.roles/]
+// TEST[s/<deciders>/$body.deciders/]

docs/reference/autoscaling/apis/put-autoscaling-policy.asciidoc

@@ -15,10 +15,9 @@ Put autoscaling policy.
 --------------------------------------------------
 PUT /_autoscaling/policy/<name>
 {
-  "policy": {
-    "deciders": {
-      "fixed": {
-      }
+  "roles": [],
+  "deciders": {
+    "fixed": {
     }
   }
 }
@@ -51,16 +50,16 @@ This API puts an autoscaling policy with the provided name.
 ==== {api-examples-title}
 
 This example puts an autoscaling policy named `my_autoscaling_policy` using the
-fixed autoscaling decider.
+fixed autoscaling decider, applying to the set of nodes having (only) the
+"data_hot" role.
 
 [source,console]
 --------------------------------------------------
 PUT /_autoscaling/policy/my_autoscaling_policy
 {
-  "policy": {
-    "deciders": {
-      "fixed": {
-      }
+  "roles" : [ "data_hot" ],
+  "deciders": {
+    "fixed": {
     }
   }
 }

x-pack/plugin/autoscaling/qa/rest/src/yamlRestTest/resources/rest-api-spec/test/autoscaling/delete_autoscaling_policy.yml

@@ -4,9 +4,9 @@
       autoscaling.put_autoscaling_policy:
         name: my_autoscaling_policy
         body:
-          policy:
-            deciders:
-              fixed: {}
+          roles: []
+          deciders:
+            fixed: {}
 
   - match: { "acknowledged": true }
 

x-pack/plugin/autoscaling/qa/rest/src/yamlRestTest/resources/rest-api-spec/test/autoscaling/get_autoscaling_capacity.yml

@@ -11,12 +11,12 @@
       autoscaling.put_autoscaling_policy:
         name: my_autoscaling_policy
         body:
-          policy:
-            deciders:
-              fixed:
-                storage: 1337b
-                memory: 7331b
-                nodes: 10
+          roles : []
+          deciders:
+            fixed:
+              storage: 1337b
+              memory: 7331b
+              nodes: 10
 
   - match: { "acknowledged": true }
 

x-pack/plugin/autoscaling/qa/rest/src/yamlRestTest/resources/rest-api-spec/test/autoscaling/get_autoscaling_policy.yml

@@ -4,9 +4,9 @@
       autoscaling.put_autoscaling_policy:
         name: my_autoscaling_policy
         body:
-          policy:
-            deciders:
-              fixed: {}
+          roles: [ master ]
+          deciders:
+            fixed: {}
 
   - match: { "acknowledged": true }
 
@@ -14,7 +14,8 @@
       autoscaling.get_autoscaling_policy:
         name: my_autoscaling_policy
 
-  - match: { policy.deciders.fixed: {} }
+  - match: { roles: [ master ] }
+  - match: { deciders.fixed: {} }
 
   # test cleanup
   - do:

x-pack/plugin/autoscaling/qa/rest/src/yamlRestTest/resources/rest-api-spec/test/autoscaling/put_autoscaling_policy.yml

@@ -4,25 +4,88 @@
       autoscaling.put_autoscaling_policy:
         name: my_autoscaling_policy
         body:
-          policy:
-            deciders:
-              fixed: {}
+          roles: [ master ]
 
   - match: { "acknowledged": true }
 
+  # update deciders
+  - do:
+      autoscaling.put_autoscaling_policy:
+        name: my_autoscaling_policy
+        body:
+          deciders:
+            fixed: {}
+
+  - do:
+      autoscaling.get_autoscaling_policy:
+        name: my_autoscaling_policy
+  - match: { roles: [ master ] }
+  - match: { deciders.fixed: {} }
+
+  # update roles
+  - do:
+      autoscaling.put_autoscaling_policy:
+        name: my_autoscaling_policy
+        body:
+          roles: [ data ]
+
+  - do:
+      autoscaling.get_autoscaling_policy:
+        name: my_autoscaling_policy
+  - match: { roles: [ data ] }
+  - match: { deciders.fixed: {} }
+
   # test cleanup
   - do:
       autoscaling.delete_autoscaling_policy:
         name: my_autoscaling_policy
 
 ---
+"Test add autoscaling policy with deciders":
+  - do:
+      autoscaling.put_autoscaling_policy:
+        name: my_autoscaling_policy
+        body:
+          roles: [ master ]
+          deciders:
+            fixed: {}
+
+  - match: { "acknowledged": true }
+
+  - do:
+      autoscaling.get_autoscaling_policy:
+        name: my_autoscaling_policy
+  - match: { roles: [ master ] }
+  - match: { deciders.fixed: {} }
+
+  # test cleanup
+  - do:
+      autoscaling.delete_autoscaling_policy:
+        name: my_autoscaling_policy
+---
 "Test put autoscaling policy with non-existent decider":
   - do:
       catch: bad_request
       autoscaling.put_autoscaling_policy:
         name: my_autoscaling_policy
         body:
-          policy:
-            deciders:
-              does_not_exist: {}
+          roles: [ master ]
+          deciders:
+            does_not_exist: {}
 
+---
+"Test put autoscaling policy with non-existent roles":
+  - do:
+      catch: bad_request
+      autoscaling.put_autoscaling_policy:
+        name: my_autoscaling_policy
+        body:
+          roles: [ non-existing ]
+
+---
+"Test add autoscaling policy with no roles":
+  - do:
+      catch: bad_request
+      autoscaling.put_autoscaling_policy:
+        name: my_autoscaling_policy
+        body: {}

x-pack/plugin/autoscaling/src/internalClusterTest/java/org/elasticsearch/xpack/autoscaling/action/TransportDeleteAutoscalingPolicyActionIT.java

@@ -23,7 +23,11 @@ public class TransportDeleteAutoscalingPolicyActionIT extends AutoscalingIntegTe
 
     public void testDeletePolicy() {
         final AutoscalingPolicy policy = randomAutoscalingPolicy();
-        final PutAutoscalingPolicyAction.Request putRequest = new PutAutoscalingPolicyAction.Request(policy);
+        final PutAutoscalingPolicyAction.Request putRequest = new PutAutoscalingPolicyAction.Request(
+            policy.name(),
+            policy.roles(),
+            policy.deciders()
+        );
         assertAcked(client().execute(PutAutoscalingPolicyAction.INSTANCE, putRequest).actionGet());
         // we trust that the policy is in the cluster state since we have tests for putting policies
         final DeleteAutoscalingPolicyAction.Request deleteRequest = new DeleteAutoscalingPolicyAction.Request(policy.name());

x-pack/plugin/autoscaling/src/internalClusterTest/java/org/elasticsearch/xpack/autoscaling/action/TransportGetAutoscalingPolicyActionIT.java

@@ -20,7 +20,11 @@ public class TransportGetAutoscalingPolicyActionIT extends AutoscalingIntegTestC
     public void testGetPolicy() {
         final String name = randomAlphaOfLength(8);
         final AutoscalingPolicy expectedPolicy = randomAutoscalingPolicyOfName(name);
-        final PutAutoscalingPolicyAction.Request putRequest = new PutAutoscalingPolicyAction.Request(expectedPolicy);
+        final PutAutoscalingPolicyAction.Request putRequest = new PutAutoscalingPolicyAction.Request(
+            expectedPolicy.name(),
+            expectedPolicy.roles(),
+            expectedPolicy.deciders()
+        );
         assertAcked(client().execute(PutAutoscalingPolicyAction.INSTANCE, putRequest).actionGet());
         // we trust that the policy is in the cluster state since we have tests for putting policies
         final GetAutoscalingPolicyAction.Request getRequest = new GetAutoscalingPolicyAction.Request(name);

x-pack/plugin/autoscaling/src/internalClusterTest/java/org/elasticsearch/xpack/autoscaling/action/TransportPutAutoscalingPolicyActionIT.java

@@ -10,6 +10,7 @@ import org.elasticsearch.cluster.ClusterState;
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.xpack.autoscaling.AutoscalingIntegTestCase;
 import org.elasticsearch.xpack.autoscaling.AutoscalingMetadata;
+import org.elasticsearch.xpack.autoscaling.AutoscalingTestCase;
 import org.elasticsearch.xpack.autoscaling.policy.AutoscalingPolicy;
 
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
@@ -32,7 +33,11 @@ public class TransportPutAutoscalingPolicyActionIT extends AutoscalingIntegTestC
 
     public void testUpdatePolicy() {
         final AutoscalingPolicy policy = putRandomAutoscalingPolicy();
-        final AutoscalingPolicy updatedPolicy = new AutoscalingPolicy(policy.name(), mutateAutoscalingDeciders(policy.deciders()));
+        final AutoscalingPolicy updatedPolicy = new AutoscalingPolicy(
+            policy.name(),
+            AutoscalingTestCase.randomRoles(),
+            mutateAutoscalingDeciders(policy.deciders())
+        );
         putAutoscalingPolicy(updatedPolicy);
         final ClusterState state = client().admin().cluster().prepareState().get().getState();
         final AutoscalingMetadata metadata = state.metadata().custom(AutoscalingMetadata.NAME);
@@ -59,7 +64,11 @@ public class TransportPutAutoscalingPolicyActionIT extends AutoscalingIntegTestC
     }
 
     private void putAutoscalingPolicy(final AutoscalingPolicy policy) {
-        final PutAutoscalingPolicyAction.Request request = new PutAutoscalingPolicyAction.Request(policy);
+        final PutAutoscalingPolicyAction.Request request = new PutAutoscalingPolicyAction.Request(
+            policy.name(),
+            policy.roles(),
+            policy.deciders()
+        );
         assertAcked(client().execute(PutAutoscalingPolicyAction.INSTANCE, request).actionGet());
     }
 

x-pack/plugin/autoscaling/src/main/java/org/elasticsearch/xpack/autoscaling/action/GetAutoscalingPolicyAction.java

@@ -94,11 +94,7 @@ public class GetAutoscalingPolicyAction extends ActionType<GetAutoscalingPolicyA
 
         @Override
         public XContentBuilder toXContent(final XContentBuilder builder, final Params params) throws IOException {
-            builder.startObject();
-            {
-                builder.field("policy", policy);
-            }
-            builder.endObject();
+            policy.toXContent(builder, params);
             return builder;
         }
 

x-pack/plugin/autoscaling/src/main/java/org/elasticsearch/xpack/autoscaling/action/PutAutoscalingPolicyAction.java

@@ -10,15 +10,26 @@ import org.elasticsearch.action.ActionRequestValidationException;
 import org.elasticsearch.action.ActionType;
 import org.elasticsearch.action.support.master.AcknowledgedRequest;
 import org.elasticsearch.action.support.master.AcknowledgedResponse;
-import org.elasticsearch.common.ParseField;
+import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.util.set.Sets;
 import org.elasticsearch.common.xcontent.ConstructingObjectParser;
 import org.elasticsearch.common.xcontent.XContentParser;
+import org.elasticsearch.xpack.autoscaling.capacity.AutoscalingDeciderConfiguration;
 import org.elasticsearch.xpack.autoscaling.policy.AutoscalingPolicy;
 
 import java.io.IOException;
+import java.util.AbstractMap;
+import java.util.List;
+import java.util.Map;
 import java.util.Objects;
+import java.util.SortedMap;
+import java.util.SortedSet;
+import java.util.TreeMap;
+import java.util.function.Function;
+import java.util.function.Predicate;
+import java.util.stream.Collectors;
 
 public class PutAutoscalingPolicyAction extends ActionType<AcknowledgedResponse> {
 
@@ -31,62 +42,126 @@ public class PutAutoscalingPolicyAction extends ActionType<AcknowledgedResponse>
 
     public static class Request extends AcknowledgedRequest<Request> {
 
-        static final ParseField POLICY_FIELD = new ParseField("policy");
-
         @SuppressWarnings("unchecked")
-        private static final ConstructingObjectParser<Request, String> PARSER = new ConstructingObjectParser<>(
-            "put_autoscaling_policy_request",
-            a -> new Request((AutoscalingPolicy) a[0])
-        );
+        private static final ConstructingObjectParser<Request, String> PARSER;
 
         static {
-            PARSER.declareObject(ConstructingObjectParser.constructorArg(), AutoscalingPolicy::parse, POLICY_FIELD);
+            PARSER = new ConstructingObjectParser<>("put_autocaling_policy_request", false, (c, name) -> {
+                @SuppressWarnings("unchecked")
+                final List<String> roles = (List<String>) c[0];
+                @SuppressWarnings("unchecked")
+                final var deciders = (List<Map.Entry<String, AutoscalingDeciderConfiguration>>) c[1];
+                return new Request(
+                    name,
+                    roles != null ? roles.stream().collect(Sets.toUnmodifiableSortedSet()) : null,
+                    deciders != null
+                        ? new TreeMap<>(deciders.stream().collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)))
+                        : null
+                );
+            });
+            PARSER.declareStringArray(ConstructingObjectParser.optionalConstructorArg(), AutoscalingPolicy.ROLES_FIELD);
+            PARSER.declareNamedObjects(
+                ConstructingObjectParser.optionalConstructorArg(),
+                (p, c, n) -> new AbstractMap.SimpleEntry<>(n, p.namedObject(AutoscalingDeciderConfiguration.class, n, null)),
+                AutoscalingPolicy.DECIDERS_FIELD
+            );
         }
 
+        private final String name;
+        private final SortedSet<String> roles;
+        private final SortedMap<String, AutoscalingDeciderConfiguration> deciders;
+
         public static Request parse(final XContentParser parser, final String name) {
             return PARSER.apply(parser, name);
         }
 
-        private final AutoscalingPolicy policy;
-
-        public AutoscalingPolicy policy() {
-            return policy;
-        }
-
-        public Request(final AutoscalingPolicy policy) {
-            this.policy = Objects.requireNonNull(policy);
+        public Request(
+            final String name,
+            final SortedSet<String> roles,
+            final SortedMap<String, AutoscalingDeciderConfiguration> deciders
+        ) {
+            this.name = name;
+            this.roles = roles;
+            this.deciders = deciders;
         }
 
         public Request(final StreamInput in) throws IOException {
             super(in);
-            policy = new AutoscalingPolicy(in);
+            name = in.readString();
+            if (in.readBoolean()) {
+                roles = in.readSet(StreamInput::readString).stream().collect(Sets.toUnmodifiableSortedSet());
+            } else {
+                roles = null;
+            }
+            if (in.readBoolean()) {
+                deciders = new TreeMap<>(
+                    in.readNamedWriteableList(AutoscalingDeciderConfiguration.class)
+                        .stream()
+                        .collect(Collectors.toMap(AutoscalingDeciderConfiguration::name, Function.identity()))
+                );
+            } else {
+                deciders = null;
+            }
         }
 
         @Override
         public void writeTo(final StreamOutput out) throws IOException {
             super.writeTo(out);
-            policy.writeTo(out);
+            out.writeString(name);
+            if (roles != null) {
+                out.writeBoolean(true);
+                out.writeCollection(roles, StreamOutput::writeString);
+            } else {
+                out.writeBoolean(false);
+            }
+            if (deciders != null) {
+                out.writeBoolean(true);
+                out.writeNamedWriteableList(deciders.values().stream().collect(Collectors.toUnmodifiableList()));
+            } else {
+                out.writeBoolean(false);
+            }
+        }
+
+        public String name() {
+            return name;
+        }
+
+        public SortedSet<String> roles() {
+            return roles;
+        }
+
+        public SortedMap<String, AutoscalingDeciderConfiguration> deciders() {
+            return deciders;
         }
 
         @Override
         public ActionRequestValidationException validate() {
-            // TODO: validate that the policy deciders are non-empty
+            if (roles != null) {
+                List<String> errors = roles.stream()
+                    .filter(Predicate.not(DiscoveryNode.getPossibleRoleNames()::contains))
+                    .collect(Collectors.toList());
+                if (errors.isEmpty() == false) {
+                    ActionRequestValidationException exception = new ActionRequestValidationException();
+                    exception.addValidationErrors(errors);
+                    return exception;
+                }
+            }
+
             return null;
         }
 
         @Override
-        public boolean equals(final Object o) {
+        public boolean equals(Object o) {
             if (this == o) return true;
             if (o == null || getClass() != o.getClass()) return false;
-            final Request request = (Request) o;
-            return policy.equals(request.policy);
+            Request request = (Request) o;
+            return name.equals(request.name) && Objects.equals(roles, request.roles) && Objects.equals(deciders, request.deciders);
         }
 
         @Override
         public int hashCode() {
-            return Objects.hash(policy);
+            return Objects.hash(name, roles, deciders);
         }
-
     }
 
 }

x-pack/plugin/autoscaling/src/main/java/org/elasticsearch/xpack/autoscaling/action/TransportPutAutoscalingPolicyAction.java

@@ -27,6 +27,7 @@ import org.elasticsearch.xpack.autoscaling.AutoscalingMetadata;
 import org.elasticsearch.xpack.autoscaling.policy.AutoscalingPolicy;
 import org.elasticsearch.xpack.autoscaling.policy.AutoscalingPolicyMetadata;
 
+import java.util.Collections;
 import java.util.SortedMap;
 import java.util.TreeMap;
 
@@ -70,7 +71,7 @@ public class TransportPutAutoscalingPolicyAction extends AcknowledgedTransportMa
 
             @Override
             public ClusterState execute(final ClusterState currentState) {
-                return putAutoscalingPolicy(currentState, request.policy(), logger);
+                return putAutoscalingPolicy(currentState, request, logger);
             }
 
         });
@@ -81,7 +82,14 @@ public class TransportPutAutoscalingPolicyAction extends AcknowledgedTransportMa
         return state.blocks().globalBlockedException(ClusterBlockLevel.METADATA_WRITE);
     }
 
-    static ClusterState putAutoscalingPolicy(final ClusterState currentState, final AutoscalingPolicy policy, final Logger logger) {
+    static ClusterState putAutoscalingPolicy(
+        final ClusterState currentState,
+        final PutAutoscalingPolicyAction.Request request,
+        final Logger logger
+    ) {
+        // we allow putting policies with roles that not all nodes in the cluster may understand currently (but the current master must
+        // know it). The expectation is that the mixed cluster situation will be healed soon. See also
+        // AutoscalingCalculateCapacityService#hasUnknownRoles where we shortcut decision making if master node does not know all roles.
         final ClusterState.Builder builder = ClusterState.builder(currentState);
         final AutoscalingMetadata currentMetadata;
         if (currentState.metadata().custom(AutoscalingMetadata.NAME) != null) {
@@ -89,20 +97,41 @@ public class TransportPutAutoscalingPolicyAction extends AcknowledgedTransportMa
         } else {
             currentMetadata = AutoscalingMetadata.EMPTY;
         }
+        final AutoscalingPolicy updatedPolicy;
+        AutoscalingPolicyMetadata existingPolicyMetadata = currentMetadata.policies().get(request.name());
+        if (existingPolicyMetadata == null) {
+            if (request.roles() == null) {
+                throw new IllegalArgumentException(
+                    "new policy " + request.name() + " with no roles defined, must provide empty list for no roles"
+                );
+            }
+            updatedPolicy = new AutoscalingPolicy(
+                request.name(),
+                request.roles(),
+                request.deciders() != null ? request.deciders() : Collections.emptySortedMap()
+            );
+        } else {
+            AutoscalingPolicy existing = existingPolicyMetadata.policy();
+            updatedPolicy = new AutoscalingPolicy(
+                request.name(),
+                request.roles() != null ? request.roles() : existing.roles(),
+                request.deciders() != null ? request.deciders() : existing.deciders()
+            );
+        }
+
         final SortedMap<String, AutoscalingPolicyMetadata> newPolicies = new TreeMap<>(currentMetadata.policies());
-        final AutoscalingPolicyMetadata newPolicyMetadata = new AutoscalingPolicyMetadata(policy);
-        final AutoscalingPolicyMetadata oldPolicyMetadata = newPolicies.put(policy.name(), newPolicyMetadata);
+        final AutoscalingPolicyMetadata newPolicyMetadata = new AutoscalingPolicyMetadata(updatedPolicy);
+        final AutoscalingPolicyMetadata oldPolicyMetadata = newPolicies.put(request.name(), newPolicyMetadata);
         if (oldPolicyMetadata == null) {
-            logger.info("adding autoscaling policy [{}]", policy.name());
+            logger.info("adding autoscaling policy [{}]", request.name());
         } else if (oldPolicyMetadata.equals(newPolicyMetadata)) {
-            logger.info("skipping updating autoscaling policy [{}] due to no change in policy", policy.name());
+            logger.info("skipping updating autoscaling policy [{}] due to no change in policy", request.name());
             return currentState;
         } else {
-            logger.info("updating autoscaling policy [{}]", policy.name());
+            logger.info("updating autoscaling policy [{}]", request.name());
         }
         final AutoscalingMetadata newMetadata = new AutoscalingMetadata(newPolicies);
         builder.metadata(Metadata.builder(currentState.getMetadata()).putCustom(AutoscalingMetadata.NAME, newMetadata).build());
         return builder.build();
     }
-
 }

x-pack/plugin/autoscaling/src/main/java/org/elasticsearch/xpack/autoscaling/capacity/AutoscalingCalculateCapacityService.java

@@ -15,6 +15,7 @@ import org.elasticsearch.cluster.node.DiscoveryNodeRole;
 import org.elasticsearch.common.collect.ImmutableOpenMap;
 import org.elasticsearch.common.collect.Tuple;
 import org.elasticsearch.common.unit.ByteSizeValue;
+import org.elasticsearch.common.util.set.Sets;
 import org.elasticsearch.xpack.autoscaling.Autoscaling;
 import org.elasticsearch.xpack.autoscaling.AutoscalingMetadata;
 import org.elasticsearch.xpack.autoscaling.policy.AutoscalingPolicy;
@@ -23,6 +24,7 @@ import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
 import java.util.SortedMap;
+import java.util.SortedSet;
 import java.util.TreeMap;
 import java.util.function.Function;
 import java.util.stream.Collectors;
@@ -57,7 +59,6 @@ public class AutoscalingCalculateCapacityService {
     }
 
     public SortedMap<String, AutoscalingDeciderResults> calculate(ClusterState state, ClusterInfo clusterInfo) {
-
         AutoscalingMetadata autoscalingMetadata = state.metadata().custom(AutoscalingMetadata.NAME);
         if (autoscalingMetadata != null) {
             return new TreeMap<>(
@@ -73,7 +74,13 @@ public class AutoscalingCalculateCapacityService {
     }
 
     private AutoscalingDeciderResults calculateForPolicy(AutoscalingPolicy policy, ClusterState state, ClusterInfo clusterInfo) {
-        DefaultAutoscalingDeciderContext context = new DefaultAutoscalingDeciderContext(policy.name(), state, clusterInfo);
+        if (hasUnknownRoles(policy)) {
+            return new AutoscalingDeciderResults(
+                AutoscalingCapacity.ZERO,
+                new TreeMap<>(Map.of("_unknown_role", new AutoscalingDeciderResult(null, null)))
+            );
+        }
+        DefaultAutoscalingDeciderContext context = new DefaultAutoscalingDeciderContext(policy.roles(), state, clusterInfo);
         SortedMap<String, AutoscalingDeciderResult> results = policy.deciders()
             .entrySet()
             .stream()
@@ -82,6 +89,14 @@ public class AutoscalingCalculateCapacityService {
         return new AutoscalingDeciderResults(context.currentCapacity, results);
     }
 
+    /**
+     * Check if the policy has unknown roles. This can only happen in mixed clusters, where one master can accept a policy but if it fails
+     * over to an older master before it is also upgraded, one of the roles might not be known.
+     */
+    private boolean hasUnknownRoles(AutoscalingPolicy policy) {
+        return DiscoveryNode.getPossibleRoleNames().containsAll(policy.roles()) == false;
+    }
+
     private <T extends AutoscalingDeciderConfiguration> AutoscalingDeciderResult calculateForDecider(
         T decider,
         AutoscalingDeciderContext context
@@ -94,14 +109,14 @@ public class AutoscalingCalculateCapacityService {
 
     static class DefaultAutoscalingDeciderContext implements AutoscalingDeciderContext {
 
-        private final String tier;
+        private final SortedSet<DiscoveryNodeRole> roles;
         private final ClusterState state;
         private final ClusterInfo clusterInfo;
         private final AutoscalingCapacity currentCapacity;
         private final boolean currentCapacityAccurate;
 
-        DefaultAutoscalingDeciderContext(String tier, ClusterState state, ClusterInfo clusterInfo) {
-            this.tier = tier;
+        DefaultAutoscalingDeciderContext(SortedSet<String> roles, ClusterState state, ClusterInfo clusterInfo) {
+            this.roles = roles.stream().map(DiscoveryNode::getRoleFromRoleName).collect(Sets.toUnmodifiableSortedSet());
             Objects.requireNonNull(state);
             Objects.requireNonNull(clusterInfo);
             this.state = state;
@@ -124,9 +139,14 @@ public class AutoscalingCalculateCapacityService {
             }
         }
 
+        @Override
+        public Set<DiscoveryNode> nodes() {
+            return StreamSupport.stream(state.nodes().spliterator(), false).filter(this::rolesFilter).collect(Collectors.toSet());
+        }
+
         private boolean calculateCurrentCapacityAccurate() {
             return StreamSupport.stream(state.nodes().spliterator(), false)
-                .filter(this::informalTierFilter)
+                .filter(this::rolesFilter)
                 .allMatch(this::nodeHasAccurateCapacity);
         }
 
@@ -137,7 +157,7 @@ public class AutoscalingCalculateCapacityService {
 
         private AutoscalingCapacity calculateCurrentCapacity() {
             return StreamSupport.stream(state.nodes().spliterator(), false)
-                .filter(this::informalTierFilter)
+                .filter(this::rolesFilter)
                 .map(this::resourcesFor)
                 .map(c -> new AutoscalingCapacity(c, c))
                 .reduce(
@@ -167,9 +187,8 @@ public class AutoscalingCalculateCapacityService {
             return diskUsage != null ? diskUsage.getTotalBytes() : -1;
         }
 
-        private boolean informalTierFilter(DiscoveryNode discoveryNode) {
-            return discoveryNode.getRoles().stream().map(DiscoveryNodeRole::roleName).anyMatch(tier::equals)
-                || tier.equals(discoveryNode.getAttributes().get("data"));
+        private boolean rolesFilter(DiscoveryNode discoveryNode) {
+            return discoveryNode.getRoles().equals(roles);
         }
     }
 }

x-pack/plugin/autoscaling/src/main/java/org/elasticsearch/xpack/autoscaling/capacity/AutoscalingDeciderContext.java

@@ -7,13 +7,21 @@
 package org.elasticsearch.xpack.autoscaling.capacity;
 
 import org.elasticsearch.cluster.ClusterState;
+import org.elasticsearch.cluster.node.DiscoveryNode;
+
+import java.util.Set;
 
 public interface AutoscalingDeciderContext {
     ClusterState state();
 
     /**
-     * Return current capacity of tier. Can be null if the capacity of some nodes is unavailable. If a decider relies on this value and
-     * gets a null current capacity, it should return a result with a null requiredCapacity (undecided).
+     * Return current capacity of nodes governed by the policy. Can be null if the capacity of some nodes is unavailable. If a decider
+     * relies on this value and gets a null current capacity, it should return a result with a null requiredCapacity (undecided).
      */
     AutoscalingCapacity currentCapacity();
+
+    /**
+     * Return the nodes governed by the policy.
+     */
+    Set<DiscoveryNode> nodes();
 }

x-pack/plugin/autoscaling/src/main/java/org/elasticsearch/xpack/autoscaling/policy/AutoscalingPolicy.java

@@ -11,6 +11,7 @@ import org.elasticsearch.cluster.Diffable;
 import org.elasticsearch.common.ParseField;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.util.set.Sets;
 import org.elasticsearch.common.xcontent.ConstructingObjectParser;
 import org.elasticsearch.common.xcontent.ToXContentObject;
 import org.elasticsearch.common.xcontent.XContentBuilder;
@@ -23,6 +24,7 @@ import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.SortedMap;
+import java.util.SortedSet;
 import java.util.TreeMap;
 import java.util.function.Function;
 import java.util.stream.Collectors;
@@ -31,6 +33,7 @@ public class AutoscalingPolicy extends AbstractDiffable<AutoscalingPolicy> imple
 
     public static final String NAME = "autoscaling_policy";
 
+    public static final ParseField ROLES_FIELD = new ParseField("roles");
     public static final ParseField DECIDERS_FIELD = new ParseField("deciders");
 
     private static final ConstructingObjectParser<AutoscalingPolicy, String> PARSER;
@@ -38,12 +41,16 @@ public class AutoscalingPolicy extends AbstractDiffable<AutoscalingPolicy> imple
     static {
         PARSER = new ConstructingObjectParser<>(NAME, false, (c, name) -> {
             @SuppressWarnings("unchecked")
-            final var deciders = (List<Map.Entry<String, AutoscalingDeciderConfiguration>>) c[0];
+            final List<String> roles = (List<String>) c[0];
+            @SuppressWarnings("unchecked")
+            final var deciders = (List<Map.Entry<String, AutoscalingDeciderConfiguration>>) c[1];
             return new AutoscalingPolicy(
                 name,
+                roles.stream().collect(Sets.toUnmodifiableSortedSet()),
                 new TreeMap<>(deciders.stream().collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)))
             );
         });
+        PARSER.declareStringArray(ConstructingObjectParser.constructorArg(), ROLES_FIELD);
         PARSER.declareNamedObjects(
             ConstructingObjectParser.constructorArg(),
             (p, c, n) -> new AbstractMap.SimpleEntry<>(n, p.namedObject(AutoscalingDeciderConfiguration.class, n, null)),
@@ -61,20 +68,31 @@ public class AutoscalingPolicy extends AbstractDiffable<AutoscalingPolicy> imple
         return name;
     }
 
+    private final SortedSet<String> roles;
+
+    public SortedSet<String> roles() {
+        return roles;
+    }
+
     private final SortedMap<String, AutoscalingDeciderConfiguration> deciders;
 
     public SortedMap<String, AutoscalingDeciderConfiguration> deciders() {
         return deciders;
     }
 
-    public AutoscalingPolicy(final String name, final SortedMap<String, AutoscalingDeciderConfiguration> deciders) {
+    public AutoscalingPolicy(
+        final String name,
+        SortedSet<String> roles,
+        final SortedMap<String, AutoscalingDeciderConfiguration> deciders
+    ) {
         this.name = Objects.requireNonNull(name);
-        // TODO: validate that the policy deciders are non-empty
+        this.roles = Objects.requireNonNull(roles);
         this.deciders = Objects.requireNonNull(deciders);
     }
 
     public AutoscalingPolicy(final StreamInput in) throws IOException {
         name = in.readString();
+        roles = in.readSet(StreamInput::readString).stream().collect(Sets.toUnmodifiableSortedSet());
         deciders = new TreeMap<>(
             in.readNamedWriteableList(AutoscalingDeciderConfiguration.class)
                 .stream()
@@ -85,6 +103,7 @@ public class AutoscalingPolicy extends AbstractDiffable<AutoscalingPolicy> imple
     @Override
     public void writeTo(final StreamOutput out) throws IOException {
         out.writeString(name);
+        out.writeCollection(roles, StreamOutput::writeString);
         out.writeNamedWriteableList(deciders.values().stream().collect(Collectors.toUnmodifiableList()));
     }
 
@@ -92,6 +111,7 @@ public class AutoscalingPolicy extends AbstractDiffable<AutoscalingPolicy> imple
     public XContentBuilder toXContent(final XContentBuilder builder, final Params params) throws IOException {
         builder.startObject();
         {
+            builder.array(ROLES_FIELD.getPreferredName(), roles.toArray(String[]::new));
             builder.startObject(DECIDERS_FIELD.getPreferredName());
             {
                 for (final Map.Entry<String, AutoscalingDeciderConfiguration> entry : deciders.entrySet()) {
@@ -109,12 +129,12 @@ public class AutoscalingPolicy extends AbstractDiffable<AutoscalingPolicy> imple
         if (this == o) return true;
         if (o == null || getClass() != o.getClass()) return false;
         final AutoscalingPolicy that = (AutoscalingPolicy) o;
-        return name.equals(that.name) && deciders.equals(that.deciders);
+        return name.equals(that.name) && roles.equals(that.roles) && deciders.equals(that.deciders);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(name, deciders);
+        return Objects.hash(name, roles, deciders);
     }
 
 }

x-pack/plugin/autoscaling/src/test/java/org/elasticsearch/xpack/autoscaling/AutoscalingTestCase.java

@@ -6,10 +6,12 @@
 
 package org.elasticsearch.xpack.autoscaling;
 
+import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.common.collect.Tuple;
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.unit.ByteSizeValue;
+import org.elasticsearch.common.util.set.Sets;
 import org.elasticsearch.common.xcontent.NamedXContentRegistry;
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.xpack.autoscaling.capacity.AutoscalingCapacity;
@@ -21,9 +23,11 @@ import org.elasticsearch.xpack.autoscaling.capacity.FixedAutoscalingDeciderServi
 import org.elasticsearch.xpack.autoscaling.policy.AutoscalingPolicy;
 import org.elasticsearch.xpack.autoscaling.policy.AutoscalingPolicyMetadata;
 
+import java.util.BitSet;
 import java.util.List;
 import java.util.Map;
 import java.util.SortedMap;
+import java.util.SortedSet;
 import java.util.TreeMap;
 import java.util.function.Function;
 import java.util.stream.Collectors;
@@ -109,18 +113,28 @@ public abstract class AutoscalingTestCase extends ESTestCase {
     }
 
     public static AutoscalingPolicy randomAutoscalingPolicyOfName(final String name) {
-        return new AutoscalingPolicy(name, randomAutoscalingDeciders());
+        return new AutoscalingPolicy(name, randomRoles(), randomAutoscalingDeciders());
     }
 
     public static AutoscalingPolicy mutateAutoscalingPolicy(final AutoscalingPolicy instance) {
-        final SortedMap<String, AutoscalingDeciderConfiguration> deciders;
-        if (randomBoolean()) {
-            // if the policy name did not change, or randomly, use a mutated set of deciders
-            deciders = mutateAutoscalingDeciders(instance.deciders());
-        } else {
-            deciders = instance.deciders();
+        String name = instance.name();
+        SortedSet<String> roles = instance.roles();
+        SortedMap<String, AutoscalingDeciderConfiguration> deciders = instance.deciders();
+        BitSet choice = BitSet.valueOf(new long[] { randomIntBetween(1, 7) });
+        if (choice.get(0)) {
+            name = randomValueOtherThan(instance.name(), () -> randomAlphaOfLength(8));
+        }
+        if (choice.get(1)) {
+            roles = mutateRoles(roles);
         }
-        return new AutoscalingPolicy(randomValueOtherThan(instance.name(), () -> randomAlphaOfLength(8)), deciders);
+        if (choice.get(2)) {
+            deciders = mutateAutoscalingDeciders(deciders);
+        }
+        return new AutoscalingPolicy(name, roles, deciders);
+    }
+
+    protected static SortedSet<String> mutateRoles(SortedSet<String> roles) {
+        return randomValueOtherThan(roles, AutoscalingTestCase::randomRoles);
     }
 
     public static SortedMap<String, AutoscalingDeciderConfiguration> mutateAutoscalingDeciders(
@@ -151,6 +165,10 @@ public abstract class AutoscalingTestCase extends ESTestCase {
         return new AutoscalingMetadata(policies);
     }
 
+    public static SortedSet<String> randomRoles() {
+        return randomSubsetOf(DiscoveryNode.getPossibleRoleNames()).stream().collect(Sets.toUnmodifiableSortedSet());
+    }
+
     public static NamedWriteableRegistry getAutoscalingNamedWriteableRegistry() {
         return new NamedWriteableRegistry(new Autoscaling(Settings.EMPTY).getNamedWriteables());
     }

x-pack/plugin/autoscaling/src/test/java/org/elasticsearch/xpack/autoscaling/action/PutAutoscalingPolicyActionRequestWireSerializingTests.java

@@ -11,8 +11,6 @@ import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.test.AbstractWireSerializingTestCase;
 import org.elasticsearch.xpack.autoscaling.AutoscalingTestCase;
 
-import static org.elasticsearch.xpack.autoscaling.AutoscalingTestCase.randomAutoscalingPolicy;
-
 public class PutAutoscalingPolicyActionRequestWireSerializingTests extends AbstractWireSerializingTestCase<
     PutAutoscalingPolicyAction.Request> {
 
@@ -23,7 +21,7 @@ public class PutAutoscalingPolicyActionRequestWireSerializingTests extends Abstr
 
     @Override
     protected PutAutoscalingPolicyAction.Request createTestInstance() {
-        return new PutAutoscalingPolicyAction.Request(randomAutoscalingPolicy());
+        return TransportPutAutoscalingPolicyActionTests.randomPutAutoscalingPolicyRequest();
     }
 
     @Override

x-pack/plugin/autoscaling/src/test/java/org/elasticsearch/xpack/autoscaling/action/TransportPutAutoscalingPolicyActionTests.java

@@ -54,7 +54,7 @@ public class TransportPutAutoscalingPolicyActionTests extends AutoscalingTestCas
             )
             .build();
         final ClusterState state = ClusterState.builder(new ClusterName(randomAlphaOfLength(8))).blocks(blocks).build();
-        final ClusterBlockException e = action.checkBlock(new PutAutoscalingPolicyAction.Request(randomAutoscalingPolicy()), state);
+        final ClusterBlockException e = action.checkBlock(randomPutAutoscalingPolicyRequest(), state);
         assertThat(e, not(nullValue()));
     }
 
@@ -68,7 +68,7 @@ public class TransportPutAutoscalingPolicyActionTests extends AutoscalingTestCas
         );
         final ClusterBlocks blocks = ClusterBlocks.builder().build();
         final ClusterState state = ClusterState.builder(new ClusterName(randomAlphaOfLength(8))).blocks(blocks).build();
-        final ClusterBlockException e = action.checkBlock(new PutAutoscalingPolicyAction.Request(randomAutoscalingPolicy()), state);
+        final ClusterBlockException e = action.checkBlock(randomPutAutoscalingPolicyRequest(), state);
         assertThat(e, nullValue());
     }
 
@@ -82,16 +82,21 @@ public class TransportPutAutoscalingPolicyActionTests extends AutoscalingTestCas
             currentState = builder.build();
         }
         // put an entirely new policy
-        final AutoscalingPolicy policy = randomAutoscalingPolicy();
+        final PutAutoscalingPolicyAction.Request request = randomPutAutoscalingPolicyRequest();
         final Logger mockLogger = mock(Logger.class);
-        final ClusterState state = TransportPutAutoscalingPolicyAction.putAutoscalingPolicy(currentState, policy, mockLogger);
+        final ClusterState state = TransportPutAutoscalingPolicyAction.putAutoscalingPolicy(currentState, request, mockLogger);
 
         // ensure the new policy is in the updated cluster state
         final AutoscalingMetadata metadata = state.metadata().custom(AutoscalingMetadata.NAME);
         assertNotNull(metadata);
-        assertThat(metadata.policies(), hasKey(policy.name()));
-        assertThat(metadata.policies().get(policy.name()).policy(), equalTo(policy));
-        verify(mockLogger).info("adding autoscaling policy [{}]", policy.name());
+        assertThat(metadata.policies(), hasKey(request.name()));
+        assertThat(metadata.policies().get(request.name()).policy().roles(), equalTo(request.roles()));
+        if (request.deciders() != null) {
+            assertThat(metadata.policies().get(request.name()).policy().deciders(), equalTo(request.deciders()));
+        } else {
+            assertThat(metadata.policies().get(request.name()).policy().deciders(), equalTo(Map.of()));
+        }
+        verify(mockLogger).info("adding autoscaling policy [{}]", request.name());
         verifyNoMoreInteractions(mockLogger);
 
         // ensure that existing policies were preserved
@@ -104,6 +109,24 @@ public class TransportPutAutoscalingPolicyActionTests extends AutoscalingTestCas
         }
     }
 
+    public void testAddPolicyWithNoRoles() {
+        PutAutoscalingPolicyAction.Request request = new PutAutoscalingPolicyAction.Request(
+            randomAlphaOfLength(8),
+            null,
+            randomAutoscalingDeciders()
+        );
+
+        final Logger mockLogger = mock(Logger.class);
+        IllegalArgumentException exception = expectThrows(
+            IllegalArgumentException.class,
+            () -> TransportPutAutoscalingPolicyAction.putAutoscalingPolicy(ClusterState.EMPTY_STATE, request, mockLogger)
+        );
+        assertThat(
+            exception.getMessage(),
+            equalTo("new policy " + request.name() + " with no roles defined, must provide empty list for " + "no roles")
+        );
+    }
+
     public void testUpdatePolicy() {
         final ClusterState currentState;
         {
@@ -116,19 +139,25 @@ public class TransportPutAutoscalingPolicyActionTests extends AutoscalingTestCas
         final AutoscalingMetadata currentMetadata = currentState.metadata().custom(AutoscalingMetadata.NAME);
         final String name = randomFrom(currentMetadata.policies().keySet());
         // add to the existing deciders, to ensure the policy has changed
-        final AutoscalingPolicy policy = new AutoscalingPolicy(
+        final PutAutoscalingPolicyAction.Request request = new PutAutoscalingPolicyAction.Request(
             name,
+            randomBoolean() ? randomRoles() : null,
             mutateAutoscalingDeciders(currentMetadata.policies().get(name).policy().deciders())
         );
+        final AutoscalingPolicy expectedPolicy = new AutoscalingPolicy(
+            name,
+            request.roles() != null ? request.roles() : currentMetadata.policies().get(name).policy().roles(),
+            request.deciders()
+        );
         final Logger mockLogger = mock(Logger.class);
-        final ClusterState state = TransportPutAutoscalingPolicyAction.putAutoscalingPolicy(currentState, policy, mockLogger);
+        final ClusterState state = TransportPutAutoscalingPolicyAction.putAutoscalingPolicy(currentState, request, mockLogger);
 
         // ensure the updated policy is in the updated cluster state
         final AutoscalingMetadata metadata = state.metadata().custom(AutoscalingMetadata.NAME);
         assertNotNull(metadata);
-        assertThat(metadata.policies(), hasKey(policy.name()));
-        assertThat(metadata.policies().get(policy.name()).policy(), equalTo(policy));
-        verify(mockLogger).info("updating autoscaling policy [{}]", policy.name());
+        assertThat(metadata.policies(), hasKey(request.name()));
+        assertThat(metadata.policies().get(request.name()).policy(), equalTo(expectedPolicy));
+        verify(mockLogger).info("updating autoscaling policy [{}]", request.name());
         verifyNoMoreInteractions(mockLogger);
 
         // ensure that existing policies were otherwise preserved
@@ -153,12 +182,24 @@ public class TransportPutAutoscalingPolicyActionTests extends AutoscalingTestCas
         // randomly put an existing policy
         final AutoscalingMetadata currentMetadata = currentState.metadata().custom(AutoscalingMetadata.NAME);
         final AutoscalingPolicy policy = randomFrom(currentMetadata.policies().values()).policy();
+        final PutAutoscalingPolicyAction.Request request = new PutAutoscalingPolicyAction.Request(
+            policy.name(),
+            randomBoolean() ? policy.roles() : null,
+            randomBoolean() ? policy.deciders() : null
+        );
         final Logger mockLogger = mock(Logger.class);
-        final ClusterState state = TransportPutAutoscalingPolicyAction.putAutoscalingPolicy(currentState, policy, mockLogger);
+        final ClusterState state = TransportPutAutoscalingPolicyAction.putAutoscalingPolicy(currentState, request, mockLogger);
 
         assertThat(state, sameInstance(currentState));
         verify(mockLogger).info("skipping updating autoscaling policy [{}] due to no change in policy", policy.name());
         verifyNoMoreInteractions(mockLogger);
     }
 
+    static PutAutoscalingPolicyAction.Request randomPutAutoscalingPolicyRequest() {
+        return new PutAutoscalingPolicyAction.Request(
+            randomAlphaOfLength(8),
+            randomRoles(),
+            randomBoolean() ? randomAutoscalingDeciders() : null
+        );
+    }
 }

x-pack/plugin/autoscaling/src/test/java/org/elasticsearch/xpack/autoscaling/capacity/AutoscalingDeciderResultServiceTests.java → x-pack/plugin/autoscaling/src/test/java/org/elasticsearch/xpack/autoscaling/capacity/AutoscalingCalculateCapacityServiceTests.java

@@ -13,6 +13,7 @@ import org.elasticsearch.cluster.ClusterState;
 import org.elasticsearch.cluster.DiskUsage;
 import org.elasticsearch.cluster.metadata.Metadata;
 import org.elasticsearch.cluster.node.DiscoveryNode;
+import org.elasticsearch.cluster.node.DiscoveryNodeRole;
 import org.elasticsearch.cluster.node.DiscoveryNodes;
 import org.elasticsearch.common.collect.ImmutableOpenMap;
 import org.elasticsearch.common.collect.Tuple;
@@ -22,16 +23,19 @@ import org.elasticsearch.xpack.autoscaling.AutoscalingTestCase;
 import org.elasticsearch.xpack.autoscaling.policy.AutoscalingPolicy;
 import org.elasticsearch.xpack.autoscaling.policy.AutoscalingPolicyMetadata;
 
+import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
 import java.util.SortedMap;
+import java.util.SortedSet;
 import java.util.TreeMap;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
+import java.util.stream.StreamSupport;
 
 import static org.hamcrest.Matchers.equalTo;
 
-public class AutoscalingDeciderResultServiceTests extends AutoscalingTestCase {
+public class AutoscalingCalculateCapacityServiceTests extends AutoscalingTestCase {
     public void testMultiplePoliciesFixedCapacity() {
         AutoscalingCalculateCapacityService service = new AutoscalingCalculateCapacityService(Set.of(new FixedAutoscalingDeciderService()));
         Set<String> policyNames = IntStream.range(0, randomIntBetween(1, 10))
@@ -40,7 +44,7 @@ public class AutoscalingDeciderResultServiceTests extends AutoscalingTestCase {
 
         SortedMap<String, AutoscalingPolicyMetadata> policies = new TreeMap<>(
             policyNames.stream()
-                .map(s -> Tuple.tuple(s, new AutoscalingPolicyMetadata(new AutoscalingPolicy(s, randomFixedDeciders()))))
+                .map(s -> Tuple.tuple(s, new AutoscalingPolicyMetadata(new AutoscalingPolicy(s, randomRoles(), randomFixedDeciders()))))
                 .collect(Collectors.toMap(Tuple::v1, Tuple::v2))
         );
         ClusterState state = ClusterState.builder(ClusterName.DEFAULT)
@@ -106,49 +110,70 @@ public class AutoscalingDeciderResultServiceTests extends AutoscalingTestCase {
     }
 
     public void testContext() {
-        String tier = randomAlphaOfLength(5);
         ClusterState state = ClusterState.builder(ClusterName.DEFAULT).build();
         ClusterInfo info = ClusterInfo.EMPTY;
+        SortedSet<String> roleNames = randomRoles();
         AutoscalingCalculateCapacityService.DefaultAutoscalingDeciderContext context =
-            new AutoscalingCalculateCapacityService.DefaultAutoscalingDeciderContext(tier, state, info);
+            new AutoscalingCalculateCapacityService.DefaultAutoscalingDeciderContext(roleNames, state, info);
 
         assertSame(state, context.state());
-        // there is no nodes in any tier.
+
+        assertThat(context.nodes(), equalTo(Set.of()));
         assertThat(context.currentCapacity(), equalTo(AutoscalingCapacity.ZERO));
 
-        tier = "data";
+        Set<DiscoveryNodeRole> roles = roleNames.stream().map(DiscoveryNode::getRoleFromRoleName).collect(Collectors.toSet());
+        Set<DiscoveryNodeRole> otherRoles = mutateRoles(roleNames).stream()
+            .map(DiscoveryNode::getRoleFromRoleName)
+            .collect(Collectors.toSet());
         state = ClusterState.builder(ClusterName.DEFAULT)
-            .nodes(DiscoveryNodes.builder().add(new DiscoveryNode("nodeId", buildNewFakeTransportAddress(), Version.CURRENT)))
+            .nodes(
+                DiscoveryNodes.builder().add(new DiscoveryNode("nodeId", buildNewFakeTransportAddress(), Map.of(), roles, Version.CURRENT))
+            )
             .build();
-        context = new AutoscalingCalculateCapacityService.DefaultAutoscalingDeciderContext(tier, state, info);
+        context = new AutoscalingCalculateCapacityService.DefaultAutoscalingDeciderContext(roleNames, state, info);
 
+        assertThat(context.nodes().size(), equalTo(1));
+        assertThat(context.nodes(), equalTo(StreamSupport.stream(state.nodes().spliterator(), false).collect(Collectors.toSet())));
         assertNull(context.currentCapacity());
 
         ImmutableOpenMap.Builder<String, DiskUsage> leastUsages = ImmutableOpenMap.<String, DiskUsage>builder();
         ImmutableOpenMap.Builder<String, DiskUsage> mostUsages = ImmutableOpenMap.<String, DiskUsage>builder();
         DiscoveryNodes.Builder nodes = DiscoveryNodes.builder();
+        Set<DiscoveryNode> expectedNodes = new HashSet<>();
         long sumTotal = 0;
         long maxTotal = 0;
         for (int i = 0; i < randomIntBetween(1, 5); ++i) {
             String nodeId = "nodeId" + i;
-            nodes.add(new DiscoveryNode(nodeId, buildNewFakeTransportAddress(), Version.CURRENT));
+            boolean useOtherRoles = randomBoolean();
+            DiscoveryNode node = new DiscoveryNode(
+                nodeId,
+                buildNewFakeTransportAddress(),
+                Map.of(),
+                useOtherRoles ? otherRoles : roles,
+                Version.CURRENT
+            );
+            nodes.add(node);
 
             long total = randomLongBetween(1, 1L << 40);
             long total1 = randomBoolean() ? total : randomLongBetween(0, total);
             long total2 = total1 != total ? total : randomLongBetween(0, total);
             leastUsages.fPut(nodeId, new DiskUsage(nodeId, null, null, total1, randomLongBetween(0, total)));
             mostUsages.fPut(nodeId, new DiskUsage(nodeId, null, null, total2, randomLongBetween(0, total)));
-            sumTotal += total;
-            maxTotal = Math.max(total, maxTotal);
+            if (useOtherRoles == false) {
+                sumTotal += total;
+                maxTotal = Math.max(total, maxTotal);
+                expectedNodes.add(node);
+            }
         }
         state = ClusterState.builder(ClusterName.DEFAULT).nodes(nodes).build();
         info = new ClusterInfo(leastUsages.build(), mostUsages.build(), null, null, null);
-        context = new AutoscalingCalculateCapacityService.DefaultAutoscalingDeciderContext(tier, state, info);
+        context = new AutoscalingCalculateCapacityService.DefaultAutoscalingDeciderContext(roleNames, state, info);
 
+        assertThat(context.nodes(), equalTo(expectedNodes));
         AutoscalingCapacity capacity = context.currentCapacity();
         assertThat(capacity.node().storage(), equalTo(new ByteSizeValue(maxTotal)));
         assertThat(capacity.tier().storage(), equalTo(new ByteSizeValue(sumTotal)));
-        // todo: fix these once we know memory of all node on master.
+        // todo: fix these once we know memory of all nodes on master.
         assertThat(capacity.node().memory(), equalTo(ByteSizeValue.ZERO));
         assertThat(capacity.tier().memory(), equalTo(ByteSizeValue.ZERO));
     }