add trie to store Rules in memory

Signed-off-by: Kaushal Kumar <[email protected]>

plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/structure/FastPrefixMatchingStructure.java

@@ -8,6 +8,12 @@
 
 package org.opensearch.plugin.wlm.rule.structure;
 
+import java.util.List;
+
 public interface FastPrefixMatchingStructure {
-    void add(String s);
+    void insert(String key, String value);
+
+    List<String> search(String key);
+
+    boolean delete(String key);
 }

plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/structure/RuleAttributeTrie.java

@@ -8,47 +8,59 @@
 
 package org.opensearch.plugin.wlm.rule.structure;
 
-import java.util.HashMap;
-import java.util.Map;
-import java.util.regex.Pattern;
+import java.util.List;
 
 public class RuleAttributeTrie implements FastPrefixMatchingStructure {
-    private static final Pattern ALLOWED_ATTRIBUTE_VALUES = Pattern.compile("^[a-zA-Z0-9-_]+\\*?$");
-    @Override
-    public void add(String s) {
+    private static final String ALLOWED_ATTRIBUTE_VALUES = "^[a-zA-Z0-9-_]+\\*?$";
+    private TrieNode root;
 
+    /**
+     * Constructs an empty AugmentedTrie.
+     */
+    public RuleAttributeTrie() {
+        root = new TrieNode("");
     }
 
-    public enum RuleAttributeName {
-        USERNAME("username"),
-        INDEX_PATTERN("index_pattern");
-        private final String name;
-
-        RuleAttributeName(String name) {
-            this.name = name;
-        }
-
-        public String getName() { return name; }
-
-        public static RuleAttributeName fromString(String name) {
-            for (RuleAttributeName attributeName : RuleAttributeName.values()) {
-                if (attributeName.getName().equals(name)) {
-                    return attributeName;
-                }
-            }
-            throw new IllegalArgumentException("Invalid rule attribute name [" + name + "]");
+    /**
+     * Inserts a key-value pair into the trie.
+     *
+     * @param key   The key to be inserted.
+     * @param value The value associated with the key.
+     */
+    public void insert(String key, String value) {
+        if (!IsvalidValue(value)) {
+            throw new IllegalArgumentException(
+                "Invalid attribute value: " + value + " it should match the regex " + ALLOWED_ATTRIBUTE_VALUES
+            );
         }
-
+        TrieInserter inserter = new TrieInserter(root, key, value);
+        root = inserter.insert();
     }
 
-    public static class LabeledNode {
-        private String label;
-        private Map<String, LabeledNode> children = new HashMap<>();
+    private boolean IsvalidValue(String value) {
+        return ALLOWED_ATTRIBUTE_VALUES.matches(value);
+    }
 
-        public LabeledNode() {}
+    /**
+     * Searches for a key in the trie.
+     *
+     * @param key The key to search for.
+     * @return A list of string values associated with the key or its prefixes.
+     *         Returns an empty list if no matches are found.
+     */
+    public List<String> search(String key) {
+        TrieSearcher searcher = new TrieSearcher(root, key);
+        return searcher.search();
+    }
 
-        public LabeledNode(String label) {
-            this.label = label;
-        }
+    /**
+     * Deletes a key from the trie.
+     *
+     * @param key The key to be deleted.
+     * @return true if the key was successfully deleted, false otherwise.
+     */
+    public boolean delete(String key) {
+        TrieDeleter deleter = new TrieDeleter(root, key);
+        return deleter.delete();
     }
 }

plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/structure/TrieDeleter.java

@@ -0,0 +1,66 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.plugin.wlm.rule.structure;
+
+/**
+ * Handles the deletion operation for the Augmented Trie.
+ */
+class TrieDeleter {
+    private TrieNode root;
+    private String key;
+
+    /**
+     * Constructs a TrieDeleter with the given root and key.
+     *
+     * @param root The root node of the trie.
+     * @param key  The key to be deleted.
+     */
+    public TrieDeleter(TrieNode root, String key) {
+        this.root = root;
+        this.key = key;
+    }
+
+    /**
+     * Performs the deletion operation.
+     *
+     * @return true if the key was successfully deleted, false otherwise.
+     */
+    public boolean delete() {
+        TrieNode current = root;
+        TrieNode parent = null;
+        String remainingKey = key;
+        while (!remainingKey.isEmpty()) {
+            TrieNode childNode = current.findCommonPrefixChild(remainingKey);
+
+            if (childNode == null) {
+                return false;
+            }
+            parent = current;
+            current = childNode;
+            remainingKey = remainingKey.substring(childNode.getKey().length());
+        }
+        final boolean deleted = current.isEndOfWord();
+
+        if (deleted) {
+            current.setEndOfWord(false);
+            current.setValue(null);
+            if (current.getChildren().isEmpty()) {
+                deleteLeafNode(parent, current);
+            }
+        }
+
+        return deleted;
+    }
+
+    private static void deleteLeafNode(TrieNode parent, TrieNode current) {
+        if (parent != null) {
+            parent.getChildren().remove(current.getKey());
+        }
+    }
+}

plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/structure/TrieInserter.java

@@ -0,0 +1,90 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.plugin.wlm.rule.structure;
+
+/**
+ * Handles the insertion operation for the Augmented Trie.
+ */
+class TrieInserter {
+    private TrieNode root;
+    private String key;
+    private String value;
+
+    /**
+     * Constructs a TrieInserter with the given root, key, and value.
+     *
+     * @param root  The root node of the trie.
+     * @param key   The key to be inserted.
+     * @param value The value associated with the key.
+     */
+    public TrieInserter(TrieNode root, String key, String value) {
+        this.root = root;
+        this.key = key;
+        this.value = value;
+    }
+
+    /**
+     * Performs the insertion operation.
+     * <ol>Method should handle 3 cases
+     * <li>Simple addition of new child </li>
+     * <li>insert splits a node</li>
+     * <li>inserted key is a prefix to existing key|s, this could either mark a node as endOfWord or it could also split the node</li>
+     * </ol>
+     * @return The root node of the trie after insertion.
+     */
+    public TrieNode insert() {
+        TrieNode current = root;
+        String remainingKey = key;
+        while (!remainingKey.isEmpty()) {
+            TrieNode child = current.findCommonPrefixChild(remainingKey);
+
+            if (child == null) {
+                boolean partialMatch = false;
+                // partial match
+                for (String childKey : current.getChildren().keySet()) {
+                    int commonPrefixLength = getLongestCommonPrefixLength(childKey, remainingKey);
+                    if (commonPrefixLength > 0) {
+                        TrieNode newNode = current.splitNode(childKey, commonPrefixLength);
+
+                        remainingKey = remainingKey.substring(commonPrefixLength);
+
+                        current = newNode;
+                        partialMatch = true;
+                        break;
+                    }
+                }
+                // no match
+                if (!partialMatch) {
+                    current = current.addNewChild(remainingKey);
+                    remainingKey = "";
+                }
+            } else {
+                current = child;
+                remainingKey = remainingKey.substring(child.getKey().length());
+            }
+        }
+        updateNodeValue(current);
+        return root;
+    }
+
+    private void updateNodeValue(TrieNode node) {
+        node.setValue(value);
+        node.setEndOfWord(true);
+    }
+
+    private int getLongestCommonPrefixLength(String str1, String str2) {
+        int minLength = Math.min(str1.length(), str2.length());
+        for (int i = 0; i < minLength; i++) {
+            if (str1.charAt(i) != str2.charAt(i)) {
+                return i;
+            }
+        }
+        return minLength;
+    }
+}

plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/rule/structure/TrieNode.java

@@ -0,0 +1,118 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.plugin.wlm.rule.structure;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Map;
+import java.util.Queue;
+
+/**
+ * Represents a node in the Augmented Trie.
+ * Each node contains a key, an optional value, and references to child nodes.
+ */
+class TrieNode {
+    public static final int CLOSEST_LIMIT = 5;
+    private Map<String, TrieNode> children;
+    private String key;
+    private String value;
+    private boolean isEndOfWord;
+
+    /**
+     * Constructs a TrieNode with the given key.
+     *
+     * @param key The key associated with this node.
+     */
+    public TrieNode(String key) {
+        this.children = new HashMap<>();
+        this.key = key;
+        this.value = null;
+        this.isEndOfWord = false;
+    }
+
+    // Getters and setters
+    public Map<String, TrieNode> getChildren() {
+        return children;
+    }
+
+    public String getKey() {
+        return key;
+    }
+
+    public void setKey(String key) {
+        this.key = key;
+    }
+
+    public String getValue() {
+        return value;
+    }
+
+    public void setValue(String value) {
+        this.value = value;
+    }
+
+    public boolean isEndOfWord() {
+        return isEndOfWord;
+    }
+
+    public void setEndOfWord(boolean endOfWord) {
+        isEndOfWord = endOfWord;
+    }
+
+    public TrieNode addNewChild(String key) {
+        TrieNode newNode = new TrieNode(key);
+        newNode.setValue(value);
+        newNode.setEndOfWord(true);
+        getChildren().put(key, newNode);
+        return newNode;
+    }
+
+    public TrieNode splitNode(String childKey, int commonPrefixLength) {
+        String commonPrefix = childKey.substring(0, commonPrefixLength);
+        TrieNode newNode = new TrieNode(commonPrefix);
+        TrieNode childNode = getChildren().get(childKey);
+
+        // remove the existing partially matching child node since we will split that
+        getChildren().remove(childKey);
+        // re-attach common prefix as direct child
+        getChildren().put(commonPrefix, newNode);
+
+        childNode.setKey(childKey.substring(commonPrefixLength));
+
+        newNode.getChildren().put(childKey.substring(commonPrefixLength), childNode);
+        return newNode;
+    }
+
+    public TrieNode findCommonPrefixChild(String key) {
+        return getChildren().entrySet()
+            .stream()
+            .filter(entry -> key.startsWith(entry.getKey()))
+            .findFirst()
+            .map(Map.Entry::getValue)
+            .orElse(null);
+    }
+
+    public List<String> findTopFiveClosest() {
+        List<String> ans = new ArrayList<>(CLOSEST_LIMIT);
+        Queue<TrieNode> queue = new LinkedList<>();
+        queue.offer(this);
+
+        while (!queue.isEmpty() && ans.size() < CLOSEST_LIMIT) {
+            TrieNode current = queue.poll();
+            if (current.isEndOfWord()) {
+                ans.add(current.getValue());
+            }
+            queue.addAll(current.getChildren().values());
+        }
+
+        return ans;
+    }
+}