Traverse PathTrie instead.

Signed-off-by: dblock <[email protected]>
opensearch-project · Jan 11, 2024 · f44b777 · f44b777
1 parent 5cae02a
commit f44b777
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 20 deletions.
diff --git a/server/src/main/java/org/opensearch/common/path/PathTrie.java b/server/src/main/java/org/opensearch/common/path/PathTrie.java
@@ -36,6 +36,7 @@
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.Map;
+import java.util.Stack;
 import java.util.NoSuchElementException;
 import java.util.function.BiFunction;
 import java.util.function.Supplier;
@@ -405,4 +406,28 @@ public T next() {
             }
         };
     }
+
+    public Iterator<T> retrieveAll() {
+        Stack<TrieNode> stack = new Stack<>();
+        stack.add(root);
+
+        return new Iterator<T>() {
+            @Override
+            public boolean hasNext() {
+                return !stack.empty();
+            }
+
+            @Override
+            public T next() {
+                while (!stack.empty()) {
+                    TrieNode node = stack.pop();
+                    stack.addAll(node.children.values());
+                    if (node.value != null) {
+                        return node.value;
+                    }
+                }
+                throw new NoSuchElementException("encountered a terminal node without a value");
+            }
+        };
+    }
 }
diff --git a/server/src/main/java/org/opensearch/rest/MethodHandlers.java b/server/src/main/java/org/opensearch/rest/MethodHandlers.java
@@ -48,6 +48,13 @@ final class MethodHandlers {
     private final String path;
     private final Map<RestRequest.Method, RestHandler> methodHandlers;
 
+    /**
+     * Complete path of the method handlers.
+     */
+    public String getPath() {
+        return path;
+    }
+
     MethodHandlers(String path, RestHandler handler, RestRequest.Method... methods) {
         this.path = path;
         this.methodHandlers = new HashMap<>(methods.length);

diff --git a/server/src/main/java/org/opensearch/rest/RestController.java b/server/src/main/java/org/opensearch/rest/RestController.java
@@ -66,7 +66,12 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.net.URI;
-import java.util.*;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Supplier;
 import java.util.function.UnaryOperator;
@@ -104,11 +109,6 @@ public class RestController implements HttpServerTransport.Dispatcher, ToXConten
     }
 
     private final PathTrie<MethodHandlers> handlers = new PathTrie<>(RestUtils.REST_DECODER);
-    private final HashMap<String, Set<RestRequest.Method>> apis = new HashMap<>();
-
-    public HashMap<String, Set<RestRequest.Method>> getApis() {
-        return apis;
-    }
 
     private final UnaryOperator<RestHandler> handlerWrapper;
 
@@ -220,15 +220,6 @@ protected void registerHandler(RestRequest.Method method, String path, RestHandl
     }
 
     private void registerHandlerNoWrap(RestRequest.Method method, String path, RestHandler maybeWrappedHandler) {
-
-        // TODO: traverse pathtrie
-        Set<RestRequest.Method> methods = apis.getOrDefault(path, null);
-        if (methods == null) {
-            methods = new HashSet<RestRequest.Method>();
-            apis.put(path, methods);
-        }
-        methods.add(method);
-
         handlers.insertOrUpdate(
             path,
             new MethodHandlers(path, maybeWrappedHandler, method),
@@ -585,11 +576,12 @@ private Set<RestRequest.Method> getValidHandlerMethodSet(String rawPath) {
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
         builder.startObject("paths");
 
-        for (Map.Entry<String, Set<RestRequest.Method>> entry : getApis().entrySet()) {
-            String key = entry.getKey();
-            Set<RestRequest.Method> methods = entry.getValue();
-            builder.startObject(key.replace("{", ":").replace("}", ""));
-            for(RestRequest.Method method : methods) {
+        Iterator<MethodHandlers> all = handlers.retrieveAll();
+
+        while(all.hasNext()) {
+            MethodHandlers handlers = all.next();
+            builder.startObject(handlers.getPath().replace("{", ":").replace("}", ""));
+            for (RestRequest.Method method : handlers.getValidMethods()) {
                 builder
                     .startObject(method.name().toLowerCase())
                     //.field("summary", "")