Expose RestController all handlers iterator. (#11876)

dblock · dblock · commit 3ba8abee6129 · 2024-01-18T17:20:19.000-05:00
* Expose RestController all handlers iterator. Signed-off-by: dblock <dblock@amazon.com> * Extract MethodHandlers interface. Signed-off-by: dblock <dblock@amazon.com> * Added CHANGELOG. Signed-off-by: dblock <dblock@amazon.com> --------- Signed-off-by: dblock <dblock@amazon.com> (cherry picked from commit 6ccb46b)
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -43,6 +43,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Add copy ingest processor ([#11870](https://github.com/opensearch-project/OpenSearch/pull/11870))
 - Introduce new feature flag "WRITEABLE_REMOTE_INDEX" to gate the writeable remote index functionality ([#11717](https://github.com/opensearch-project/OpenSearch/pull/11170))
 - Bump OpenTelemetry from 1.32.0 to 1.34.1 ([#11891](https://github.com/opensearch-project/OpenSearch/pull/11891))
+- Add `org.opensearch.rest.MethodHandlers` and `RestController#getAllHandlers` ([11876](https://github.com/opensearch-project/OpenSearch/pull/11876))
 
 ### Dependencies
 - Bumps jetty version to 9.4.52.v20230823 to fix GMS-2023-1857 ([#9822](https://github.com/opensearch-project/OpenSearch/pull/9822))
diff --git a/server/src/main/java/org/opensearch/common/path/PathTrie.java b/server/src/main/java/org/opensearch/common/path/PathTrie.java
@@ -37,6 +37,7 @@
 import java.util.Iterator;
 import java.util.Map;
 import java.util.NoSuchElementException;
+import java.util.Stack;
 import java.util.function.BiFunction;
 import java.util.function.Supplier;
 
@@ -405,4 +406,45 @@ public T next() {
             }
         };
     }
+
+    public Iterator<T> retrieveAll() {
+        Stack<TrieNode> stack = new Stack<>();
+        stack.add(root);
+
+        return new Iterator<T>() {
+            @Override
+            public boolean hasNext() {
+                while (!stack.empty()) {
+                    TrieNode node = stack.peek();
+
+                    if (node.value != null) {
+                        return true;
+                    }
+
+                    advance();
+                }
+
+                return false;
+            }
+
+            @Override
+            public T next() {
+                while (!stack.empty()) {
+                    TrieNode node = advance();
+
+                    if (node.value != null) {
+                        return node.value;
+                    }
+                }
+
+                throw new NoSuchElementException("called next() without validating hasNext()! no more nodes available");
+            }
+
+            private TrieNode advance() {
+                TrieNode node = stack.pop();
+                stack.addAll(node.children.values());
+                return node;
+            }
+        };
+    }
 }
diff --git a/server/src/main/java/org/opensearch/rest/MethodHandlers.java b/server/src/main/java/org/opensearch/rest/MethodHandlers.java
@@ -6,82 +6,24 @@
  * compatible open source license.
  */
 
-/*
- * Licensed to Elasticsearch under one or more contributor
- * license agreements. See the NOTICE file distributed with
- * this work for additional information regarding copyright
- * ownership. Elasticsearch licenses this file to you under
- * the Apache License, Version 2.0 (the "License"); you may
- * not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied.  See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-/*
- * Modifications Copyright OpenSearch Contributors. See
- * GitHub history for details.
- */
-
 package org.opensearch.rest;
 
-import org.opensearch.common.Nullable;
+import org.opensearch.common.annotation.PublicApi;
 
-import java.util.HashMap;
-import java.util.Map;
 import java.util.Set;
 
 /**
- * Encapsulate multiple handlers for the same path, allowing different handlers for different HTTP verbs.
- *
- * @opensearch.api
+ * A collection of REST method handlers.
  */
-final class MethodHandlers {
-
-    private final String path;
-    private final Map<RestRequest.Method, RestHandler> methodHandlers;
-
-    MethodHandlers(String path, RestHandler handler, RestRequest.Method... methods) {
-        this.path = path;
-        this.methodHandlers = new HashMap<>(methods.length);
-        for (RestRequest.Method method : methods) {
-            methodHandlers.put(method, handler);
-        }
-    }
-
-    /**
-     * Add a handler for an additional array of methods. Note that {@code MethodHandlers}
-     * does not allow replacing the handler for an already existing method.
-     */
-    MethodHandlers addMethods(RestHandler handler, RestRequest.Method... methods) {
-        for (RestRequest.Method method : methods) {
-            RestHandler existing = methodHandlers.putIfAbsent(method, handler);
-            if (existing != null) {
-                throw new IllegalArgumentException("Cannot replace existing handler for [" + path + "] for method: " + method);
-            }
-        }
-        return this;
-    }
-
+@PublicApi(since = "2.12.0")
+public interface MethodHandlers {
     /**
-     * Returns the handler for the given method or {@code null} if none exists.
+     * Return a set of all valid HTTP methods for the particular path.
      */
-    @Nullable
-    RestHandler getHandler(RestRequest.Method method) {
-        return methodHandlers.get(method);
-    }
+    Set<RestRequest.Method> getValidMethods();
 
     /**
-     * Return a set of all valid HTTP methods for the particular path
+     * Returns the relative HTTP path of the set of method handlers.
      */
-    Set<RestRequest.Method> getValidMethods() {
-        return methodHandlers.keySet();
-    }
+    String getPath();
 }
diff --git a/server/src/main/java/org/opensearch/rest/RestController.java b/server/src/main/java/org/opensearch/rest/RestController.java
@@ -65,6 +65,7 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.net.URI;
+import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Iterator;
@@ -107,7 +108,7 @@ public class RestController implements HttpServerTransport.Dispatcher {
         }
     }
 
-    private final PathTrie<MethodHandlers> handlers = new PathTrie<>(RestUtils.REST_DECODER);
+    private final PathTrie<RestMethodHandlers> handlers = new PathTrie<>(RestUtils.REST_DECODER);
 
     private final UnaryOperator<RestHandler> handlerWrapper;
 
@@ -144,6 +145,16 @@ public RestController(
         );
     }
 
+    /**
+     * Returns an iterator over registered REST method handlers.
+     * @return {@link Iterator} of {@link MethodHandlers}
+     */
+    public Iterator<MethodHandlers> getAllHandlers() {
+        List<MethodHandlers> methodHandlers = new ArrayList<>();
+        handlers.retrieveAll().forEachRemaining(methodHandlers::add);
+        return methodHandlers.iterator();
+    }
+
     /**
      * Registers a REST handler to be executed when the provided {@code method} and {@code path} match the request.
      *
@@ -221,7 +232,7 @@ protected void registerHandler(RestRequest.Method method, String path, RestHandl
     private void registerHandlerNoWrap(RestRequest.Method method, String path, RestHandler maybeWrappedHandler) {
         handlers.insertOrUpdate(
             path,
-            new MethodHandlers(path, maybeWrappedHandler, method),
+            new RestMethodHandlers(path, maybeWrappedHandler, method),
             (mHandlers, newMHandler) -> mHandlers.addMethods(maybeWrappedHandler, method)
         );
     }
@@ -392,10 +403,10 @@ private void tryAllHandlers(final RestRequest request, final RestChannel channel
             // Resolves the HTTP method and fails if the method is invalid
             requestMethod = request.method();
             // Loop through all possible handlers, attempting to dispatch the request
-            Iterator<MethodHandlers> allHandlers = getAllHandlers(request.params(), rawPath);
+            Iterator<RestMethodHandlers> allHandlers = getAllRestMethodHandlers(request.params(), rawPath);
             while (allHandlers.hasNext()) {
                 final RestHandler handler;
-                final MethodHandlers handlers = allHandlers.next();
+                final RestMethodHandlers handlers = allHandlers.next();
                 if (handlers == null) {
                     handler = null;
                 } else {
@@ -423,7 +434,7 @@ private void tryAllHandlers(final RestRequest request, final RestChannel channel
         handleBadRequest(uri, requestMethod, channel);
     }
 
-    Iterator<MethodHandlers> getAllHandlers(@Nullable Map<String, String> requestParamsRef, String rawPath) {
+    Iterator<RestMethodHandlers> getAllRestMethodHandlers(@Nullable Map<String, String> requestParamsRef, String rawPath) {
         final Supplier<Map<String, String>> paramsSupplier;
         if (requestParamsRef == null) {
             paramsSupplier = () -> null;
@@ -561,7 +572,7 @@ private boolean handleAuthenticateUser(final RestRequest request, final RestChan
      */
     private Set<RestRequest.Method> getValidHandlerMethodSet(String rawPath) {
         Set<RestRequest.Method> validMethods = new HashSet<>();
-        Iterator<MethodHandlers> allHandlers = getAllHandlers(null, rawPath);
+        Iterator<RestMethodHandlers> allHandlers = getAllRestMethodHandlers(null, rawPath);
         while (allHandlers.hasNext()) {
             final MethodHandlers methodHandlers = allHandlers.next();
             if (methodHandlers != null) {
diff --git a/server/src/main/java/org/opensearch/rest/RestMethodHandlers.java b/server/src/main/java/org/opensearch/rest/RestMethodHandlers.java
@@ -0,0 +1,92 @@
+/*
+ * 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.
+ */
+
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*
+ * Modifications Copyright OpenSearch Contributors. See
+ * GitHub history for details.
+ */
+
+package org.opensearch.rest;
+
+import org.opensearch.common.Nullable;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Set;
+
+/**
+ * Encapsulate multiple handlers for the same path, allowing different handlers for different HTTP verbs.
+ */
+final class RestMethodHandlers implements MethodHandlers {
+
+    private final String path;
+    private final Map<RestRequest.Method, RestHandler> methodHandlers;
+
+    RestMethodHandlers(String path, RestHandler handler, RestRequest.Method... methods) {
+        this.path = path;
+        this.methodHandlers = new HashMap<>(methods.length);
+        for (RestRequest.Method method : methods) {
+            methodHandlers.put(method, handler);
+        }
+    }
+
+    /**
+     * Add a handler for an additional array of methods. Note that {@code MethodHandlers}
+     * does not allow replacing the handler for an already existing method.
+     */
+    public RestMethodHandlers addMethods(RestHandler handler, RestRequest.Method... methods) {
+        for (RestRequest.Method method : methods) {
+            RestHandler existing = methodHandlers.putIfAbsent(method, handler);
+            if (existing != null) {
+                throw new IllegalArgumentException("Cannot replace existing handler for [" + path + "] for method: " + method);
+            }
+        }
+        return this;
+    }
+
+    /**
+     * Returns the handler for the given method or {@code null} if none exists.
+     */
+    @Nullable
+    public RestHandler getHandler(RestRequest.Method method) {
+        return methodHandlers.get(method);
+    }
+
+    /**
+     * Return a set of all valid HTTP methods for the particular path.
+     */
+    public Set<RestRequest.Method> getValidMethods() {
+        return methodHandlers.keySet();
+    }
+
+    /**
+     * Returns the relative HTTP path of the set of method handlers.
+     */
+    public String getPath() {
+        return path;
+    }
+}
diff --git a/server/src/test/java/org/opensearch/common/path/PathTrieTests.java b/server/src/test/java/org/opensearch/common/path/PathTrieTests.java
@@ -36,8 +36,10 @@
 import org.opensearch.rest.RestUtils;
 import org.opensearch.test.OpenSearchTestCase;
 
+import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.Iterator;
+import java.util.List;
 import java.util.Map;
 
 import static org.hamcrest.Matchers.equalTo;
@@ -286,4 +288,33 @@ public void testEscapedSlashWithinUrl() {
         assertThat(params.get("type"), equalTo("type"));
         assertThat(params.get("id"), equalTo("id"));
     }
+
+    public void testRetrieveAllEmpty() {
+        PathTrie<String> trie = new PathTrie<>(NO_DECODER);
+        Iterator<String> allPaths = trie.retrieveAll();
+        assertFalse(allPaths.hasNext());
+    }
+
+    public void testRetrieveAll() {
+        PathTrie<String> trie = new PathTrie<>(NO_DECODER);
+        trie.insert("{testA}", "test1");
+        trie.insert("{testA}/{testB}", "test2");
+        trie.insert("a/{testB}", "test3");
+        trie.insert("{testA}/b", "test4");
+        trie.insert("{testA}/b/c", "test5");
+
+        Iterator<String> iterator = trie.retrieveAll();
+        assertTrue(iterator.hasNext());
+        List<String> paths = new ArrayList<>();
+        iterator.forEachRemaining(paths::add);
+        assertEquals(paths, List.of("test1", "test4", "test5", "test2", "test3"));
+        assertFalse(iterator.hasNext());
+    }
+
+    public void testRetrieveAllWithNllValue() {
+        PathTrie<String> trie = new PathTrie<>(NO_DECODER);
+        trie.insert("{testA}", null);
+        Iterator<String> iterator = trie.retrieveAll();
+        assertFalse(iterator.hasNext());
+    }
 }
diff --git a/server/src/test/java/org/opensearch/rest/RestControllerTests.java b/server/src/test/java/org/opensearch/rest/RestControllerTests.java
