Tempo: Add filtering for service graph query (#41162)

* Add filter based on AdHocFilter element

* Add tests

* Cancel layout in case we have have new data or we unmount node graph

* Fix typing

* Fix test

public/app/features/variables/adhoc/picker/AdHocFilter.test.tsx

@@ -0,0 +1,109 @@
+import React from 'react';
+import { render, screen } from '@testing-library/react';
+import userEvent from '@testing-library/user-event';
+import selectEvent from 'react-select-event';
+import { AdHocFilter } from './AdHocFilter';
+import { AdHocVariableFilter } from '../../types';
+import { setDataSourceSrv } from '../../../../../../packages/grafana-runtime';
+
+describe('AdHocFilter', () => {
+  it('renders filters', async () => {
+    setup();
+    expect(screen.getByText('key1')).toBeInTheDocument();
+    expect(screen.getByText('val1')).toBeInTheDocument();
+    expect(screen.getByText('key2')).toBeInTheDocument();
+    expect(screen.getByText('val2')).toBeInTheDocument();
+    expect(screen.getByLabelText('Add Filter')).toBeInTheDocument();
+  });
+
+  it('adds filter', async () => {
+    const { addFilter } = setup();
+
+    // Select key
+    userEvent.click(screen.getByLabelText('Add Filter'));
+    const selectEl = screen.getByTestId('AdHocFilterKey-add-key-wrapper');
+    expect(selectEl).toBeInTheDocument();
+    await selectEvent.select(selectEl, 'key3', { container: document.body });
+
+    // Select value
+    userEvent.click(screen.getByText('select value'));
+    screen.debug(screen.getAllByTestId('AdHocFilterValue-value-wrapper'));
+    // There are already some filters rendered
+    const selectEl2 = screen.getAllByTestId('AdHocFilterValue-value-wrapper')[2];
+    await selectEvent.select(selectEl2, 'val3', { container: document.body });
+
+    // Only after value is selected the addFilter is called
+    expect(addFilter).toBeCalled();
+  });
+
+  it('removes filter', async () => {
+    const { removeFilter } = setup();
+
+    // Select key
+    userEvent.click(screen.getByText('key1'));
+    const selectEl = screen.getAllByTestId('AdHocFilterKey-key-wrapper')[0];
+    expect(selectEl).toBeInTheDocument();
+    await selectEvent.select(selectEl, '-- remove filter --', { container: document.body });
+
+    // Only after value is selected the addFilter is called
+    expect(removeFilter).toBeCalled();
+  });
+
+  it('changes filter', async () => {
+    const { changeFilter } = setup();
+
+    // Select key
+    userEvent.click(screen.getByText('val1'));
+    const selectEl = screen.getAllByTestId('AdHocFilterValue-value-wrapper')[0];
+    expect(selectEl).toBeInTheDocument();
+    await selectEvent.select(selectEl, 'val4', { container: document.body });
+
+    // Only after value is selected the addFilter is called
+    expect(changeFilter).toBeCalled();
+  });
+});
+
+function setup() {
+  setDataSourceSrv({
+    get() {
+      return {
+        getTagKeys() {
+          return [{ text: 'key3' }];
+        },
+        getTagValues() {
+          return [{ text: 'val3' }, { text: 'val4' }];
+        },
+      };
+    },
+  } as any);
+
+  const filters: AdHocVariableFilter[] = [
+    {
+      key: 'key1',
+      operator: '=',
+      value: 'val1',
+      condition: '',
+    },
+    {
+      key: 'key2',
+      operator: '=',
+      value: 'val2',
+      condition: '',
+    },
+  ];
+  const addFilter = jest.fn();
+  const removeFilter = jest.fn();
+  const changeFilter = jest.fn();
+
+  render(
+    <AdHocFilter
+      datasource={{ uid: 'test' }}
+      filters={filters}
+      addFilter={addFilter}
+      removeFilter={removeFilter}
+      changeFilter={changeFilter}
+    />
+  );
+
+  return { addFilter, removeFilter, changeFilter };
+}

public/app/features/variables/adhoc/picker/AdHocFilter.tsx

@@ -0,0 +1,82 @@
+import React, { PureComponent, ReactNode } from 'react';
+import { AdHocVariableFilter } from 'app/features/variables/types';
+import { DataSourceRef, SelectableValue } from '@grafana/data';
+import { AdHocFilterBuilder } from './AdHocFilterBuilder';
+import { ConditionSegment } from './ConditionSegment';
+import { REMOVE_FILTER_KEY } from './AdHocFilterKey';
+import { AdHocFilterRenderer } from './AdHocFilterRenderer';
+
+interface Props {
+  datasource: DataSourceRef | null;
+  filters: AdHocVariableFilter[];
+  addFilter: (filter: AdHocVariableFilter) => void;
+  removeFilter: (index: number) => void;
+  changeFilter: (index: number, newFilter: AdHocVariableFilter) => void;
+}
+
+/**
+ * Simple filtering component that automatically uses datasource APIs to get available labels and it's values, for
+ * dynamic visual filtering without need for much setup. Instead of having single onChange prop this reports all the
+ * change events with separate props so it is usable with AdHocPicker.
+ *
+ * Note: There isn't API on datasource to suggest the operators here so that is hardcoded to use prometheus style
+ * operators. Also filters are assumed to be joined with `AND` operator, which is also hardcoded.
+ */
+export class AdHocFilter extends PureComponent<Props> {
+  onChange = (index: number, prop: string) => (key: SelectableValue<string | null>) => {
+    const { filters } = this.props;
+    const { value } = key;
+
+    if (key.value === REMOVE_FILTER_KEY) {
+      return this.props.removeFilter(index);
+    }
+
+    return this.props.changeFilter(index, {
+      ...filters[index],
+      [prop]: value,
+    });
+  };
+
+  appendFilterToVariable = (filter: AdHocVariableFilter) => {
+    this.props.addFilter(filter);
+  };
+
+  render() {
+    const { filters } = this.props;
+
+    return (
+      <div className="gf-form-inline">
+        {this.renderFilters(filters)}
+        <AdHocFilterBuilder
+          datasource={this.props.datasource!}
+          appendBefore={filters.length > 0 ? <ConditionSegment label="AND" /> : null}
+          onCompleted={this.appendFilterToVariable}
+        />
+      </div>
+    );
+  }
+
+  renderFilters(filters: AdHocVariableFilter[]) {
+    return filters.reduce((segments: ReactNode[], filter, index) => {
+      if (segments.length > 0) {
+        segments.push(<ConditionSegment label="AND" key={`condition-${index}`} />);
+      }
+      segments.push(this.renderFilterSegments(filter, index));
+      return segments;
+    }, []);
+  }
+
+  renderFilterSegments(filter: AdHocVariableFilter, index: number) {
+    return (
+      <React.Fragment key={`filter-${index}`}>
+        <AdHocFilterRenderer
+          datasource={this.props.datasource!}
+          filter={filter}
+          onKeyChange={this.onChange(index, 'key')}
+          onOperatorChange={this.onChange(index, 'operator')}
+          onValueChange={this.onChange(index, 'value')}
+        />
+      </React.Fragment>
+    );
+  }
+}

public/app/features/variables/adhoc/picker/AdHocFilterKey.tsx

@@ -16,7 +16,7 @@ export const AdHocFilterKey: FC<Props> = ({ datasource, onChange, filterKey }) =
 
   if (filterKey === null) {
     return (
-      <div className="gf-form">
+      <div className="gf-form" data-testid="AdHocFilterKey-add-key-wrapper">
         <SegmentAsync
           className="query-segment-key"
           Component={plusSegment}
@@ -30,7 +30,7 @@ export const AdHocFilterKey: FC<Props> = ({ datasource, onChange, filterKey }) =
   }
 
   return (
-    <div className="gf-form">
+    <div className="gf-form" data-testid="AdHocFilterKey-key-wrapper">
       <SegmentAsync
         className="query-segment-key"
         value={filterKey}
@@ -46,7 +46,7 @@ export const REMOVE_FILTER_KEY = '-- remove filter --';
 const REMOVE_VALUE = { label: REMOVE_FILTER_KEY, value: REMOVE_FILTER_KEY };
 
 const plusSegment: ReactElement = (
-  <a className="gf-form-label query-part">
+  <a className="gf-form-label query-part" aria-label="Add Filter">
     <Icon name="plus" />
   </a>
 );

public/app/features/variables/adhoc/picker/AdHocFilterValue.tsx

@@ -15,7 +15,7 @@ export const AdHocFilterValue: FC<Props> = ({ datasource, onChange, filterKey, f
   const loadValues = () => fetchFilterValues(datasource, filterKey);
 
   return (
-    <div className="gf-form">
+    <div className="gf-form" data-testid="AdHocFilterValue-value-wrapper">
       <SegmentAsync
         className="query-segment-value"
         placeholder={placeHolder}

public/app/features/variables/adhoc/picker/AdHocPicker.tsx

@@ -1,13 +1,9 @@
-import React, { PureComponent, ReactNode } from 'react';
+import React, { PureComponent } from 'react';
 import { connect, ConnectedProps } from 'react-redux';
 import { AdHocVariableFilter, AdHocVariableModel } from 'app/features/variables/types';
 import { VariablePickerProps } from '../../pickers/types';
-import { SelectableValue } from '@grafana/data';
-import { AdHocFilterBuilder } from './AdHocFilterBuilder';
-import { ConditionSegment } from './ConditionSegment';
 import { addFilter, changeFilter, removeFilter } from '../actions';
-import { REMOVE_FILTER_KEY } from './AdHocFilterKey';
+import { AdHocFilter } from './AdHocFilter';
-import { AdHocFilterRenderer } from './AdHocFilterRenderer';
 
 const mapDispatchToProps = {
   addFilter,
@@ -21,65 +17,37 @@ interface OwnProps extends VariablePickerProps<AdHocVariableModel> {}
 
 type Props = OwnProps & ConnectedProps<typeof connector>;
 
+/**
+ * Thin wrapper over AdHocFilter to add redux actions and change the props so it can be used for ad hoc variable
+ * control.
+ */
 export class AdHocPickerUnconnected extends PureComponent<Props> {
-  onChange = (index: number, prop: string) => (key: SelectableValue<string | null>) => {
+  addFilter = (filter: AdHocVariableFilter) => {
-    const { id, filters } = this.props.variable;
+    this.props.addFilter(this.props.variable.id, filter);
-    const { value } = key;
+  };
 
-    if (key.value === REMOVE_FILTER_KEY) {
+  removeFilter = (index: number) => {
-      return this.props.removeFilter(id, index);
+    this.props.removeFilter(this.props.variable.id, index);
-    }
+  };
 
-    return this.props.changeFilter(id, {
+  changeFilter = (index: number, filter: AdHocVariableFilter) => {
+    this.props.changeFilter(this.props.variable.id, {
       index,
-      filter: {
+      filter,
-        ...filters[index],
-        [prop]: value,
-      },
     });
   };
 
-  appendFilterToVariable = (filter: AdHocVariableFilter) => {
-    const { id } = this.props.variable;
-    this.props.addFilter(id, filter);
-  };
-
   render() {
-    const { filters } = this.props.variable;
+    const { filters, datasource } = this.props.variable;
 
     return (
-      <div className="gf-form-inline">
+      <AdHocFilter
-        {this.renderFilters(filters)}
+        datasource={datasource}
-        <AdHocFilterBuilder
+        filters={filters}
-          datasource={this.props.variable.datasource!}
+        addFilter={this.addFilter}
-          appendBefore={filters.length > 0 ? <ConditionSegment label="AND" /> : null}
+        removeFilter={this.removeFilter}
-          onCompleted={this.appendFilterToVariable}
+        changeFilter={this.changeFilter}
-        />
+      />
-      </div>
-    );
-  }
-
-  renderFilters(filters: AdHocVariableFilter[]) {
-    return filters.reduce((segments: ReactNode[], filter, index) => {
-      if (segments.length > 0) {
-        segments.push(<ConditionSegment label="AND" key={`condition-${index}`} />);
-      }
-      segments.push(this.renderFilterSegments(filter, index));
-      return segments;
-    }, []);
-  }
-
-  renderFilterSegments(filter: AdHocVariableFilter, index: number) {
-    return (
-      <React.Fragment key={`filter-${index}`}>
-        <AdHocFilterRenderer
-          datasource={this.props.variable.datasource!}
-          filter={filter}
-          onKeyChange={this.onChange(index, 'key')}
-          onOperatorChange={this.onChange(index, 'operator')}
-          onValueChange={this.onChange(index, 'value')}
-        />
-      </React.Fragment>
     );
   }
 }

public/app/plugins/datasource/prometheus/datasource.ts

@@ -300,7 +300,7 @@ export class PrometheusDatasource extends DataSourceWithBackend<PromQuery, PromO
   };
 
   shouldRunExemplarQuery(target: PromQuery): boolean {
-    /* We want to run exemplar query only for histogram metrics: 
+    /* We want to run exemplar query only for histogram metrics:
     1. If we haven't processd histogram metrics yet, we need to check if expr includes "_bucket" which means that it is probably histogram metric (can rarely lead to false positive).
     2. If we have processed histogram metrics, check if it is part of query expr.
     */
@@ -790,7 +790,7 @@ export class PrometheusDatasource extends DataSourceWithBackend<PromQuery, PromO
     return result?.data?.data?.map((value: any) => ({ text: value })) ?? [];
   }
 
-  async getTagValues(options: any = {}) {
+  async getTagValues(options: { key?: string } = {}) {
     const result = await this.metadataRequest(`/api/v1/label/${options.key}/values`);
     return result?.data?.data?.map((value: any) => ({ text: value })) ?? [];
   }

public/app/plugins/datasource/tempo/QueryEditor/NativeSearch.tsx

@@ -12,13 +12,13 @@ import {
   Alert,
   useStyles2,
 } from '@grafana/ui';
-import { tokenizer } from './syntax';
+import { tokenizer } from '../syntax';
 import Prism from 'prismjs';
 import { Node } from 'slate';
 import { css } from '@emotion/css';
 import { GrafanaTheme2, isValidGoDuration, SelectableValue } from '@grafana/data';
-import TempoLanguageProvider from './language_provider';
+import TempoLanguageProvider from '../language_provider';
-import { TempoDatasource, TempoQuery } from './datasource';
+import { TempoDatasource, TempoQuery } from '../datasource';
 import { debounce } from 'lodash';
 import { dispatch } from 'app/store/store';
 import { notifyApp } from 'app/core/actions';

public/app/plugins/datasource/tempo/QueryEditor/QueryField.tsx

@@ -1,6 +1,6 @@
 import { css } from '@emotion/css';
-import { DataSourceApi, QueryEditorProps, SelectableValue } from '@grafana/data';
+import { QueryEditorProps, SelectableValue } from '@grafana/data';
-import { config, getDataSourceSrv } from '@grafana/runtime';
+import { config } from '@grafana/runtime';
 import {
   Badge,
   FileDropzone,
@@ -14,13 +14,15 @@ import {
 } from '@grafana/ui';
 import { TraceToLogsOptions } from 'app/core/components/TraceToLogsSettings';
 import React from 'react';
-import { LokiQueryField } from '../loki/components/LokiQueryField';
+import { LokiQueryField } from '../../loki/components/LokiQueryField';
-import { LokiQuery } from '../loki/types';
+import { LokiQuery } from '../../loki/types';
-import { TempoDatasource, TempoQuery, TempoQueryType } from './datasource';
+import { TempoDatasource, TempoQuery, TempoQueryType } from '../datasource';
-import LokiDatasource from '../loki/datasource';
+import LokiDatasource from '../../loki/datasource';
-import { PrometheusDatasource } from '../prometheus/datasource';
+import { PrometheusDatasource } from '../../prometheus/datasource';
 import useAsync from 'react-use/lib/useAsync';
 import NativeSearch from './NativeSearch';
+import { getDS } from './utils';
+import { ServiceGraphSection } from './ServiceGraphSection';
 
 interface Props extends QueryEditorProps<TempoDatasource, TempoQuery>, Themeable2 {}
 
@@ -188,36 +190,14 @@ class TempoQueryFieldComponent extends React.PureComponent<Props, State> {
             </InlineField>
           </InlineFieldRow>
         )}
-        {query.queryType === 'serviceMap' && <ServiceGraphSection graphDatasourceUid={graphDatasourceUid} />}
+        {query.queryType === 'serviceMap' && (
+          <ServiceGraphSection graphDatasourceUid={graphDatasourceUid} query={query} onChange={onChange} />
+        )}
       </>
     );
   }
 }
 
-function ServiceGraphSection({ graphDatasourceUid }: { graphDatasourceUid?: string }) {
-  const dsState = useAsync(() => getDS(graphDatasourceUid), [graphDatasourceUid]);
-  if (dsState.loading) {
-    return null;
-  }
-
-  const ds = dsState.value as LokiDatasource;
-
-  if (!graphDatasourceUid) {
-    return <div className="text-warning">Please set up a service graph datasource in the datasource settings.</div>;
-  }
-
-  if (graphDatasourceUid && !ds) {
-    return (
-      <div className="text-warning">
-        Service graph datasource is configured but the data source no longer exists. Please configure existing data
-        source to use the service graph functionality.
-      </div>
-    );
-  }
-
-  return null;
-}
-
 interface SearchSectionProps {
   linkedDatasourceUid?: string;
   onChange: (value: LokiQuery) => void;
@@ -264,18 +244,4 @@ function SearchSection({ linkedDatasourceUid, onChange, onRunQuery, query }: Sea
   return null;
 }
 
-async function getDS(uid?: string): Promise<DataSourceApi | undefined> {
-  if (!uid) {
-    return undefined;
-  }
-
-  const dsSrv = getDataSourceSrv();
-  try {
-    return await dsSrv.get(uid);
-  } catch (error) {
-    console.error('Failed to load data source', error);
-    return undefined;
-  }
-}
-
 export const TempoQueryField = withTheme2(TempoQueryFieldComponent);

public/app/plugins/datasource/tempo/QueryEditor/ServiceGraphSection.tsx

@@ -0,0 +1,87 @@
+import React from 'react';
+import useAsync from 'react-use/lib/useAsync';
+import { getDS } from './utils';
+import { InlineField, InlineFieldRow } from '@grafana/ui';
+import { AdHocVariableFilter } from '../../../../features/variables/types';
+import { TempoQuery } from '../datasource';
+import { AdHocFilter } from '../../../../features/variables/adhoc/picker/AdHocFilter';
+import { PrometheusDatasource } from '../../prometheus/datasource';
+
+export function ServiceGraphSection({
+  graphDatasourceUid,
+  query,
+  onChange,
+}: {
+  graphDatasourceUid?: string;
+  query: TempoQuery;
+  onChange: (value: TempoQuery) => void;
+}) {
+  const dsState = useAsync(() => getDS(graphDatasourceUid), [graphDatasourceUid]);
+  if (dsState.loading) {
+    return null;
+  }
+
+  const ds = dsState.value as PrometheusDatasource;
+
+  if (!graphDatasourceUid) {
+    return <div className="text-warning">Please set up a service graph datasource in the datasource settings.</div>;
+  }
+
+  if (graphDatasourceUid && !ds) {
+    return (
+      <div className="text-warning">
+        Service graph datasource is configured but the data source no longer exists. Please configure existing data
+        source to use the service graph functionality.
+      </div>
+    );
+  }
+  const filters = queryToFilter(query.serviceMapQuery || '');
+
+  return (
+    <div>
+      <InlineFieldRow>
+        <InlineField label="Filter" labelWidth={14} grow>
+          <AdHocFilter
+            datasource={{ uid: graphDatasourceUid }}
+            filters={filters}
+            addFilter={(filter: AdHocVariableFilter) => {
+              onChange({
+                ...query,
+                serviceMapQuery: filtersToQuery([...filters, filter]),
+              });
+            }}
+            removeFilter={(index: number) => {
+              const newFilters = [...filters];
+              newFilters.splice(index, 1);
+              onChange({ ...query, serviceMapQuery: filtersToQuery(newFilters) });
+            }}
+            changeFilter={(index: number, filter: AdHocVariableFilter) => {
+              const newFilters = [...filters];
+              newFilters.splice(index, 1, filter);
+              onChange({ ...query, serviceMapQuery: filtersToQuery(newFilters) });
+            }}
+          />
+        </InlineField>
+      </InlineFieldRow>
+    </div>
+  );
+}
+
+function queryToFilter(query: string): AdHocVariableFilter[] {
+  let match;
+  let filters: AdHocVariableFilter[] = [];
+  const re = /([\w_]+)(=|!=|<|>|=~|!~)"(.*?)"/g;
+  while ((match = re.exec(query)) !== null) {
+    filters.push({
+      key: match[1],
+      operator: match[2],
+      value: match[3],
+      condition: '',
+    });
+  }
+  return filters;
+}
+
+function filtersToQuery(filters: AdHocVariableFilter[]): string {
+  return `{${filters.map((f) => `${f.key}${f.operator}"${f.value}"`).join(',')}}`;
+}

public/app/plugins/datasource/tempo/QueryEditor/utils.ts

@@ -0,0 +1,16 @@
+import { DataSourceApi } from '@grafana/data';
+import { getDataSourceSrv } from '@grafana/runtime';
+
+export async function getDS(uid?: string): Promise<DataSourceApi | undefined> {
+  if (!uid) {
+    return undefined;
+  }
+
+  const dsSrv = getDataSourceSrv();
+  try {
+    return await dsSrv.get(uid);
+  } catch (error) {
+    console.error('Failed to load data source', error);
+    return undefined;
+  }
+}

public/app/plugins/datasource/tempo/datasource.ts

@@ -54,6 +54,7 @@ export type TempoQuery = {
   minDuration?: string;
   maxDuration?: string;
   limit?: number;
+  serviceMapQuery?: string;
 } & DataQuery;
 
 export const DEFAULT_LIMIT = 20;
@@ -268,6 +269,16 @@ export class TempoDatasource extends DataSourceWithBackend<TempoQuery, TempoJson
     const tagsQueryObject = tagsQuery.reduce((tagQuery, item) => ({ ...tagQuery, ...item }), {});
     return { ...tagsQueryObject, ...tempoQuery };
   }
+
+  async getServiceGraphLabels() {
+    const ds = await getDatasourceSrv().get(this.serviceMap!.datasourceUid);
+    return ds.getTagKeys!();
+  }
+
+  async getServiceGraphLabelValues(key: string) {
+    const ds = await getDatasourceSrv().get(this.serviceMap!.datasourceUid);
+    return ds.getTagValues!({ key });
+  }
 }
 
 function queryServiceMapPrometheus(request: DataQueryRequest<PromQuery>, datasourceUid: string) {
@@ -324,7 +335,9 @@ function makePromServiceMapRequest(options: DataQueryRequest<TempoQuery>): DataQ
     targets: serviceMapMetrics.map((metric) => {
       return {
         refId: metric,
-        expr: `delta(${metric}[$__range])`,
+        // options.targets[0] is not correct here, but not sure what should happen if you have multiple queries for
+        // service map at the same time anyway
+        expr: `delta(${metric}${options.targets[0].serviceMapQuery || ''}[$__range])`,
         instant: true,
       };
     }),

public/app/plugins/datasource/tempo/module.ts

@@ -2,7 +2,7 @@ import { DataSourcePlugin } from '@grafana/data';
 import CheatSheet from './CheatSheet';
 import { ConfigEditor } from './configuration/ConfigEditor';
 import { TempoDatasource } from './datasource';
-import { TempoQueryField } from './QueryField';
+import { TempoQueryField } from './QueryEditor/QueryField';
 
 export const plugin = new DataSourcePlugin(TempoDatasource)
   .setConfigEditor(ConfigEditor)

public/app/plugins/panel/nodeGraph/layout.test.ts

@@ -0,0 +1,88 @@
+import { renderHook } from '@testing-library/react-hooks';
+import { useLayout } from './layout';
+import { EdgeDatum, NodeDatum } from './types';
+
+let onmessage: jest.MockedFunction<any>;
+let postMessage: jest.MockedFunction<any>;
+let terminate: jest.MockedFunction<any>;
+
+jest.mock('./createLayoutWorker', () => {
+  return {
+    __esModule: true,
+    createWorker: () => {
+      onmessage = jest.fn();
+      postMessage = jest.fn();
+      terminate = jest.fn();
+      return {
+        onmessage: onmessage,
+        postMessage: postMessage,
+        terminate: terminate,
+      };
+    },
+  };
+});
+
+describe('layout', () => {
+  it('doesnt fail without any data', async () => {
+    const nodes: NodeDatum[] = [];
+    const edges: EdgeDatum[] = [];
+
+    const { result } = renderHook(() => {
+      return useLayout(nodes, edges, undefined, 100, 1000);
+    });
+    expect(result.current.nodes).toEqual([]);
+    expect(result.current.edges).toEqual([]);
+    expect(postMessage).toBeUndefined();
+  });
+
+  it('cancels worker', async () => {
+    const { result, rerender } = renderHook(
+      ({ nodes, edges }) => {
+        return useLayout(nodes, edges, undefined, 100, 1000);
+      },
+      {
+        initialProps: {
+          nodes: [makeNode(0, 0), makeNode(1, 1)],
+          edges: [makeEdge(0, 1)],
+        },
+      }
+    );
+    expect(postMessage).toBeCalledTimes(1);
+    // Bit convoluted but we cannot easily access the worker instance as we only export constructor so the default
+    // export is class and we only store latest instance of the methods as jest.fn here as module local variables.
+    // So we capture the terminate function from current worker so that when we call rerender and new worker is created
+    // we can still access and check the method from the old one that we assume should be canceled.
+    const localTerminate = terminate;
+
+    rerender({
+      nodes: [],
+      edges: [],
+    });
+
+    expect(result.current.nodes).toEqual([]);
+    expect(result.current.edges).toEqual([]);
+    expect(localTerminate).toBeCalledTimes(1);
+  });
+});
+
+function makeNode(index: number, incoming: number): NodeDatum {
+  return {
+    id: `n${index}`,
+    title: `n${index}`,
+    subTitle: '',
+    dataFrameRowIndex: 0,
+    incoming,
+    arcSections: [],
+  };
+}
+
+function makeEdge(source: number, target: number): EdgeDatum {
+  return {
+    id: `${source}-${target}`,
+    source: 'n' + source,
+    target: 'n' + target,
+    mainStat: '',
+    secondaryStat: '',
+    dataFrameRowIndex: 0,
+  };
+}

public/app/plugins/panel/nodeGraph/layout.ts

@@ -1,10 +1,11 @@
-import { useEffect, useMemo, useState } from 'react';
+import { useEffect, useMemo, useRef, useState } from 'react';
 import { EdgeDatum, EdgeDatumLayout, NodeDatum } from './types';
 import { Field } from '@grafana/data';
 import { useNodeLimit } from './useNodeLimit';
 import useMountedState from 'react-use/lib/useMountedState';
 import { graphBounds } from './utils';
 import { createWorker } from './createLayoutWorker';
+import { useUnmount } from 'react-use';
 
 export interface Config {
   linkDistance: number;
@@ -52,6 +53,13 @@ export function useLayout(
   const [loading, setLoading] = useState(false);
 
   const isMounted = useMountedState();
+  const layoutWorkerCancelRef = useRef<(() => void) | undefined>();
+
+  useUnmount(() => {
+    if (layoutWorkerCancelRef.current) {
+      layoutWorkerCancelRef.current();
+    }
+  });
 
   // Also we compute both layouts here. Grid layout should not add much time and we can more easily just cache both
   // so this should happen only once for a given response data.
@@ -69,6 +77,7 @@ export function useLayout(
     if (rawNodes.length === 0) {
       setNodesGraph([]);
       setEdgesGraph([]);
+      setLoading(false);
       return;
     }
 
@@ -76,14 +85,15 @@ export function useLayout(
 
     // This is async but as I wanted to still run the sync grid layout and you cannot return promise from effect so
     // having callback seems ok here.
-    defaultLayout(rawNodes, rawEdges, ({ nodes, edges }) => {
+    const cancel = defaultLayout(rawNodes, rawEdges, ({ nodes, edges }) => {
-      // TODO: it would be better to cancel the worker somehow but probably not super important right now.
       if (isMounted()) {
         setNodesGraph(nodes);
         setEdgesGraph(edges as EdgeDatumLayout[]);
         setLoading(false);
       }
     });
+    layoutWorkerCancelRef.current = cancel;
+    return cancel;
   }, [rawNodes, rawEdges, isMounted]);
 
   // Compute grid separately as it is sync and do not need to be inside effect. Also it is dependant on width while
@@ -128,6 +138,7 @@ export function useLayout(
 
 /**
  * Wraps the layout code in a worker as it can take long and we don't want to block the main thread.
+ * Returns a cancel function to terminate the worker.
  */
 function defaultLayout(
   nodes: NodeDatum[],
@@ -154,6 +165,10 @@ function defaultLayout(
     edges,
     config: defaultConfig,
   });
+
+  return () => {
+    worker.terminate();
+  };
 }
 
 /**

public/test/mocks/workers.ts

@@ -1,15 +1,21 @@
 const { layout } = jest.requireActual('../../app/plugins/panel/nodeGraph/layout.worker.js');
 
 class LayoutMockWorker {
+  timeout: number | undefined;
   constructor() {}
-
   postMessage(data: any) {
     const { nodes, edges, config } = data;
-    setTimeout(() => {
+    this.timeout = setTimeout(() => {
+      this.timeout = undefined;
       layout(nodes, edges, config);
       // @ts-ignore
       this.onmessage({ data: { nodes, edges } });
-    }, 1);
+    }, 1) as any;
+  }
+  terminate() {
+    if (this.timeout) {
+      clearTimeout(this.timeout);
+    }
   }
 }