From 44bd81873da9e7917507f35976f7ec5df223e49f Mon Sep 17 00:00:00 2001
From: brharrington <brharrington@gmail.com>
Date: Sat, 2 Dec 2023 13:25:14 -0600
Subject: [PATCH] atlas: add method to find hotspots in index (#1098)

Add helper method to QueryIndex to help find hot spots
in the index.
---
 .../spectator/atlas/impl/QueryIndex.java      | 66 +++++++++++++++++++
 .../spectator/atlas/impl/QueryIndexTest.java  | 18 +++++
 2 files changed, 84 insertions(+)

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 609a67a0e..a0c21179c 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
@@ -19,12 +19,15 @@
 import com.netflix.spectator.api.Registry;
 import com.netflix.spectator.impl.Cache;
 
+import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.Deque;
 import java.util.List;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.CopyOnWriteArraySet;
+import java.util.function.BiConsumer;
 import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Supplier;
@@ -533,6 +536,69 @@ public void forEachMatch(Function<String, String> tags, Consumer<T> consumer) {
     }
   }
 
+  /**
+   * Find hot spots in the index where there is a large set of linear matches, e.g. a bunch
+   * of regex queries for a given key.
+   *
+   * @param threshold
+   *     Threshold for the number of entries in the other checks sub-tree to be considered
+   *     a hot spot.
+   * @param consumer
+   *     Function that will be invoked with a path and set of queries for the hot spot.
+   */
+  public void findHotSpots(int threshold, BiConsumer<List<String>, List<Query.KeyQuery>> consumer) {
+    Deque<String> path = new ArrayDeque<>();
+    findHotSpots(threshold, path, consumer);
+  }
+
+  private void findHotSpots(
+      int threshold,
+      Deque<String> path,
+      BiConsumer<List<String>, List<Query.KeyQuery>> consumer
+  ) {
+    if (key != null) {
+      path.addLast("K=" + key);
+
+      equalChecks.forEach((v, idx) -> {
+        path.addLast(key + "," + v + ",:eq");
+        idx.findHotSpots(threshold, path, consumer);
+        path.removeLast();
+      });
+
+      path.addLast("other-checks");
+      if (otherChecks.size() > threshold) {
+        List<Query.KeyQuery> queries = new ArrayList<>(otherChecks.keySet());
+        consumer.accept(new ArrayList<>(path), queries);
+      }
+      otherChecks.forEach((q, idx) -> {
+        path.addLast(q.toString());
+        idx.findHotSpots(threshold, path, consumer);
+        path.removeLast();
+      });
+      path.removeLast();
+
+      if (hasKeyIdx != null) {
+        path.addLast("has");
+        hasKeyIdx.findHotSpots(threshold, path, consumer);
+        path.removeLast();
+      }
+
+      path.removeLast();
+    }
+
+    if (otherKeysIdx != null) {
+      path.addLast("other-keys");
+      otherKeysIdx.findHotSpots(threshold, path, consumer);
+      path.removeLast();
+    }
+
+    if (missingKeysIdx != null) {
+      path.addLast("missing-keys");
+      missingKeysIdx.findHotSpots(threshold, path, consumer);
+      path.removeLast();
+    }
+  }
+
   @Override public String toString() {
     StringBuilder builder = new StringBuilder();
     buildString(builder, 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 6e04f14ea..17aaefb67 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
@@ -533,4 +533,22 @@ public void addRemoveFuzz() {
       Assertions.assertTrue(idx.isEmpty());
     }
   }
+
+  @Test
+  public void findHotSpots() {
+    Registry registry = new NoopRegistry();
+    QueryIndex<Integer> idx = QueryIndex.newInstance(registry);
+    for (int i = 0; i < 5; ++i) {
+      Query q = Parser.parseQuery("name,foo,:re,i," + i + ",:re,:and");
+      idx.add(q, i);
+    }
+
+    idx.findHotSpots(4, (path, queries) -> {
+      Assertions.assertEquals(
+          "K=name > other-checks > name,foo,:re > K=i > other-checks",
+          String.join(" > ", path)
+      );
+      Assertions.assertEquals(5, queries.size());
+    });
+  }
 }
\ No newline at end of file