From f974cb12b565f1eabd8208395391b1c71bcf96af Mon Sep 17 00:00:00 2001
From: Bryan Huhta <32787160+bryanhuhta@users.noreply.github.com>
Date: Mon, 16 Jun 2025 00:33:53 -0700
Subject: [PATCH] FlameGraph: Add support for regex search patterns and
 multiple search terms (#106347)

* "or" search terms with commas

* Add regex support to search bar

* Don't try match empty search terms

* Fix lint error
---
 .../src/FlameGraphContainer.test.tsx          | 109 +++++++++++++++++-
 .../src/FlameGraphContainer.tsx               |  69 +++++++++--
 2 files changed, 160 insertions(+), 18 deletions(-)

diff --git a/packages/grafana-flamegraph/src/FlameGraphContainer.test.tsx b/packages/grafana-flamegraph/src/FlameGraphContainer.test.tsx
index 97a434172fd..94fcfa71a1c 100644
--- a/packages/grafana-flamegraph/src/FlameGraphContainer.test.tsx
+++ b/packages/grafana-flamegraph/src/FlameGraphContainer.test.tsx
@@ -4,8 +4,9 @@ import { useRef, useCallback } from 'react';
 
 import { createDataFrame, createTheme } from '@grafana/data';
 
+import { FlameGraphDataContainer } from './FlameGraph/dataTransform';
 import { data } from './FlameGraph/testData/dataNestedSet';
-import FlameGraphContainer from './FlameGraphContainer';
+import FlameGraphContainer, { labelSearch } from './FlameGraphContainer';
 import { MIN_WIDTH_TO_SHOW_BOTH_TOPTABLE_AND_FLAMEGRAPH } from './constants';
 
 jest.mock('react-use', () => ({
@@ -99,21 +100,117 @@ describe('FlameGraphContainer', () => {
     render(<FlameGraphContainerWithProps />);
 
     // Checking for presence of this function before filter
-    const matchingText = 'net/http.HandlerFunc.ServeHTTP';
+    const matchingText1 = 'net/http.HandlerFunc.ServeHTTP';
+    const matchingText2 = 'runtime.gcBgMarkWorker';
     const nonMatchingText = 'runtime.systemstack';
 
-    expect(screen.queryAllByText(matchingText).length).toBe(1);
+    expect(screen.queryAllByText(matchingText1).length).toBe(1);
+    expect(screen.queryAllByText(matchingText2).length).toBe(1);
     expect(screen.queryAllByText(nonMatchingText).length).toBe(1);
 
     // Apply the filter
-    const searchInput = await screen.getByPlaceholderText('Search...');
-    await userEvent.type(searchInput, 'Handler serve');
+    const searchInput = screen.getByPlaceholderText('Search...');
+    await userEvent.type(searchInput, 'Handler serve,gcBgMarkWorker');
 
     // We have to wait for filter to take effect
     await waitFor(() => {
       expect(screen.queryAllByText(nonMatchingText).length).toBe(0);
     });
     // Check we didn't lose the one that should match
-    expect(screen.queryAllByText(matchingText).length).toBe(1);
+    expect(screen.queryAllByText(matchingText1).length).toBe(1);
+    expect(screen.queryAllByText(matchingText2).length).toBe(1);
+  });
+});
+
+describe('labelSearch', () => {
+  let container: FlameGraphDataContainer;
+
+  beforeEach(() => {
+    const df = createDataFrame(data);
+    df.meta = {
+      custom: {
+        ProfileTypeID: 'cpu:foo:bar',
+      },
+    };
+
+    container = new FlameGraphDataContainer(df, { collapsing: false });
+  });
+
+  describe('fuzzy', () => {
+    it('single term', () => {
+      const search = 'test pkg';
+      let found = labelSearch(search, container);
+      expect(found.size).toBe(45);
+    });
+
+    it('multiple terms', () => {
+      const search = 'test pkg,compress';
+      let found = labelSearch(search, container);
+      expect(found.size).toBe(107);
+    });
+
+    it('falls back to fuzzy with malformed regex', () => {
+      const search = 'deduplicatingSlice[.';
+      let found = labelSearch(search, container);
+      expect(found.size).toBe(1);
+    });
+
+    it('no results', () => {
+      const search = 'term_not_found';
+      let found = labelSearch(search, container);
+      expect(found.size).toBe(0);
+    });
+  });
+
+  describe('regex', () => {
+    it('single pattern', () => {
+      const term = '\\d$';
+      let found = labelSearch(term, container);
+      expect(found.size).toBe(61);
+    });
+
+    it('multiple patterns', () => {
+      const term = '\\d$,^go';
+      let found = labelSearch(term, container);
+      expect(found.size).toBe(62);
+    });
+
+    it('no results', () => {
+      const term = 'pattern_not_found';
+      let found = labelSearch(term, container);
+      expect(found.size).toBe(0);
+    });
+  });
+
+  describe('fuzzy and regex', () => {
+    it('regex found, fuzzy found', () => {
+      const term = '\\d$,test pkg';
+      let found = labelSearch(term, container);
+      expect(found.size).toBe(98);
+    });
+
+    it('regex not found, fuzzy found', () => {
+      const term = 'not_found_suffix$,test pkg';
+      let found = labelSearch(term, container);
+      expect(found.size).toBe(45);
+    });
+
+    it('regex found, fuzzy not found', () => {
+      const term = '\\d$,not_found_fuzzy';
+      let found = labelSearch(term, container);
+      expect(found.size).toBe(61);
+    });
+
+    it('regex not found, fuzzy not found', () => {
+      const term = 'not_found_suffix$,not_found_fuzzy';
+      let found = labelSearch(term, container);
+      expect(found.size).toBe(0);
+    });
+
+    it('does not match empty terms', () => {
+      const search = ',,,,,';
+      let found = labelSearch(search, container);
+      expect(found.size).toBe(0);
+    });
   });
 });
diff --git a/packages/grafana-flamegraph/src/FlameGraphContainer.tsx b/packages/grafana-flamegraph/src/FlameGraphContainer.tsx
index 4d23da838e2..706702142db 100644
--- a/packages/grafana-flamegraph/src/FlameGraphContainer.tsx
+++ b/packages/grafana-flamegraph/src/FlameGraphContainer.tsx
@@ -317,26 +317,71 @@ function useColorScheme(dataContainer: FlameGraphDataContainer | undefined) {
 /**
  * Based on the search string it does a fuzzy search over all the unique labels, so we can highlight them later.
  */
-function useLabelSearch(
+export function useLabelSearch(
   search: string | undefined,
   data: FlameGraphDataContainer | undefined
 ): Set<string> | undefined {
   return useMemo(() => {
-    if (search && data) {
-      const foundLabels = new Set<string>();
-      let idxs = ufuzzy.filter(data.getUniqueLabels(), search);
+    if (!search || !data) {
+      // In this case undefined means there was no search so no attempt to
+      // highlighting anything should be made.
+      return undefined;
+    }
 
-      if (idxs) {
-        for (let idx of idxs) {
-          foundLabels.add(data.getUniqueLabels()[idx]);
-        }
+    return labelSearch(search, data);
+  }, [search, data]);
+}
+
+export function labelSearch(search: string, data: FlameGraphDataContainer): Set<string> {
+  const foundLabels = new Set<string>();
+  const terms = search.split(',');
+
+  const regexFilter = (labels: string[], pattern: string): boolean => {
+    let regex: RegExp;
+    try {
+      regex = new RegExp(pattern);
+    } catch (e) {
+      return false;
+    }
+
+    let foundMatch = false;
+    for (let label of labels) {
+      if (!regex.test(label)) {
+        continue;
       }
 
-      return foundLabels;
+      foundLabels.add(label);
+      foundMatch = true;
     }
-    // In this case undefined means there was no search so no attempt to highlighting anything should be made.
-    return undefined;
-  }, [search, data]);
+    return foundMatch;
+  };
+
+  const fuzzyFilter = (labels: string[], term: string): boolean => {
+    let idxs = ufuzzy.filter(labels, term);
+    if (!idxs) {
+      return false;
+    }
+
+    let foundMatch = false;
+    for (let idx of idxs) {
+      foundLabels.add(labels[idx]);
+      foundMatch = true;
+    }
+    return foundMatch;
+  };
+
+  for (let term of terms) {
+    if (!term) {
+      continue;
+    }
+
+    const found = regexFilter(data.getUniqueLabels(), term);
+    if (!found) {
+      fuzzyFilter(data.getUniqueLabels(), term);
+    }
+  }
+
+  return foundLabels;
 }
 
 function getStyles(theme: GrafanaTheme2) {