Refactor PathTrie and RestController to use a single trie for all methods (#25459)

* Refactor PathTrie and RestController to use a single trie for all methods

This changes `PathTrie` and `RestController` to use a single `PathTrie` for all
endpoints, it also allows retrieving the endpoints' supported HTTP methods more
easily.

This is a spin-off and prerequisite of #24437

* Use EnumSet instead of multiple if conditions

* Make MethodHandlers package-private and final

* Remove duplicate registerHandler method

* Remove public modifier

Author: dakrone

core/src/main/java/org/elasticsearch/common/path/PathTrie.java

@@ -19,14 +19,49 @@
 
 package org.elasticsearch.common.path;
 
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.EnumSet;
 import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.List;
 import java.util.Map;
+import java.util.NoSuchElementException;
+import java.util.Set;
+import java.util.function.BiFunction;
+import java.util.function.Supplier;
 
 import static java.util.Collections.emptyMap;
 import static java.util.Collections.unmodifiableMap;
 
 public class PathTrie<T> {
 
+    enum TrieMatchingMode {
+        /*
+         * Retrieve only explicitly mapped nodes, no wildcards are
+         * matched.
+         */
+        EXPLICIT_NODES_ONLY,
+        /*
+         * Retrieve only explicitly mapped nodes, with wildcards
+         * allowed as root nodes.
+         */
+        WILDCARD_ROOT_NODES_ALLOWED,
+        /*
+         * Retrieve only explicitly mapped nodes, with wildcards
+         * allowed as leaf nodes.
+         */
+        WILDCARD_LEAF_NODES_ALLOWED,
+        /*
+         * Retrieve both explicitly mapped and wildcard nodes.
+         */
+        WILDCARD_NODES_ALLOWED
+    }
+
+    static EnumSet<TrieMatchingMode> EXPLICIT_OR_ROOT_WILDCARD =
+            EnumSet.of(TrieMatchingMode.EXPLICIT_NODES_ONLY, TrieMatchingMode.WILDCARD_ROOT_NODES_ALLOWED);
+
     public interface Decoder {
         String decode(String value);
     }
@@ -107,15 +142,15 @@ public synchronized void insert(String[] path, int index, T value) {
                 if (isNamedWildcard(token)) {
                     node.updateKeyWithNamedWildcard(token);
                 }
-
-                // in case the target(last) node already exist but without a value
-                // than the value should be updated.
+                /*
+                 * If the target node already exists, but is without a value,
+                 *  then the value should be updated.
+                 */
                 if (index == (path.length - 1)) {
                     if (node.value != null) {
                         throw new IllegalArgumentException("Path [" + String.join("/", path)+ "] already has a value ["
                                 + node.value + "]");
-                    }
-                    if (node.value == null) {
+                    } else {
                         node.value = value;
                     }
                 }
@@ -124,6 +159,40 @@ public synchronized void insert(String[] path, int index, T value) {
             node.insert(path, index + 1, value);
         }
 
+        public synchronized void insertOrUpdate(String[] path, int index, T value, BiFunction<T, T, T> updater) {
+            if (index >= path.length)
+                return;
+
+            String token = path[index];
+            String key = token;
+            if (isNamedWildcard(token)) {
+                key = wildcard;
+            }
+            TrieNode node = children.get(key);
+            if (node == null) {
+                T nodeValue = index == path.length - 1 ? value : null;
+                node = new TrieNode(token, nodeValue, wildcard);
+                addInnerChild(key, node);
+            } else {
+                if (isNamedWildcard(token)) {
+                    node.updateKeyWithNamedWildcard(token);
+                }
+                /*
+                 * If the target node already exists, but is without a value,
+                 *  then the value should be updated.
+                 */
+                if (index == (path.length - 1)) {
+                    if (node.value != null) {
+                        node.value = updater.apply(node.value, value);
+                    } else {
+                        node.value = value;
+                    }
+                }
+            }
+
+            node.insertOrUpdate(path, index + 1, value, updater);
+        }
+
         private boolean isNamedWildcard(String key) {
             return key.indexOf('{') != -1 && key.indexOf('}') != -1;
         }
@@ -136,23 +205,57 @@ private boolean isNamedWildcard() {
             return namedWildcard != null;
         }
 
-        public T retrieve(String[] path, int index, Map<String, String> params) {
+        public T retrieve(String[] path, int index, Map<String, String> params, TrieMatchingMode trieMatchingMode) {
             if (index >= path.length)
                 return null;
 
             String token = path[index];
             TrieNode node = children.get(token);
             boolean usedWildcard;
+
             if (node == null) {
-                node = children.get(wildcard);
-                if (node == null) {
+                if (trieMatchingMode == TrieMatchingMode.WILDCARD_NODES_ALLOWED) {
+                    node = children.get(wildcard);
+                    if (node == null) {
+                        return null;
+                    }
+                    usedWildcard = true;
+                } else if (trieMatchingMode == TrieMatchingMode.WILDCARD_ROOT_NODES_ALLOWED && index == 1) {
+                    /*
+                     * Allow root node wildcard matches.
+                     */
+                    node = children.get(wildcard);
+                    if (node == null) {
+                        return null;
+                    }
+                    usedWildcard = true;
+                } else if (trieMatchingMode == TrieMatchingMode.WILDCARD_LEAF_NODES_ALLOWED && index + 1 == path.length) {
+                    /*
+                     * Allow leaf node wildcard matches.
+                     */
+                    node = children.get(wildcard);
+                    if (node == null) {
+                        return null;
+                    }
+                    usedWildcard = true;
+                } else {
                     return null;
                 }
-                usedWildcard = true;
             } else {
-                // If we are at the end of the path, the current node does not have a value but there
-                // is a child wildcard node, use the child wildcard node
-                if (index + 1 == path.length && node.value == null && children.get(wildcard) != null) {
+                if (index + 1 == path.length && node.value == null && children.get(wildcard) != null
+                        && EXPLICIT_OR_ROOT_WILDCARD.contains(trieMatchingMode) == false) {
+                    /*
+                     * If we are at the end of the path, the current node does not have a value but
+                     * there is a child wildcard node, use the child wildcard node.
+                     */
+                    node = children.get(wildcard);
+                    usedWildcard = true;
+                } else if (index == 1 && node.value == null && children.get(wildcard) != null
+                        && trieMatchingMode == TrieMatchingMode.WILDCARD_ROOT_NODES_ALLOWED) {
+                    /*
+                     * If we are at the root, and root wildcards are allowed, use the child wildcard
+                     * node.
+                     */
                     node = children.get(wildcard);
                     usedWildcard = true;
                 } else {
@@ -166,16 +269,16 @@ public T retrieve(String[] path, int index, Map<String, String> params) {
                 return node.value;
             }
 
-            T res = node.retrieve(path, index + 1, params);
-            if (res == null && !usedWildcard) {
+            T nodeValue = node.retrieve(path, index + 1, params, trieMatchingMode);
+            if (nodeValue == null && !usedWildcard && trieMatchingMode != TrieMatchingMode.EXPLICIT_NODES_ONLY) {
                 node = children.get(wildcard);
                 if (node != null) {
                     put(params, node, token);
-                    res = node.retrieve(path, index + 1, params);
+                    nodeValue = node.retrieve(path, index + 1, params, trieMatchingMode);
                 }
             }
 
-            return res;
+            return nodeValue;
         }
 
         private void put(Map<String, String> params, TrieNode node, String value) {
@@ -200,18 +303,47 @@ public void insert(String path, T value) {
             return;
         }
         int index = 0;
-        // supports initial delimiter.
+        // Supports initial delimiter.
         if (strings.length > 0 && strings[0].isEmpty()) {
             index = 1;
         }
         root.insert(strings, index, value);
     }
 
+    /**
+     * Insert a value for the given path. If the path already exists, replace the value with:
+     * <pre>
+     * value = updater.apply(oldValue, newValue);
+     * </pre>
+     * allowing the value to be updated if desired.
+     */
+    public void insertOrUpdate(String path, T value, BiFunction<T, T, T> updater) {
+        String[] strings = path.split(SEPARATOR);
+        if (strings.length == 0) {
+            if (rootValue != null) {
+                rootValue = updater.apply(rootValue, value);
+            } else {
+                rootValue = value;
+            }
+            return;
+        }
+        int index = 0;
+        // Supports initial delimiter.
+        if (strings.length > 0 && strings[0].isEmpty()) {
+            index = 1;
+        }
+        root.insertOrUpdate(strings, index, value, updater);
+    }
+
     public T retrieve(String path) {
-        return retrieve(path, null);
+        return retrieve(path, null, TrieMatchingMode.WILDCARD_NODES_ALLOWED);
     }
 
     public T retrieve(String path, Map<String, String> params) {
+        return retrieve(path, params, TrieMatchingMode.WILDCARD_NODES_ALLOWED);
+    }
+
+    public T retrieve(String path, Map<String, String> params, TrieMatchingMode trieMatchingMode) {
         if (path.length() == 0) {
             return rootValue;
         }
@@ -220,10 +352,56 @@ public T retrieve(String path, Map<String, String> params) {
             return rootValue;
         }
         int index = 0;
-        // supports initial delimiter.
+
+        // Supports initial delimiter.
         if (strings.length > 0 && strings[0].isEmpty()) {
             index = 1;
         }
-        return root.retrieve(strings, index, params);
+
+        return root.retrieve(strings, index, params, trieMatchingMode);
+    }
+
+    /**
+     * Returns an iterator of the objects stored in the {@code PathTrie}, using
+     * all possible {@code TrieMatchingMode} modes. The {@code paramSupplier}
+     * is called between each invocation of {@code next()} to supply a new map
+     * of parameters.
+     */
+    public Iterator<T> retrieveAll(String path, Supplier<Map<String, String>> paramSupplier) {
+        return new PathTrieIterator<>(this, path, paramSupplier);
+    }
+
+    class PathTrieIterator<T> implements Iterator<T> {
+
+        private final List<TrieMatchingMode> modes;
+        private final Supplier<Map<String, String>> paramSupplier;
+        private final PathTrie<T> trie;
+        private final String path;
+
+        PathTrieIterator(PathTrie trie, String path, Supplier<Map<String, String>> paramSupplier) {
+            this.path = path;
+            this.trie = trie;
+            this.paramSupplier = paramSupplier;
+            this.modes = new ArrayList<>(Arrays.asList(TrieMatchingMode.EXPLICIT_NODES_ONLY,
+                TrieMatchingMode.WILDCARD_ROOT_NODES_ALLOWED,
+                TrieMatchingMode.WILDCARD_LEAF_NODES_ALLOWED,
+                TrieMatchingMode.WILDCARD_NODES_ALLOWED));
+            assert TrieMatchingMode.values().length == 4 : "missing trie matching mode";
+        }
+
+        @Override
+        public boolean hasNext() {
+            return modes.isEmpty() == false;
+        }
+
+        @Override
+        public T next() {
+            if (modes.isEmpty()) {
+                throw new NoSuchElementException("called next() without validating hasNext()! no more modes available");
+            }
+            TrieMatchingMode mode = modes.remove(0);
+            Map<String, String> params = paramSupplier.get();
+            return trie.retrieve(path, params, mode);
+        }
     }
 }

core/src/main/java/org/elasticsearch/rest/MethodHandlers.java

@@ -0,0 +1,80 @@
+/*
+ * 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.
+ */
+
+package org.elasticsearch.rest;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+
+/**
+ * Encapsulate multiple handlers for the same path, allowing different handlers for different HTTP verbs.
+ */
+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 an additional method and handler for an existing path. Note that {@code MethodHandlers}
+     * does not allow replacing the handler for an already existing method.
+     */
+    public MethodHandlers addMethod(RestRequest.Method method, RestHandler handler) {
+        RestHandler existing = methodHandlers.putIfAbsent(method, handler);
+        if (existing != null) {
+            throw new IllegalArgumentException("Cannot replace existing handler for [" + path + "] for method: " + method);
+        }
+        return this;
+    }
+
+    /**
+     * 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 MethodHandlers addMethods(RestHandler handler, RestRequest.Method... methods) {
+        for (RestRequest.Method method : methods) {
+            addMethod(method, handler);
+        }
+        return this;
+    }
+
+    /**
+     * Return an Optional-wrapped handler for a method, or an empty optional if
+     * there is no handler.
+     */
+    public Optional<RestHandler> getHandler(RestRequest.Method method) {
+        return Optional.ofNullable(methodHandlers.get(method));
+    }
+
+    /**
+     * Return a set of all valid HTTP methods for the particular path
+     */
+    public Set<RestRequest.Method> getValidMethods() {
+        return methodHandlers.keySet();
+    }
+}