Enforce node-level WLM request throttling (node_limit)

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

@@ -117,14 +117,32 @@ public LogicalOperator getLogicalOperator() {
             }
         }
 
+        AttributeExtractor<String> principalExtractor = 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());
+            principalExtractor = attributeExtensions.get(attribute).getAttributeExtractor();
+            attributeExtractors.add(principalExtractor);
         }
 
         Optional<String> label = ruleProcessingService.evaluateLabel(attributeExtractors);
         label.ifPresent(s -> threadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, s));
+        // Propagate the principal so core-side throttling can build per-username / per-role buckets.
+        if (principalExtractor != null) {
+            stashPrincipal(principalExtractor);
+        }
         chain.proceed(task, action, request, listener);
     }
+
+    /**
+     * Stashes the extractor's {@code subfield|value} principal tokens (joined by
+     * {@link WorkloadGroupTask#WORKLOAD_GROUP_PRINCIPAL_VALUE_DELIMITER}) into the thread context for core-side
+     * throttling. No header is set when the extractor yields nothing, leaving username/role throttling to fail open.
+     */
+    private void stashPrincipal(AttributeExtractor<String> principalExtractor) {
+        String principal = String.join(WorkloadGroupTask.WORKLOAD_GROUP_PRINCIPAL_VALUE_DELIMITER, principalExtractor.extract());
+        if (principal.isEmpty() == false) {
+            threadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_PRINCIPAL_HEADER, principal);
+        }
+    }
 }

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

@@ -18,6 +18,7 @@
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.action.ActionResponse;
 import org.opensearch.core.common.io.stream.StreamOutput;
+import org.opensearch.plugin.wlm.spi.AttributeExtractorExtension;
 import org.opensearch.rule.InMemoryRuleProcessingService;
 import org.opensearch.rule.RuleAttribute;
 import org.opensearch.rule.attribute_extractor.AttributeExtractor;
@@ -38,6 +39,7 @@
 import java.util.Map;
 import java.util.Optional;
 
+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 +102,91 @@ public void testApplyForInValidRequest() {
         verify(ruleProcessingService, times(0)).evaluateLabel(anyList());
     }
 
+    public void testApplyStashesPrincipalHeaderWhenExtractorPresent() {
+        // 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);
+        try (ThreadContext.StoredContext ctx = threadPool.getThreadContext().stashContext()) {
+            when(svc.evaluateLabel(anyList())).thenReturn(Optional.of("QG"));
+            filter.apply(mock(Task.class), "Test", request, ActionRequestMetadata.empty(), null, chain);
+
+            // Both principal tokens are joined into the propagated header for core-side throttling.
+            assertEquals(
+                "username|alice,role|admin",
+                threadPool.getThreadContext().getHeader(WorkloadGroupTask.WORKLOAD_GROUP_PRINCIPAL_HEADER)
+            );
+        }
+    }
+
+    public void testApplyDoesNotStashPrincipalHeaderWhenNoExtractor() {
+        // Default filter from setUp has no principal attribute/extractor -> header must be absent.
+        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"));
+            autoTaggingActionFilter.apply(mock(Task.class), "Test", request, ActionRequestMetadata.empty(), null, chain);
+            assertNull(threadPool.getThreadContext().getHeader(WorkloadGroupTask.WORKLOAD_GROUP_PRINCIPAL_HEADER));
+        }
+    }
+
     public void testApplyForScrollRequestWithOriginalIndices() {
         SearchScrollRequest request = mock(SearchScrollRequest.class);
         ActionFilterChain<ActionRequest, ActionResponse> chain = mock(TestActionFilterChain.class);

server/src/main/java/org/opensearch/action/search/StreamTransportSearchAction.java

@@ -30,6 +30,7 @@
 import org.opensearch.transport.StreamTransportService;
 import org.opensearch.transport.Transport;
 import org.opensearch.transport.client.node.NodeClient;
+import org.opensearch.wlm.WorkloadGroupService;
 
 import java.util.Map;
 import java.util.Set;
@@ -59,7 +60,8 @@ public StreamTransportSearchAction(
         SearchRequestOperationsCompositeListenerFactory searchRequestOperationsCompositeListenerFactory,
         Tracer tracer,
         TaskResourceTrackingService taskResourceTrackingService,
-        IndicesService indicesService
+        IndicesService indicesService,
+        WorkloadGroupService workloadGroupService
     ) {
         super(
             client,
@@ -78,7 +80,8 @@ public StreamTransportSearchAction(
             searchRequestOperationsCompositeListenerFactory,
             tracer,
             taskResourceTrackingService,
-            indicesService
+            indicesService,
+            workloadGroupService
         );
     }
 

server/src/main/java/org/opensearch/action/search/TransportSearchAction.java

@@ -57,6 +57,7 @@
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.Nullable;
 import org.opensearch.common.inject.Inject;
+import org.opensearch.common.lease.Releasable;
 import org.opensearch.common.settings.Setting;
 import org.opensearch.common.settings.Setting.Property;
 import org.opensearch.common.unit.TimeValue;
@@ -67,6 +68,7 @@
 import org.opensearch.core.common.breaker.CircuitBreaker;
 import org.opensearch.core.common.io.stream.NamedWriteableRegistry;
 import org.opensearch.core.common.io.stream.Writeable;
+import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException;
 import org.opensearch.core.index.Index;
 import org.opensearch.core.index.shard.ShardId;
 import org.opensearch.core.indices.breaker.CircuitBreakerService;
@@ -107,6 +109,7 @@
 import org.opensearch.transport.client.Client;
 import org.opensearch.transport.client.OriginSettingClient;
 import org.opensearch.transport.client.node.NodeClient;
+import org.opensearch.wlm.WorkloadGroupService;
 import org.opensearch.wlm.WorkloadGroupTask;
 
 import java.util.ArrayList;
@@ -188,6 +191,7 @@ public class TransportSearchAction extends HandledTransportAction<SearchRequest,
     private final MetricsRegistry metricsRegistry;
 
     private TaskResourceTrackingService taskResourceTrackingService;
+    private final WorkloadGroupService workloadGroupService;
 
     @Inject
     public TransportSearchAction(
@@ -207,7 +211,8 @@ public TransportSearchAction(
         SearchRequestOperationsCompositeListenerFactory searchRequestOperationsCompositeListenerFactory,
         Tracer tracer,
         TaskResourceTrackingService taskResourceTrackingService,
-        IndicesService indicesService
+        IndicesService indicesService,
+        WorkloadGroupService workloadGroupService
     ) {
         super(SearchAction.NAME, transportService, actionFilters, (Writeable.Reader<SearchRequest>) SearchRequest::new);
         this.client = client;
@@ -231,6 +236,7 @@ public TransportSearchAction(
         this.tracer = tracer;
         this.taskResourceTrackingService = taskResourceTrackingService;
         this.indicesService = indicesService;
+        this.workloadGroupService = workloadGroupService;
     }
 
     private Map<String, AliasFilter> buildPerIndexAliasFilter(
@@ -463,7 +469,7 @@ void executeRequest(
         final Span requestSpan = tracer.startSpan(SpanBuilder.from(task, actionName));
         try (final SpanScope spanScope = tracer.withSpanInScope(requestSpan)) {
             SearchRequestOperationsListener.CompositeListener requestOperationsListeners;
-            final ActionListener<SearchResponse> updatedListener = TraceableActionListener.create(originalListener, requestSpan, tracer);
+            ActionListener<SearchResponse> updatedListener = TraceableActionListener.create(originalListener, requestSpan, tracer);
             requestOperationsListeners = searchRequestOperationsCompositeListenerFactory.buildCompositeListener(
                 originalSearchRequest,
                 logger,
@@ -474,14 +480,31 @@ void executeRequest(
                 originalSearchRequest,
                 taskResourceTrackingService::getTaskResourceUsageFromThreadContext
             );
-            searchRequestContext.getSearchRequestOperationsListener().onRequestStart(searchRequestContext);
 
             // At this point either the QUERY_GROUP_ID header will be present in ThreadContext either via ActionFilter
             // or HTTP header (HTTP header will be deprecated once ActionFilter is implemented)
             if (task instanceof WorkloadGroupTask) {
                 ((WorkloadGroupTask) task).setWorkloadGroupId(threadPool.getThreadContext());
+                // Node-level throttle admission. Runs before onRequestStart so a rejection doesn't leave the
+                // request-operations gauges incremented (they're only decremented on request end/failure, which the
+                // early return below skips). Principal header is null unless the security plugin's extractor set it.
+                try {
+                    String principal = threadPool.getThreadContext().getHeader(WorkloadGroupTask.WORKLOAD_GROUP_PRINCIPAL_HEADER);
+                    Releasable throttlePermit = workloadGroupService.acquireThrottleOrReject(
+                        ((WorkloadGroupTask) task).getWorkloadGroupId(),
+                        principal
+                    );
+                    if (throttlePermit != null) {
+                        updatedListener = ActionListener.runBefore(updatedListener, throttlePermit::close);
+                    }
+                } catch (OpenSearchRejectedExecutionException e) {
+                    updatedListener.onFailure(e);
+                    return;
+                }
             }
 
+            searchRequestContext.getSearchRequestOperationsListener().onRequestStart(searchRequestContext);
+
             PipelinedRequest searchRequest;
             ActionListener<SearchResponse> listener;
             try {
@@ -512,7 +535,12 @@ void executeRequest(
                     );
                 }
             }, listener::onFailure);
-            searchRequest.transformRequest(requestTransformListener);
+            try {
+                searchRequest.transformRequest(requestTransformListener);
+            } catch (Exception e) {
+                // Route through updatedListener so a chained throttle permit is released rather than leaked.
+                updatedListener.onFailure(e);
+            }
         }
     }
 

server/src/main/java/org/opensearch/wlm/WorkloadGroupService.java

@@ -16,7 +16,9 @@
 import org.opensearch.cluster.metadata.Metadata;
 import org.opensearch.cluster.metadata.WorkloadGroup;
 import org.opensearch.cluster.service.ClusterService;
+import org.opensearch.common.lease.Releasable;
 import org.opensearch.common.lifecycle.AbstractLifecycleComponent;
+import org.opensearch.common.settings.Settings;
 import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException;
 import org.opensearch.monitor.jvm.JvmStats;
 import org.opensearch.monitor.process.ProcessProbe;
@@ -59,6 +61,8 @@ public class WorkloadGroupService extends AbstractLifecycleComponent
     private final Set<WorkloadGroup> deletedWorkloadGroups;
     private final NodeDuressTrackers nodeDuressTrackers;
     private final WorkloadGroupsStateAccessor workloadGroupsStateAccessor;
+    // Node-local in-flight throttle counters, keyed by throttle bucket. No cross-node coordination in this tier.
+    private final WorkloadGroupThrottleTracker throttleTracker = new WorkloadGroupThrottleTracker();
 
     public WorkloadGroupService(
         WorkloadGroupTaskCancellationService taskCancellationService,
@@ -312,6 +316,82 @@ public void rejectIfNeeded(String workloadGroupId) {
         });
     }
 
+    /**
+     * Acquires one node-level throttle permit for the request, or returns {@code null} (nothing to release) when the
+     * request is not throttled: WLM disabled, default/unknown group, no {@code node_limit}, or no resolvable bucket
+     * (see {@link #resolveThrottleBucketKey}). The bucket depends on the group's throttle {@code attribute}.
+     *
+     * @param workloadGroupId the workload group the request is assigned to
+     * @param principal       the raw {@code WORKLOAD_GROUP_PRINCIPAL_HEADER} value, or {@code null} (see resolver)
+     * @return a permit to close on request completion, or {@code null} if not throttled
+     * @throws OpenSearchRejectedExecutionException if the bucket is already at its node limit
+     */
+    public Releasable acquireThrottleOrReject(String workloadGroupId, String principal) {
+        if (workloadManagementSettings.getWlmMode() != WlmMode.ENABLED) {
+            return null;
+        }
+        if (workloadGroupId == null || workloadGroupId.equals(WorkloadGroupTask.DEFAULT_WORKLOAD_GROUP_ID_SUPPLIER.get())) {
+            return null;
+        }
+        try {
+            WorkloadGroup workloadGroup = getWorkloadGroupById(workloadGroupId);
+            if (workloadGroup == null) {
+                return null;
+            }
+            Settings throttling = workloadGroup.getMutableWorkloadGroupFragment().getThrottling();
+            int nodeLimit = WorkloadGroupThrottleSettings.NODE_LIMIT.get(throttling);
+            if (nodeLimit == WorkloadGroupThrottleSettings.UNSET_LIMIT) {
+                return null;
+            }
+            // A null bucket means the request can't be attributed (e.g. username/role with no principal) -> fail open.
+            String bucketKey = resolveThrottleBucketKey(
+                workloadGroupId,
+                WorkloadGroupThrottleSettings.ATTRIBUTE.get(throttling),
+                principal
+            );
+            if (bucketKey == null) {
+                return null;
+            }
+            try {
+                return throttleTracker.acquire(bucketKey, nodeLimit);
+            } catch (OpenSearchRejectedExecutionException e) {
+                workloadGroupsStateAccessor.getWorkloadGroupState(workloadGroupId).totalThrottled.inc();
+                throw e;
+            }
+        } catch (OpenSearchRejectedExecutionException e) {
+            throw e; // the intended 429
+        } catch (Exception e) {
+            // A bug in the throttle path must never fail an otherwise-valid search, so fail open.
+            logger.warn("Skipping node-level throttle for workload group [" + workloadGroupId + "] due to an error", e);
+            return null;
+        }
+    }
+
+    /**
+     * Builds the throttle bucket key: {@code <groupId>:group} for whole-group throttling, or
+     * {@code <groupId>:<attribute>:<value>} keyed by the matching principal subfield for {@code username}/{@code role}.
+     * Returns {@code null} (fail open, not throttled) when the principal is absent or lacks a token for the subfield.
+     */
+    private String resolveThrottleBucketKey(String workloadGroupId, String attribute, String principal) {
+        if ("group".equals(attribute)) {
+            return workloadGroupId + ":group";
+        }
+        if (principal == null || principal.isEmpty()) {
+            return null;
+        }
+        String subfieldPrefix = attribute + "|";
+        for (String token : principal.split(WorkloadGroupTask.WORKLOAD_GROUP_PRINCIPAL_VALUE_DELIMITER)) {
+            String trimmed = token.trim();
+            if (trimmed.startsWith(subfieldPrefix)) {
+                String value = trimmed.substring(subfieldPrefix.length());
+                if (value.isEmpty() == false) {
+                    return workloadGroupId + ":" + attribute + ":" + value;
+                }
+            }
+        }
+        return null;
+    }
+
     private double getNormalisedRejectionThreshold(double limit, ResourceType resourceType) {
         if (resourceType == ResourceType.CPU) {
             return limit * workloadManagementSettings.getNodeLevelCpuRejectionThreshold();

server/src/main/java/org/opensearch/wlm/WorkloadGroupTask.java

@@ -30,6 +30,13 @@ public class WorkloadGroupTask extends CancellableTask {
 
     private static final Logger logger = LogManager.getLogger(WorkloadGroupTask.class);
     public static final String WORKLOAD_GROUP_ID_HEADER = "workloadGroupId";
+    /**
+     * Carries the request's principal for username/role throttling: {@code subfield|value} tokens
+     * (e.g. {@code username|alice}) joined by {@link #WORKLOAD_GROUP_PRINCIPAL_VALUE_DELIMITER}. Only set when the
+     * security plugin's principal extractor is installed; absent otherwise (username/role throttling then fails open).
+     */
+    public static final String WORKLOAD_GROUP_PRINCIPAL_HEADER = "workloadGroupPrincipal";
+    public static final String WORKLOAD_GROUP_PRINCIPAL_VALUE_DELIMITER = ",";
     public static final Supplier<String> DEFAULT_WORKLOAD_GROUP_ID_SUPPLIER = () -> "DEFAULT_WORKLOAD_GROUP";
     private final LongSupplier nanoTimeSupplier;
     private String workloadGroupId;