From cf0b717fc09db58828ad20757a87becfe6f1ebcc Mon Sep 17 00:00:00 2001 From: brharrington Date: Thu, 4 Jan 2024 07:12:08 -0600 Subject: [PATCH] use PrefixTree for :in queries (#1107) Updates the query index to leverage the prefix tree for `:in` queries. The `:in` query will only be returned in the result set if it is an exact match, so no further checks are needed. --- .../spectator/atlas/impl/PrefixTree.java | 164 +++++++++++++----- .../spectator/atlas/impl/QueryIndex.java | 22 +-- .../spectator/atlas/impl/PrefixTreeTest.java | 100 ++++++----- 3 files changed, 175 insertions(+), 111 deletions(-) diff --git a/spectator-reg-atlas/src/main/java/com/netflix/spectator/atlas/impl/PrefixTree.java b/spectator-reg-atlas/src/main/java/com/netflix/spectator/atlas/impl/PrefixTree.java index 72874ba4f..4ec3b4140 100644 --- a/spectator-reg-atlas/src/main/java/com/netflix/spectator/atlas/impl/PrefixTree.java +++ b/spectator-reg-atlas/src/main/java/com/netflix/spectator/atlas/impl/PrefixTree.java @@ -34,16 +34,55 @@ * prefix or search key, then the prefix match will only check up to the unsupported * character and the caller will need to perform further checks on the returned value. */ -final class PrefixTree { +final class PrefixTree { private final Lock lock = new ReentrantLock(); private volatile Node root; - private final Set values; + private final Set otherQueries; /** Create a new instance. */ PrefixTree() { - values = newSet(); + this.otherQueries = newSet(); + } + + private Node addQuery(Node node, Query.KeyQuery query) { + if (query instanceof Query.In) { + Query.In q = (Query.In) query; + node.inQueries.add(q); + } else { + node.otherQueries.add(query); + } + return node; + } + + private boolean removeQuery(Node node, Query.KeyQuery query) { + if (query instanceof Query.In) { + Query.In q = (Query.In) query; + return node.inQueries.remove(q); + } else { + return node.otherQueries.remove(query); + } + } + + /** + * Put a query into the tree. + * + * @param query + * Query to add, the prefix will be extracted from the query clause. + */ + void put(Query.KeyQuery query) { + if (query instanceof Query.In) { + Query.In q = (Query.In) query; + for (String v : q.values()) { + put(v, q); + } + } else if (query instanceof Query.Regex) { + Query.Regex q = (Query.Regex) query; + put(q.pattern().prefix(), q); + } else { + otherQueries.add(query); + } } /** @@ -54,15 +93,16 @@ final class PrefixTree { * @param value * Value to associate with the prefix. */ - void put(String prefix, T value) { + void put(String prefix, Query.KeyQuery value) { if (prefix == null) { - values.add(value); + otherQueries.add(value); } else { lock.lock(); try { Node node = root; if (node == null) { - root = new Node(prefix, EMPTY, asSet(value)); + root = new Node(prefix, EMPTY); + addQuery(root, value); } else { root = putImpl(node, prefix, 0, value); } @@ -72,17 +112,17 @@ void put(String prefix, T value) { } } - private Node putImpl(Node node, String key, int offset, T value) { + private Node putImpl(Node node, String key, int offset, Query.KeyQuery value) { final int prefixLength = node.prefix.length(); final int keyLength = key.length() - offset; final int commonLength = commonPrefixLength(node.prefix, key, offset); if (commonLength == 0 && prefixLength > 0) { // No common prefix - Node n = new Node(key.substring(offset), EMPTY, asSet(value)); + Node n = addQuery(new Node(key.substring(offset), EMPTY), value); return new Node("", new Node[] {n, node}); } else if (keyLength == prefixLength && commonLength == prefixLength) { // Fully matches, add the value to this node - node.values.add(value); + addQuery(node, value); return node; } else if (keyLength > prefixLength && commonLength == prefixLength) { // key.startsWith(prefix), put the value into a child @@ -92,25 +132,49 @@ private Node putImpl(Node node, String key, int offset, T value) { Node n = putImpl(node.children[pos], key, childOffset, value); return node.replaceChild(n, pos); } else { - Node n = new Node(key.substring(childOffset), EMPTY, asSet(value)); + Node n = addQuery(new Node(key.substring(childOffset), EMPTY), value); return node.addChild(n); } } else if (prefixLength > keyLength && commonLength == keyLength) { // prefix.startsWith(key), make new parent node and add this node as a child int childOffset = offset + commonLength; - Node n = new Node(node.prefix.substring(commonLength), node.children, node.values); - return new Node(key.substring(offset, childOffset), new Node[] {n}, asSet(value)); + Node n = new Node(node.prefix.substring(commonLength), node.children, node.inQueries, node.otherQueries); + return addQuery(new Node(key.substring(offset, childOffset), new Node[] {n}), value); } else { // Common prefix is a subset of both int childOffset = offset + commonLength; Node[] children = { - new Node(node.prefix.substring(commonLength), node.children, node.values), - new Node(key.substring(childOffset), EMPTY, asSet(value)) + new Node(node.prefix.substring(commonLength), node.children, node.inQueries, node.otherQueries), + addQuery(new Node(key.substring(childOffset), EMPTY), value) }; return new Node(node.prefix.substring(0, commonLength), children); } } + /** + * Remove a value from the tree with the associated prefix. + * + * @param query + * Query to remove, the prefix will be extracted from the query clause. + * @return + * Returns true if a value was removed from the tree. + */ + boolean remove(Query.KeyQuery query) { + if (query instanceof Query.In) { + boolean removed = false; + Query.In q = (Query.In) query; + for (String v : q.values()) { + removed |= remove(v, q); + } + return removed; + } else if (query instanceof Query.Regex) { + Query.Regex q = (Query.Regex) query; + return remove(q.pattern().prefix(), q); + } else { + return otherQueries.remove(query); + } + } + /** * Remove a value from the tree with the associated prefix. * @@ -121,9 +185,9 @@ private Node putImpl(Node node, String key, int offset, T value) { * @return * Returns true if a value was removed from the tree. */ - boolean remove(String prefix, T value) { + boolean remove(String prefix, Query.KeyQuery value) { if (prefix == null) { - return values.remove(value); + return otherQueries.remove(value); } else { lock.lock(); try { @@ -143,13 +207,13 @@ boolean remove(String prefix, T value) { } } - private boolean removeImpl(Node node, String key, int offset, T value) { + private boolean removeImpl(Node node, String key, int offset, Query.KeyQuery value) { final int prefixLength = node.prefix.length(); final int keyLength = key.length() - offset; final int commonLength = commonPrefixLength(node.prefix, key, offset); if (keyLength == prefixLength && commonLength == prefixLength) { // Fully matches, remove the value from this node - return node.values.remove(value); + return removeQuery(node, value); } else if (keyLength > prefixLength && commonLength == prefixLength) { // Try to remove from children int childOffset = offset + commonLength; @@ -168,8 +232,8 @@ private boolean removeImpl(Node node, String key, int offset, T value) { * @return * Values associated with a matching prefix. */ - List get(String key) { - List result = new ArrayList<>(); + List get(String key) { + List result = new ArrayList<>(); forEach(key, result::add); return result; } @@ -182,25 +246,23 @@ List get(String key) { * @param consumer * Function to call for matching values. */ - void forEach(String key, Consumer consumer) { - values.forEach(consumer); + void forEach(String key, Consumer consumer) { + // In queries cannot have an empty value, so cannot be in the root set + otherQueries.forEach(consumer); Node node = root; if (node != null) { forEachImpl(node, key, 0, consumer); } } - @SuppressWarnings("unchecked") - private void forEachImpl(Node node, String key, int offset, Consumer consumer) { + private void forEachImpl(Node node, String key, int offset, Consumer consumer) { final int prefixLength = node.prefix.length(); final int keyLength = key.length() - offset; final int commonLength = commonPrefixLength(node.prefix, key, offset); if (commonLength == prefixLength) { - // Prefix matches, consume values for this node - for (Object value : node.values) { - consumer.accept((T) value); - } + // Prefix matches, consume other queries + node.otherQueries.forEach(consumer); if (commonLength < keyLength) { // There is more to the key, check if there are also matches for child nodes @@ -209,6 +271,9 @@ private void forEachImpl(Node node, String key, int offset, Consumer consumer if (pos >= 0) { forEachImpl(node.children[pos], key, childOffset, consumer); } + } else { + // It is an exact match, consume in queries + node.inQueries.forEach(consumer); } } } @@ -217,7 +282,7 @@ private void forEachImpl(Node node, String key, int offset, Consumer consumer * Returns true if the tree is empty. */ boolean isEmpty() { - return values.isEmpty() && (root == null || root.isEmpty()); + return otherQueries.isEmpty() && (root == null || root.isEmpty()); } /** @@ -226,7 +291,7 @@ boolean isEmpty() { */ int size() { Node r = root; - return (r == null ? 0 : r.size()) + values.size(); + return (r == null ? 0 : r.size()) + otherQueries.size(); } /** @@ -273,8 +338,8 @@ private static Set newSet() { return new CopyOnWriteArraySet<>(); } - private static Set asSet(Object value) { - Set set = newSet(); + private static Set asSet(Query.KeyQuery value) { + Set set = newSet(); set.add(value); return set; } @@ -285,31 +350,33 @@ private static Set asSet(Object value) { public boolean equals(Object o) { if (this == o) return true; if (o == null || getClass() != o.getClass()) return false; - PrefixTree that = (PrefixTree) o; + PrefixTree that = (PrefixTree) o; return Objects.equals(root, that.root) - && values.equals(that.values); + && otherQueries.equals(that.otherQueries); } @Override public int hashCode() { - return Objects.hash(root, values); + return Objects.hash(root, otherQueries); } private static class Node { final String prefix; final Node[] children; - final Set values; + final Set inQueries; + final Set otherQueries; - Node(String prefix, Node[] children, Set values) { + Node(String prefix, Node[] children, Set inQueries, Set otherQueries) { this.prefix = Preconditions.checkNotNull(prefix, "prefix"); this.children = Preconditions.checkNotNull(children, "children"); - this.values = Preconditions.checkNotNull(values, "values"); + this.inQueries = Preconditions.checkNotNull(inQueries, "inQueries"); + this.otherQueries = Preconditions.checkNotNull(otherQueries, "otherQueries"); Arrays.sort(children, Comparator.comparing(n -> n.prefix)); } Node(String prefix, Node[] children) { - this(prefix, children, newSet()); + this(prefix, children, newSet(), newSet()); } Node replaceChild(Node n, int i) { @@ -317,14 +384,14 @@ Node replaceChild(Node n, int i) { System.arraycopy(children, 0, cs, 0, i); cs[i] = n; System.arraycopy(children, i + 1, cs, i + 1, children.length - i - 1); - return new Node(prefix, cs, values); + return new Node(prefix, cs, inQueries, otherQueries); } Node addChild(Node n) { Node[] cs = new Node[children.length + 1]; System.arraycopy(children, 0, cs, 0, children.length); cs[children.length] = n; - return new Node(prefix, cs, values); + return new Node(prefix, cs, inQueries, otherQueries); } Node compress() { @@ -352,17 +419,17 @@ Node compress() { // Return compressed node. Merge nodes if intermediates have no values. if (cs == null) { return this; - } else if (values.isEmpty() && cs.size() == 1) { + } else if (inQueries.isEmpty() && otherQueries.isEmpty() && cs.size() == 1) { Node c = cs.get(0); String p = prefix + c.prefix; - return new Node(p, EMPTY, c.values); + return new Node(p, EMPTY, c.inQueries, c.otherQueries); } else { - return new Node(prefix, cs.toArray(EMPTY), values); + return new Node(prefix, cs.toArray(EMPTY), inQueries, otherQueries); } } boolean isEmpty() { - return values.isEmpty() && areAllChildrenEmpty(); + return inQueries.isEmpty() && otherQueries.isEmpty() && areAllChildrenEmpty(); } private boolean areAllChildrenEmpty() { @@ -375,7 +442,7 @@ private boolean areAllChildrenEmpty() { } int size() { - int sz = values.size(); + int sz = inQueries.size() + otherQueries.size(); for (Node child : children) { sz += child.size(); } @@ -389,12 +456,13 @@ public boolean equals(Object o) { Node node = (Node) o; return prefix.equals(node.prefix) && Arrays.equals(children, node.children) - && values.equals(node.values); + && inQueries.equals(node.inQueries) + && otherQueries.equals(node.otherQueries); } @Override public int hashCode() { - int result = Objects.hash(prefix, values); + int result = Objects.hash(prefix, inQueries, otherQueries); result = 31 * result + Arrays.hashCode(children); return result; } diff --git a/spectator-reg-atlas/src/main/java/com/netflix/spectator/atlas/impl/QueryIndex.java b/spectator-reg-atlas/src/main/java/com/netflix/spectator/atlas/impl/QueryIndex.java index 04c2234cd..5ea05d2be 100644 --- a/spectator-reg-atlas/src/main/java/com/netflix/spectator/atlas/impl/QueryIndex.java +++ b/spectator-reg-atlas/src/main/java/com/netflix/spectator/atlas/impl/QueryIndex.java @@ -116,7 +116,7 @@ private static int compare(String k1, String k2) { // filter regex and in clauses. The matching is cached to avoid expensive regex checks // as much as possible. private final ConcurrentHashMap> otherChecks; - private final PrefixTree otherChecksTree; + private final PrefixTree otherChecksTree; private final Cache>> otherChecksCache; // Index for :has queries @@ -137,7 +137,7 @@ private QueryIndex(CacheSupplier cacheSupplier, String key) { this.key = key; this.equalChecks = new ConcurrentHashMap<>(); this.otherChecks = new ConcurrentHashMap<>(); - this.otherChecksTree = new PrefixTree<>(); + this.otherChecksTree = new PrefixTree(); this.otherChecksCache = cacheSupplier.get(); this.hasKeyIdx = null; this.otherKeysIdx = null; @@ -216,12 +216,7 @@ private void add(List queries, int i, T value) { } else { QueryIndex idx = otherChecks.computeIfAbsent(kq, id -> QueryIndex.empty(cacheSupplier)); idx.add(queries, j, value); - if (kq instanceof Query.Regex) { - Query.Regex re = (Query.Regex) kq; - otherChecksTree.put(re.pattern().prefix(), kq); - } else { - otherChecksTree.put(null, kq); - } + otherChecksTree.put(kq); otherChecksCache.clear(); // Not queries should match if the key is missing from the id, so they need to @@ -308,12 +303,7 @@ private boolean remove(List queries, int i, T value) { otherChecksCache.clear(); if (idx.isEmpty()) { otherChecks.remove(kq); - if (kq instanceof Query.Regex) { - Query.Regex re = (Query.Regex) kq; - otherChecksTree.remove(re.pattern().prefix(), kq); - } else { - otherChecksTree.remove(null, kq); - } + otherChecksTree.remove(kq); } } @@ -406,7 +396,7 @@ private void forEachMatch(Id tags, int i, Consumer consumer) { if (!otherChecks.isEmpty()) { List> tmp = new ArrayList<>(); otherChecksTree.forEach(v, kq -> { - if (kq.matches(v)) { + if (kq instanceof Query.In || kq.matches(v)) { QueryIndex idx = otherChecks.get(kq); if (idx != null) { tmp.add(idx); @@ -499,7 +489,7 @@ public void forEachMatch(Function tags, Consumer consumer) { if (!otherChecks.isEmpty()) { List> tmp = new ArrayList<>(); otherChecksTree.forEach(v, kq -> { - if (matches(kq, v)) { + if (kq instanceof Query.In || matches(kq, v)) { QueryIndex idx = otherChecks.get(kq); if (idx != null) { tmp.add(idx); diff --git a/spectator-reg-atlas/src/test/java/com/netflix/spectator/atlas/impl/PrefixTreeTest.java b/spectator-reg-atlas/src/test/java/com/netflix/spectator/atlas/impl/PrefixTreeTest.java index 7d4d368c0..ac622396b 100644 --- a/spectator-reg-atlas/src/test/java/com/netflix/spectator/atlas/impl/PrefixTreeTest.java +++ b/spectator-reg-atlas/src/test/java/com/netflix/spectator/atlas/impl/PrefixTreeTest.java @@ -20,57 +20,63 @@ import java.util.Arrays; import java.util.Collections; +import java.util.Comparator; import java.util.List; +import java.util.stream.Collectors; public class PrefixTreeTest { - private List list(String... values) { - return Arrays.asList(values); + private List list(String... values) { + return Arrays.stream(values).map(PrefixTreeTest::re).collect(Collectors.toList()); } - private List sort(List values) { - values.sort(String::compareTo); + private List sort(List values) { + values.sort(Comparator.comparing(Object::toString)); return values; } - private void assertSize(PrefixTree tree, int expected) { + private static Query.KeyQuery re(String s) { + return new Query.Regex("name", "^" + s); + } + + private void assertSize(PrefixTree tree, int expected) { Assertions.assertEquals(expected, tree.size()); Assertions.assertEquals(expected == 0, tree.isEmpty()); } @Test public void nullPrefix() { - PrefixTree tree = new PrefixTree<>(); - tree.put(null, "1"); + PrefixTree tree = new PrefixTree(); + tree.put(null, re("1")); assertSize(tree, 1); Assertions.assertEquals(list("1"), tree.get("foo")); Assertions.assertEquals(list("1"), tree.get("bar")); Assertions.assertEquals(list("1"), tree.get("")); - Assertions.assertFalse(tree.remove(null, "2")); - Assertions.assertTrue(tree.remove(null, "1")); + Assertions.assertFalse(tree.remove(null, re("2"))); + Assertions.assertTrue(tree.remove(null, re("1"))); assertSize(tree, 0); Assertions.assertEquals(Collections.emptyList(), tree.get("foo")); } @Test public void emptyPrefix() { - PrefixTree tree = new PrefixTree<>(); - tree.put("", "1"); + PrefixTree tree = new PrefixTree(); + tree.put("", re("1")); Assertions.assertEquals(list("1"), tree.get("foo")); Assertions.assertEquals(list("1"), tree.get("bar")); Assertions.assertEquals(list("1"), tree.get("")); - Assertions.assertFalse(tree.remove("", "2")); - Assertions.assertTrue(tree.remove("", "1")); + Assertions.assertFalse(tree.remove("", re("2"))); + Assertions.assertTrue(tree.remove("", re("1"))); Assertions.assertEquals(Collections.emptyList(), tree.get("foo")); } @Test public void simplePrefix() { - PrefixTree tree = new PrefixTree<>(); - tree.put("abc", "1"); + PrefixTree tree = new PrefixTree(); + tree.put("abc", re("1")); assertSize(tree, 1); Assertions.assertEquals(list("1"), tree.get("abcdef")); @@ -79,19 +85,19 @@ public void simplePrefix() { Assertions.assertEquals(Collections.emptyList(), tree.get("abd")); Assertions.assertEquals(Collections.emptyList(), tree.get("ab")); - Assertions.assertTrue(tree.remove("abc", "1")); - Assertions.assertFalse(tree.remove("abc", "1")); + Assertions.assertTrue(tree.remove("abc", re("1"))); + Assertions.assertFalse(tree.remove("abc", re("1"))); assertSize(tree, 0); Assertions.assertEquals(Collections.emptyList(), tree.get("abcdef")); } @Test public void multipleMatches() { - PrefixTree tree = new PrefixTree<>(); - tree.put("abc", "1"); - tree.put("ab", "2"); - tree.put("a", "3"); - tree.put("abc", "4"); + PrefixTree tree = new PrefixTree(); + tree.put("abc", re("1")); + tree.put("ab", re("2")); + tree.put("a", re("3")); + tree.put("abc", re("4")); assertSize(tree, 4); Assertions.assertEquals(list("1", "2", "3", "4"), sort(tree.get("abcdef"))); @@ -99,16 +105,16 @@ public void multipleMatches() { Assertions.assertEquals(list("3"), tree.get("adef")); Assertions.assertEquals(Collections.emptyList(), tree.get("bcdef")); - Assertions.assertFalse(tree.remove("ab", "1")); - Assertions.assertTrue(tree.remove("abc", "1")); + Assertions.assertFalse(tree.remove("ab", re("1"))); + Assertions.assertTrue(tree.remove("abc", re("1"))); assertSize(tree, 3); Assertions.assertEquals(list("2", "3", "4"), sort(tree.get("abcdef"))); } @Test public void unicodeCharInPrefix() { - PrefixTree tree = new PrefixTree<>(); - tree.put("aβc", "1"); + PrefixTree tree = new PrefixTree(); + tree.put("aβc", re("1")); assertSize(tree, 1); Assertions.assertEquals(list("1"), tree.get("aβcdef")); @@ -122,24 +128,24 @@ public void unicodeCharInPrefix() { Assertions.assertEquals(Collections.emptyList(), tree.get("aβ")); Assertions.assertEquals(Collections.emptyList(), tree.get("b")); - Assertions.assertTrue(tree.remove("aβc", "1")); + Assertions.assertTrue(tree.remove("aβc", re("1"))); assertSize(tree, 0); Assertions.assertEquals(Collections.emptyList(), tree.get("abcdef")); } @Test public void emptyRootPrefix() { - PrefixTree tree = new PrefixTree<>(); - tree.put("ab", "ab"); - tree.put("cd", "cd"); + PrefixTree tree = new PrefixTree(); + tree.put("ab", re("ab")); + tree.put("cd", re("cd")); Assertions.assertEquals(list("ab"), tree.get("ab")); Assertions.assertEquals(list("cd"), tree.get("cd")); Assertions.assertEquals(Collections.emptyList(), tree.get("ef")); - Assertions.assertTrue(tree.remove("ab", "ab")); - PrefixTree tree2 = new PrefixTree<>(); - tree2.put("cd", "cd"); + Assertions.assertTrue(tree.remove("ab", re("ab"))); + PrefixTree tree2 = new PrefixTree(); + tree2.put("cd", re("cd")); Assertions.assertEquals(tree2, tree); Assertions.assertEquals(Collections.emptyList(), tree.get("ab")); @@ -149,19 +155,19 @@ public void emptyRootPrefix() { @Test public void updateExistingNode() { - PrefixTree tree = new PrefixTree<>(); - tree.put("abcdef", "1"); - tree.put("abcdef", "2"); + PrefixTree tree = new PrefixTree(); + tree.put("abcdef", re("1")); + tree.put("abcdef", re("2")); Assertions.assertEquals(list("1", "2"), tree.get("abcdef")); } @Test public void addChildNode() { - PrefixTree tree = new PrefixTree<>(); - tree.put("abc", "1"); - tree.put("abcdefghi", "2"); - tree.put("abcdef", "3"); + PrefixTree tree = new PrefixTree(); + tree.put("abc", re("1")); + tree.put("abcdefghi", re("2")); + tree.put("abcdef", re("3")); Assertions.assertEquals(list("1", "3"), tree.get("abcdef")); Assertions.assertEquals(list("1", "3", "2"), tree.get("abcdefghi")); @@ -169,9 +175,9 @@ public void addChildNode() { @Test public void splitInteriorNode() { - PrefixTree tree = new PrefixTree<>(); - tree.put("abcdef", "abcdef"); - tree.put("abcghi", "abcghi"); + PrefixTree tree = new PrefixTree(); + tree.put("abcdef", re("abcdef")); + tree.put("abcghi", re("abcghi")); Assertions.assertEquals(list("abcdef"), tree.get("abcdef")); Assertions.assertEquals(list("abcghi"), tree.get("abcghi")); @@ -179,10 +185,10 @@ public void splitInteriorNode() { @Test public void manyDifferentRootPrefixes() { - PrefixTree tree = new PrefixTree<>(); - tree.put("abc", "abc"); - tree.put("def", "def"); - tree.put("ghi", "ghi"); + PrefixTree tree = new PrefixTree(); + tree.put("abc", re("abc")); + tree.put("def", re("def")); + tree.put("ghi", re("ghi")); Assertions.assertEquals(list("abc"), tree.get("abcdef")); }