Add per-node request throttling to workload management

plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/AutoTaggingActionFilter.java

@@ -117,14 +117,45 @@ public LogicalOperator getLogicalOperator() {
             }
         }
 
+        List<String> principalValues = null;
         if (featureType.getAllowedAttributesRegistry().containsKey(PRINCIPAL_ATTRIBUTE_NAME)) {
             Attribute attribute = featureType.getAllowedAttributesRegistry().get(PRINCIPAL_ATTRIBUTE_NAME);
             assert attributeExtensions.containsKey(attribute);
-            attributeExtractors.add(attributeExtensions.get(attribute).getAttributeExtractor());
+            final AttributeExtractor<String> extractor = attributeExtensions.get(attribute).getAttributeExtractor();
+            // Materialize once. The value is needed both for label evaluation and for the principal header, and
+            // AttributeExtractor.extract() carries no re-iterability contract -- a stream-backed implementation would
+            // yield nothing the second time and silently disable username/role throttling.
+            final List<String> values = new ArrayList<>();
+            extractor.extract().forEach(values::add);
+            principalValues = values;
+            attributeExtractors.add(new AttributeExtractor<>() {
+                @Override
+                public Attribute getAttribute() {
+                    return extractor.getAttribute();
+                }
+
+                @Override
+                public Iterable<String> extract() {
+                    return values;
+                }
+
+                @Override
+                public LogicalOperator getLogicalOperator() {
+                    return extractor.getLogicalOperator();
+                }
+            });
         }
 
         Optional<String> label = ruleProcessingService.evaluateLabel(attributeExtractors);
         label.ifPresent(s -> threadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, s));
+        // Hand the principal to core-side throttling so it can build per-username / per-role buckets. It goes on the
+        // task, not into the thread context: see WorkloadGroupTask#setThrottlePrincipal.
+        if (principalValues != null && task instanceof WorkloadGroupTask) {
+            String principal = String.join(WorkloadGroupTask.WORKLOAD_GROUP_PRINCIPAL_VALUE_DELIMITER, principalValues);
+            if (principal.isEmpty() == false) {
+                ((WorkloadGroupTask) task).setThrottlePrincipal(principal);
+            }
+        }
         chain.proceed(task, action, request, listener);
     }
 }

plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportCreateWorkloadGroupAction.java

@@ -71,6 +71,15 @@ protected void clusterManagerOperation(
         ClusterState clusterState,
         ActionListener<CreateWorkloadGroupResponse> listener
     ) {
+        try {
+            WorkloadGroupPersistenceService.validateThrottlingIsEnforceable(
+                request.getWorkloadGroup().getMutableWorkloadGroupFragment().getThrottling(),
+                clusterState
+            );
+        } catch (Exception e) {
+            listener.onFailure(e);
+            return;
+        }
         workloadGroupPersistenceService.persistInClusterStateMetadata(request.getWorkloadGroup(), listener);
     }
 

plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportUpdateWorkloadGroupAction.java

@@ -71,6 +71,15 @@ protected void clusterManagerOperation(
         ClusterState clusterState,
         ActionListener<UpdateWorkloadGroupResponse> listener
     ) {
+        try {
+            WorkloadGroupPersistenceService.validateThrottlingIsEnforceable(
+                request.getmMutableWorkloadGroupFragment().getThrottling(),
+                clusterState
+            );
+        } catch (Exception e) {
+            listener.onFailure(e);
+            return;
+        }
         workloadGroupPersistenceService.updateInClusterStateMetadata(request, listener);
     }
 

plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/WorkloadGroupPersistenceService.java

@@ -11,6 +11,7 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.opensearch.ResourceNotFoundException;
+import org.opensearch.Version;
 import org.opensearch.action.support.clustermanager.AcknowledgedResponse;
 import org.opensearch.cluster.AckedClusterStateUpdateTask;
 import org.opensearch.cluster.ClusterState;
@@ -27,12 +28,17 @@
 import org.opensearch.common.settings.Settings;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.rest.RestStatus;
+import org.opensearch.plugin.wlm.WorkloadManagementPlugin;
 import org.opensearch.plugin.wlm.action.CreateWorkloadGroupResponse;
 import org.opensearch.plugin.wlm.action.DeleteWorkloadGroupRequest;
 import org.opensearch.plugin.wlm.action.UpdateWorkloadGroupRequest;
 import org.opensearch.plugin.wlm.action.UpdateWorkloadGroupResponse;
+import org.opensearch.plugin.wlm.rule.WorkloadGroupFeatureType;
+import org.opensearch.rule.autotagging.AutoTaggingRegistry;
+import org.opensearch.rule.autotagging.FeatureType;
 import org.opensearch.wlm.MutableWorkloadGroupFragment;
 import org.opensearch.wlm.ResourceType;
+import org.opensearch.wlm.WorkloadGroupThrottleSettings;
 
 import java.util.Collection;
 import java.util.EnumMap;
@@ -366,4 +372,59 @@ public int getMaxWorkloadGroupCount() {
     public ClusterService getClusterService() {
         return clusterService;
     }
+
+    /**
+     * Rejects a throttling config the cluster cannot actually honour. Both cases below would otherwise return a 200 for
+     * a config that silently never takes effect:
+     * <ul>
+     *   <li>a node older than {@link Version#V_3_9_0} is still in the cluster. {@code throttling} is gated on the wire,
+     *       so the config is dropped when the request or the resulting cluster state crosses that node, and the group
+     *       reads back without it.</li>
+     *   <li>the attribute keys on a principal ({@code username}/{@code role}) but no principal attribute is registered,
+     *       so no bucket can ever be resolved and the limit always fails open.</li>
+     * </ul>
+     * Called from the cluster-manager transport actions rather than from a cluster-state applier or settings update
+     * consumer on purpose: throwing while applying cluster state wedges the cluster-manager.
+     *
+     * @param throttling   the incoming throttling fragment, may be {@code null} or empty (both fine: nothing to honour)
+     * @param clusterState state used to read the oldest node version in the cluster
+     * @throws IllegalArgumentException if the config cannot be enforced
+     */
+    public static void validateThrottlingIsEnforceable(Settings throttling, ClusterState clusterState) {
+        if (throttling == null || throttling.isEmpty()) {
+            return;
+        }
+        Version minNodeVersion = clusterState.nodes().getMinNodeVersion();
+        if (minNodeVersion.before(Version.V_3_9_0)) {
+            throw new IllegalArgumentException(
+                "workload group throttling requires every node to be on "
+                    + Version.V_3_9_0
+                    + " or later, but the oldest node in the cluster is on "
+                    + minNodeVersion
+                    + ". The throttling config would be silently dropped; complete the upgrade first."
+            );
+        }
+        // ATTRIBUTE.get returns "" (its default), not null, when the key is absent -- which is the normal shape of a
+        // partial update that only changes the limit. Only an explicitly principal-keyed attribute is checked here; the
+        // merged config is validated separately.
+        String attribute = WorkloadGroupThrottleSettings.ATTRIBUTE.get(throttling);
+        if (attribute == null || attribute.isEmpty() || "group".equals(attribute)) {
+            return;
+        }
+        try {
+            FeatureType featureType = AutoTaggingRegistry.getFeatureType(WorkloadGroupFeatureType.NAME);
+            if (featureType.getAllowedAttributesRegistry().containsKey(WorkloadManagementPlugin.PRINCIPAL_ATTRIBUTE_NAME) == false) {
+                throw new IllegalArgumentException(
+                    "throttling attribute ["
+                        + attribute
+                        + "] needs a principal attribute provider (the security plugin) to be installed, otherwise the "
+                        + "limit can never be enforced. Use attribute [group] instead."
+                );
+            }
+        } catch (ResourceNotFoundException e) {
+            // Feature type not registered on this node yet. Skip rather than reject a config that is probably fine --
+            // the throttle path fails open anyway, so a false rejection here is worse than a missed warning.
+            logger.debug("WLM feature type not registered; skipping principal-attribute check for throttling config", e);
+        }
+    }
 }

plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/AutoTaggingActionFilterTests.java

@@ -18,6 +18,8 @@
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.action.ActionResponse;
 import org.opensearch.core.common.io.stream.StreamOutput;
+import org.opensearch.core.tasks.TaskId;
+import org.opensearch.plugin.wlm.spi.AttributeExtractorExtension;
 import org.opensearch.rule.InMemoryRuleProcessingService;
 import org.opensearch.rule.RuleAttribute;
 import org.opensearch.rule.attribute_extractor.AttributeExtractor;
@@ -37,7 +39,9 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.concurrent.atomic.AtomicInteger;
 
+import static org.opensearch.plugin.wlm.WorkloadManagementPlugin.PRINCIPAL_ATTRIBUTE_NAME;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.anyList;
 import static org.mockito.Mockito.doAnswer;
@@ -100,6 +104,181 @@ public void testApplyForInValidRequest() {
         verify(ruleProcessingService, times(0)).evaluateLabel(anyList());
     }
 
+    public void testApplySetsThrottlePrincipalOnTaskWhenExtractorPresent() {
+        // A feature type that includes a "principal" attribute + a matching extractor extension in the map.
+        Attribute principalAttr = new Attribute() {
+            @Override
+            public String getName() {
+                return PRINCIPAL_ATTRIBUTE_NAME;
+            }
+
+            @Override
+            public void validateAttribute() {}
+
+            @Override
+            public void writeTo(StreamOutput out) throws IOException {}
+        };
+        FeatureType featureTypeWithPrincipal = new FeatureType() {
+            @Override
+            public String getName() {
+                return "wlm";
+            }
+
+            @Override
+            public Map<Attribute, Integer> getOrderedAttributes() {
+                return Map.of(principalAttr, 1);
+            }
+        };
+        AttributeExtractor<String> principalExtractor = new AttributeExtractor<>() {
+            @Override
+            public Attribute getAttribute() {
+                return principalAttr;
+            }
+
+            @Override
+            public Iterable<String> extract() {
+                return List.of("username|alice", "role|admin");
+            }
+
+            @Override
+            public LogicalOperator getLogicalOperator() {
+                return LogicalOperator.OR;
+            }
+        };
+        AttributeExtractorExtension extension = () -> principalExtractor;
+        Map<Attribute, AttributeExtractorExtension> extensions = Map.of(principalAttr, extension);
+
+        InMemoryRuleProcessingService svc = spy(
+            new InMemoryRuleProcessingService(
+                new AttributeValueStoreFactory(featureTypeWithPrincipal, DefaultAttributeValueStore::new),
+                null
+            )
+        );
+        AutoTaggingActionFilter filter = new AutoTaggingActionFilter(
+            svc,
+            threadPool,
+            extensions,
+            mock(WlmClusterSettingValuesProvider.class),
+            featureTypeWithPrincipal
+        );
+
+        SearchRequest request = mock(SearchRequest.class);
+        when(request.indices()).thenReturn(new String[] { "foo" });
+        ActionFilterChain<ActionRequest, ActionResponse> chain = mock(TestActionFilterChain.class);
+        WorkloadGroupTask task = newWorkloadGroupTask();
+        try (ThreadContext.StoredContext ctx = threadPool.getThreadContext().stashContext()) {
+            when(svc.evaluateLabel(anyList())).thenReturn(Optional.of("QG"));
+            filter.apply(task, "Test", request, ActionRequestMetadata.empty(), null, chain);
+
+            // Both principal tokens are joined (by WORKLOAD_GROUP_PRINCIPAL_VALUE_DELIMITER) onto the task for
+            // core-side throttling.
+            assertEquals(
+                "username|alice" + WorkloadGroupTask.WORKLOAD_GROUP_PRINCIPAL_VALUE_DELIMITER + "role|admin",
+                task.getThrottlePrincipal()
+            );
+            // The principal must NOT land in the thread context: request headers are serialized onto every outgoing
+            // transport request, which would ship the caller's identity to every shard and to remote clusters.
+            assertNull(threadPool.getThreadContext().getHeader("workloadGroupPrincipal"));
+        }
+    }
+
+    public void testApplyTwiceOnOneThreadContextIsTolerated() {
+        // _msearch dispatches each sub-search through the filter chain on a single ThreadContext with no
+        // stashContext(), and ThreadContext.putHeader throws when the key is already present. The filter must therefore
+        // tolerate running more than once, otherwise every sub-request after the first fails. Carrying the principal on
+        // the task instead of in the thread context is what makes that safe, and gives each sub-request its own value.
+        Attribute principalAttr = new Attribute() {
+            @Override
+            public String getName() {
+                return PRINCIPAL_ATTRIBUTE_NAME;
+            }
+
+            @Override
+            public void validateAttribute() {}
+
+            @Override
+            public void writeTo(StreamOutput out) throws IOException {}
+        };
+        FeatureType featureTypeWithPrincipal = new FeatureType() {
+            @Override
+            public String getName() {
+                return "wlm";
+            }
+
+            @Override
+            public Map<Attribute, Integer> getOrderedAttributes() {
+                return Map.of(principalAttr, 1);
+            }
+        };
+        AtomicInteger extractCalls = new AtomicInteger();
+        AttributeExtractor<String> principalExtractor = new AttributeExtractor<>() {
+            @Override
+            public Attribute getAttribute() {
+                return principalAttr;
+            }
+
+            @Override
+            public Iterable<String> extract() {
+                extractCalls.incrementAndGet();
+                return List.of("username|alice");
+            }
+
+            @Override
+            public LogicalOperator getLogicalOperator() {
+                return LogicalOperator.OR;
+            }
+        };
+        AttributeExtractorExtension extension = () -> principalExtractor;
+        InMemoryRuleProcessingService svc = spy(
+            new InMemoryRuleProcessingService(
+                new AttributeValueStoreFactory(featureTypeWithPrincipal, DefaultAttributeValueStore::new),
+                null
+            )
+        );
+        AutoTaggingActionFilter filter = new AutoTaggingActionFilter(
+            svc,
+            threadPool,
+            Map.of(principalAttr, extension),
+            mock(WlmClusterSettingValuesProvider.class),
+            featureTypeWithPrincipal
+        );
+
+        SearchRequest request = mock(SearchRequest.class);
+        when(request.indices()).thenReturn(new String[] { "foo" });
+        ActionFilterChain<ActionRequest, ActionResponse> chain = mock(TestActionFilterChain.class);
+        WorkloadGroupTask first = newWorkloadGroupTask();
+        WorkloadGroupTask second = newWorkloadGroupTask();
+        try (ThreadContext.StoredContext ctx = threadPool.getThreadContext().stashContext()) {
+            // No label, so the (separate, pre-existing) workload-group-id header is not set and this test isolates the
+            // principal.
+            when(svc.evaluateLabel(anyList())).thenReturn(Optional.empty());
+            filter.apply(first, "Test", request, ActionRequestMetadata.empty(), null, chain);
+            filter.apply(second, "Test", request, ActionRequestMetadata.empty(), null, chain);
+
+            // Each sub-request carries its own principal, and neither run threw on a duplicate key.
+            assertEquals("username|alice", first.getThrottlePrincipal());
+            assertEquals("username|alice", second.getThrottlePrincipal());
+            assertNull(threadPool.getThreadContext().getHeader("workloadGroupPrincipal"));
+            // The principal is materialized once per request and reused for both label evaluation and the task field;
+            // extract() carries no re-iterability contract, so calling it twice per request risks yielding nothing.
+            assertEquals("extract() must be invoked once per request", 2, extractCalls.get());
+        }
+    }
+
+    public void testApplyLeavesThrottlePrincipalUnsetWhenNoExtractor() {
+        // Default filter from setUp has no principal attribute/extractor -> the task's principal stays null, which is
+        // what makes username/role throttling fail open rather than bucket everyone together.
+        SearchRequest request = mock(SearchRequest.class);
+        when(request.indices()).thenReturn(new String[] { "foo" });
+        ActionFilterChain<ActionRequest, ActionResponse> chain = mock(TestActionFilterChain.class);
+        try (ThreadContext.StoredContext ctx = threadPool.getThreadContext().stashContext()) {
+            when(ruleProcessingService.evaluateLabel(anyList())).thenReturn(Optional.of("QG"));
+            WorkloadGroupTask task = newWorkloadGroupTask();
+            autoTaggingActionFilter.apply(task, "Test", request, ActionRequestMetadata.empty(), null, chain);
+            assertNull(task.getThrottlePrincipal());
+        }
+    }
+
     public void testApplyForScrollRequestWithOriginalIndices() {
         SearchScrollRequest request = mock(SearchScrollRequest.class);
         ActionFilterChain<ActionRequest, ActionResponse> chain = mock(TestActionFilterChain.class);
@@ -169,6 +348,10 @@ public void validateAttribute() {}
         public void writeTo(StreamOutput out) throws IOException {}
     }
 
+    private static WorkloadGroupTask newWorkloadGroupTask() {
+        return new WorkloadGroupTask(1L, "transport", "Test", "test task", TaskId.EMPTY_TASK_ID, Map.of());
+    }
+
     private static class TestActionFilterChain implements ActionFilterChain<ActionRequest, ActionResponse> {
         @Override
         public void proceed(Task task, String action, ActionRequest request, ActionListener<ActionResponse> listener) {