[WLM] Introduce rule cardinality check (opensearch-project#18663)

Signed-off-by: Kaushal Kumar <ravi.kaushal97@gmail.com>

Author: kaushalmahi12

CHANGELOG.md

@@ -26,6 +26,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Add Semantic Version field type mapper and extensive unit tests([#18454](https://github.com/opensearch-project/OpenSearch/pull/18454))
 - Pass index settings to system ingest processor factories. ([#18708](https://github.com/opensearch-project/OpenSearch/pull/18708))
 - Include named queries from rescore contexts in matched_queries array ([#18697](https://github.com/opensearch-project/OpenSearch/pull/18697))
+- Add the configurable limit on rule cardinality ([#18663](https://github.com/opensearch-project/OpenSearch/pull/18663))
 
 ### Changed
 - Update Subject interface to use CheckedRunnable ([#18570](https://github.com/opensearch-project/OpenSearch/issues/18570))

modules/autotagging-commons/common/src/main/java/org/opensearch/rule/RuleQueryMapper.java

@@ -24,4 +24,11 @@ public interface RuleQueryMapper<T> {
      * @return
      */
     T from(GetRuleRequest request);
+
+    /**
+     * This method returns the cardinality query for the rule, this query should
+     * be constructed in such a way that it can be used to calculate the cardinality of the rules
+     * @return
+     */
+    T getCardinalityQuery();
 }

modules/autotagging-commons/common/src/main/java/org/opensearch/rule/service/IndexStoredRulePersistenceService.java

@@ -19,9 +19,11 @@
 import org.opensearch.action.support.clustermanager.AcknowledgedResponse;
 import org.opensearch.action.update.UpdateRequest;
 import org.opensearch.cluster.service.ClusterService;
+import org.opensearch.common.settings.Setting;
 import org.opensearch.common.util.concurrent.ThreadContext;
 import org.opensearch.common.xcontent.XContentFactory;
 import org.opensearch.core.action.ActionListener;
+import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException;
 import org.opensearch.core.xcontent.ToXContent;
 import org.opensearch.index.engine.DocumentMissingException;
 import org.opensearch.index.query.QueryBuilder;
@@ -53,8 +55,28 @@
  */
 public class IndexStoredRulePersistenceService implements RulePersistenceService {
     /**
-     * The system index name used for storing rules
+     * Default value for max rules count
      */
+    public static final int DEFAULT_MAX_ALLOWED_RULE_COUNT = 200;
+
+    /**
+     *  max wlm rules setting name
+     */
+    public static final String MAX_RULES_COUNT_SETTING_NAME = "wlm.autotagging.max_rules";
+
+    /**
+     *  max wlm rules setting
+     */
+    public static final Setting<Integer> MAX_WLM_RULES_SETTING = Setting.intSetting(
+        MAX_RULES_COUNT_SETTING_NAME,
+        DEFAULT_MAX_ALLOWED_RULE_COUNT,
+        10,
+        500,
+        Setting.Property.NodeScope,
+        Setting.Property.Dynamic
+    );
+
+    private int maxAllowedRulesCount;
     private final String indexName;
     private final Client client;
     private final ClusterService clusterService;
@@ -87,6 +109,12 @@ public IndexStoredRulePersistenceService(
         this.maxRulesPerPage = maxRulesPerPage;
         this.parser = parser;
         this.queryBuilder = queryBuilder;
+        this.maxAllowedRulesCount = MAX_WLM_RULES_SETTING.get(clusterService.getSettings());
+        clusterService.getClusterSettings().addSettingsUpdateConsumer(MAX_WLM_RULES_SETTING, this::setMaxAllowedRules);
+    }
+
+    private void setMaxAllowedRules(int maxAllowedRules) {
+        this.maxAllowedRulesCount = maxAllowedRules;
     }
 
     /**
@@ -101,12 +129,27 @@ public void createRule(CreateRuleRequest request, ActionListener<CreateRuleRespo
                 logger.error("Index {} does not exist", indexName);
                 listener.onFailure(new IllegalStateException("Index" + indexName + " does not exist"));
             } else {
+                performCardinalityCheck(listener);
                 Rule rule = request.getRule();
                 validateNoDuplicateRule(rule, ActionListener.wrap(unused -> persistRule(rule, listener), listener::onFailure));
             }
         }
     }
 
+    private void performCardinalityCheck(ActionListener<CreateRuleResponse> listener) {
+        SearchResponse searchResponse = client.prepareSearch(indexName).setQuery(queryBuilder.getCardinalityQuery()).get();
+        if (searchResponse.getHits().getTotalHits() != null && searchResponse.getHits().getTotalHits().value() >= maxAllowedRulesCount) {
+            listener.onFailure(
+                new OpenSearchRejectedExecutionException(
+                    "This create operation will violate"
+                        + " the cardinality limit of "
+                        + DEFAULT_MAX_ALLOWED_RULE_COUNT
+                        + ". Please delete some stale or redundant rules first"
+                )
+            );
+        }
+    }
+
     /**
      * Validates that no existing rule has the same attribute map as the given rule.
      * This validation must be performed one at a time to prevent writing duplicate rules.

modules/autotagging-commons/common/src/main/java/org/opensearch/rule/storage/IndexBasedRuleQueryMapper.java

@@ -53,4 +53,9 @@ public QueryBuilder from(GetRuleRequest request) {
         }
         return boolQuery;
     }
+
+    @Override
+    public QueryBuilder getCardinalityQuery() {
+        return QueryBuilders.matchAllQuery();
+    }
 }

modules/autotagging-commons/common/src/test/java/org/opensearch/rule/service/IndexStoredRulePersistenceServiceTests.java

@@ -21,10 +21,12 @@
 import org.opensearch.cluster.metadata.Metadata;
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.action.ActionFuture;
+import org.opensearch.common.settings.ClusterSettings;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.common.util.concurrent.ThreadContext;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.common.bytes.BytesArray;
+import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException;
 import org.opensearch.core.index.shard.ShardId;
 import org.opensearch.index.engine.DocumentMissingException;
 import org.opensearch.index.query.QueryBuilder;
@@ -50,6 +52,7 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -99,6 +102,11 @@ public void setUp() throws Exception {
         client = setUpMockClient(searchRequestBuilder);
         rule = mock(Rule.class);
         clusterService = mock(ClusterService.class);
+        Settings testSettings = Settings.EMPTY;
+        ClusterSettings clusterSettings = new ClusterSettings(testSettings, new HashSet<>());
+        when(clusterService.getSettings()).thenReturn(testSettings);
+        clusterSettings.registerSetting(IndexStoredRulePersistenceService.MAX_WLM_RULES_SETTING);
+        when(clusterService.getClusterSettings()).thenReturn(clusterSettings);
         ClusterState clusterState = mock(ClusterState.class);
         Metadata metadata = mock(Metadata.class);
         when(clusterService.state()).thenReturn(clusterState);
@@ -109,6 +117,7 @@ public void setUp() throws Exception {
         queryBuilder = mock(QueryBuilder.class);
         when(queryBuilder.filter(any())).thenReturn(queryBuilder);
         when(ruleQueryMapper.from(any(GetRuleRequest.class))).thenReturn(queryBuilder);
+        when(ruleQueryMapper.getCardinalityQuery()).thenReturn(mock(QueryBuilder.class));
         when(ruleEntityParser.parse(anyString())).thenReturn(rule);
 
         rulePersistenceService = new IndexStoredRulePersistenceService(
@@ -144,6 +153,25 @@ public void testCreateRuleOnExistingIndex() throws Exception {
         assertNotNull(responseCaptor.getValue().getRule());
     }
 
+    public void testCardinalityCheckBasedFailure() throws Exception {
+        CreateRuleRequest createRuleRequest = mock(CreateRuleRequest.class);
+        when(createRuleRequest.getRule()).thenReturn(rule);
+        when(rule.toXContent(any(), any())).thenAnswer(invocation -> invocation.getArgument(0));
+
+        SearchResponse searchResponse = mock(SearchResponse.class);
+        when(searchResponse.getHits()).thenReturn(
+            new SearchHits(new SearchHit[] {}, new TotalHits(10000, TotalHits.Relation.EQUAL_TO), 1.0f)
+        );
+        when(searchRequestBuilder.get()).thenReturn(searchResponse);
+
+        ActionListener<CreateRuleResponse> listener = mock(ActionListener.class);
+        rulePersistenceService.createRule(createRuleRequest, listener);
+
+        ArgumentCaptor<Exception> exceptionCaptor = ArgumentCaptor.forClass(OpenSearchRejectedExecutionException.class);
+        verify(listener).onFailure(exceptionCaptor.capture());
+        assertNotNull(exceptionCaptor.getValue());
+    }
+
     public void testConcurrentCreateDuplicateRules() throws InterruptedException {
         ExecutorService singleThreadExecutor = Executors.newSingleThreadExecutor();
         int threadCount = 10;

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

@@ -75,6 +75,8 @@
 import java.util.Map;
 import java.util.function.Supplier;
 
+import static org.opensearch.rule.service.IndexStoredRulePersistenceService.MAX_WLM_RULES_SETTING;
+
 /**
  * Plugin class for WorkloadManagement
  */
@@ -190,7 +192,11 @@ public List<RestHandler> getRestHandlers(
 
     @Override
     public List<Setting<?>> getSettings() {
-        return List.of(WorkloadGroupPersistenceService.MAX_QUERY_GROUP_COUNT, RefreshBasedSyncMechanism.RULE_SYNC_REFRESH_INTERVAL_SETTING);
+        return List.of(
+            WorkloadGroupPersistenceService.MAX_QUERY_GROUP_COUNT,
+            RefreshBasedSyncMechanism.RULE_SYNC_REFRESH_INTERVAL_SETTING,
+            MAX_WLM_RULES_SETTING
+        );
     }
 
     @Override