From 74207499e1b14b898e9f0b95c2c0fa22a048e04c Mon Sep 17 00:00:00 2001
From: Brian Harrington <brharrington@gmail.com>
Date: Thu, 17 Oct 2024 10:35:43 -0500
Subject: [PATCH] atlas: fix bug in PrefixTree

If a node only had a single child after a removal, it would
get compressed by merging the prefix. However the children
of the merged child node were getting dropped.
---
 .../spectator/atlas/impl/PrefixTree.java      |  2 +-
 .../spectator/atlas/impl/PrefixTreeTest.java  | 25 ++++++++++++++++
 .../spectator/atlas/impl/QueryIndexTest.java  | 30 +++++++++++++++++++
 3 files changed, 56 insertions(+), 1 deletion(-)

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 f5948ef1d..dd0b5cdd7 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
@@ -475,7 +475,7 @@ Node compress() {
       } 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.inQueries, c.otherQueries);
+        return new Node(p, c.children, c.inQueries, c.otherQueries);
       } else {
         return new Node(prefix, cs.toArray(EMPTY), inQueries, otherQueries);
       }
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 a4558c82c..67fae9192 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
@@ -202,6 +202,31 @@ public void manyDifferentRootPrefixes() {
     assertEquals(list("abc"), tree, "abcdef");
   }
 
+  @Test
+  public void compressWithSingleChildHavingNoDirectMatches() {
+    PrefixTree tree = new PrefixTree();
+    tree.put(re("abcdef"));
+    tree.put(re("abcghi"));
+    tree.put(re("xyz"));
+    Assertions.assertEquals(3, tree.size());
+
+    tree.remove(re("xyz"));
+    Assertions.assertEquals(2, tree.size());
+  }
+
+  @Test
+  public void compressWithSingleChildHavingDirectMatches() {
+    PrefixTree tree = new PrefixTree();
+    tree.put(re("abc"));
+    tree.put(re("abcdef"));
+    tree.put(re("abcghi"));
+    tree.put(re("xyz"));
+    Assertions.assertEquals(4, tree.size());
+
+    tree.remove(re("xyz"));
+    Assertions.assertEquals(3, tree.size());
+  }
+
   @Test
   public void commonPrefixNoMatch() {
     Assertions.assertEquals(0, PrefixTree.commonPrefixLength("abcdef", "defghi", 0));
diff --git a/spectator-reg-atlas/src/test/java/com/netflix/spectator/atlas/impl/QueryIndexTest.java b/spectator-reg-atlas/src/test/java/com/netflix/spectator/atlas/impl/QueryIndexTest.java
index aac2b82fa..d47c5856b 100644
--- a/spectator-reg-atlas/src/test/java/com/netflix/spectator/atlas/impl/QueryIndexTest.java
+++ b/spectator-reg-atlas/src/test/java/com/netflix/spectator/atlas/impl/QueryIndexTest.java
@@ -636,4 +636,34 @@ public void updateRaceCondition() throws Exception {
       }
     });
   }
+
+  @Test
+  public void otherChecksBug() {
+    Query q1 = Parser.parseQuery("name,foo,:eq,path,abcdef,:re,:and");
+    Query q2 = Parser.parseQuery("name,foo,:eq,path,abcghi,:re,:and");
+    Query q3 = Parser.parseQuery("name,foo,:eq,path,xyz,:re,:and");
+
+    QueryIndex<Integer> idx = QueryIndex.newInstance(new NoopRegistry());
+    idx.add(q1, 1);
+    idx.add(q2, 2);
+    idx.add(q3, 3);
+    Assertions.assertEquals(
+        Collections.singletonList(1),
+        idx.findMatches(id("foo", "path", "abcdef")));
+    Assertions.assertEquals(
+        Collections.singletonList(2),
+        idx.findMatches(id("foo", "path", "abcghijkl")));
+    Assertions.assertEquals(
+        Collections.singletonList(3),
+        idx.findMatches(id("foo", "path", "xyz")));
+
+    idx.remove(q3, 3);
+    Assertions.assertEquals(
+        Collections.singletonList(1),
+        idx.findMatches(id("foo", "path", "abcdef")));
+    Assertions.assertEquals(
+        Collections.singletonList(2),
+        idx.findMatches(id("foo", "path", "abcghijkl")));
+    Assertions.assertTrue(idx.findMatches(id("foo", "path", "xyz")).isEmpty());
+  }
 }
\ No newline at end of file