[ENHANCEMENT] Keep usage metrics state across renders

plugin-system/src/runtime/DataQueriesProvider/DataQueriesProvider.tsx

@@ -13,7 +13,7 @@
 
 import type { QueryType, TimeSeriesQueryDefinition } from '@perses-dev/spec';
 import type { ReactElement } from 'react';
-import { createContext, useCallback, useContext, useMemo } from 'react';
+import { createContext, useCallback, useContext, useEffect, useMemo } from 'react';
 
 import type { AlertsQueryDefinition } from '../alerts-queries';
 import { useAlertsQueries } from '../alerts-queries';
@@ -123,18 +123,6 @@ export function DataQueriesProvider(props: DataQueriesProviderProps): ReactEleme
       ...transformQueryResults(jsonResults, jsonQueries),
     ];
 
-    if (queryOptions?.enabled) {
-      for (const result of mergedQueryResults) {
-        if (!result.isLoading && !result.isFetching && !result.error) {
-          usageMetrics.markQuery(result.definition, 'success');
-        } else if (result.error) {
-          usageMetrics.markQuery(result.definition, 'error');
-        } else {
-          usageMetrics.markQuery(result.definition, 'pending');
-        }
-      }
-    }
-
     return {
       queryDefinitions: definitions,
       queryResults: mergedQueryResults,
@@ -159,10 +147,26 @@ export function DataQueriesProvider(props: DataQueriesProviderProps): ReactEleme
     traceQueries,
     traceResults,
     definitions,
-    queryOptions?.enabled,
     refetchAll,
-    usageMetrics,
   ]);
 
+  useEffect(() => {
+    if (queryOptions?.enabled) {
+      // Register the complete batch before completing any cached query.
+      for (const result of ctx.queryResults) {
+        usageMetrics.markQuery(result.definition, 'pending');
+      }
+      for (const result of ctx.queryResults) {
+        if (!result.isLoading && !result.isFetching && !result.error) {
+          usageMetrics.markQuery(result.definition, 'success');
+        } else if (result.error) {
+          usageMetrics.markQuery(result.definition, 'error');
+        } else {
+          usageMetrics.markQuery(result.definition, 'pending');
+        }
+      }
+    }
+  }, [ctx.queryResults, queryOptions?.enabled, usageMetrics]);
+
   return <DataQueriesContext.Provider value={ctx}>{children}</DataQueriesContext.Provider>;
 }

plugin-system/src/runtime/DataQueriesProvider/DataQueriesProvider.usage-metrics.test.tsx

@@ -0,0 +1,102 @@
+// Copyright The Perses Authors
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import { FetchProvider } from '@perses-dev/client';
+import type { QueryDefinition } from '@perses-dev/spec';
+import { render } from '@testing-library/react';
+import type { ReactElement } from 'react';
+import { StrictMode, Suspense } from 'react';
+
+import { UsageMetricsContext, UsageMetricsProvider } from '../UsageMetricsProvider';
+import { DataQueriesProvider } from './DataQueriesProvider';
+
+const mocks = vi.hoisted(() => ({
+  timeSeriesResults: [] as unknown[],
+  noResults: [] as unknown[],
+}));
+
+vi.mock('../time-series-queries', () => ({ useTimeSeriesQueries: (): unknown[] => mocks.timeSeriesResults }));
+vi.mock('../trace-queries', () => ({ useTraceQueries: (): unknown[] => mocks.noResults }));
+vi.mock('../profile-queries', () => ({ useProfileQueries: (): unknown[] => mocks.noResults }));
+vi.mock('../log-queries', () => ({ useLogQueries: (): unknown[] => mocks.noResults }));
+vi.mock('../alerts-queries', () => ({ useAlertsQueries: (): unknown[] => mocks.noResults }));
+vi.mock('../silences-queries', () => ({ useSilencesQueries: (): unknown[] => mocks.noResults }));
+vi.mock('../json-queries', () => ({ useJsonQueries: (): unknown[] => mocks.noResults }));
+
+const cachedQuery: QueryDefinition = { kind: 'TimeSeriesQuery', spec: { plugin: { kind: 'Cached', spec: {} } } };
+const slowQuery: QueryDefinition = { kind: 'TimeSeriesQuery', spec: { plugin: { kind: 'Slow', spec: {} } } };
+const definitions = [cachedQuery, slowQuery];
+const queryOptions = { enabled: true };
+
+const done = { data: {}, isFetching: false, isLoading: false, error: null };
+const loading = { data: undefined, isFetching: true, isLoading: true, error: null };
+
+const markQuery = vi.fn();
+const usageMetrics = { project: 'project', dashboard: 'dashboard', markQuery };
+const neverResolves = new Promise<never>(() => undefined);
+function Suspend(): ReactElement {
+  throw neverResolves;
+}
+
+describe('DataQueriesProvider usage metrics', () => {
+  afterEach(() => {
+    mocks.timeSeriesResults = [];
+    vi.restoreAllMocks();
+  });
+
+  it('waits for the whole batch when a cached query completes first', () => {
+    vi.spyOn(Date, 'now').mockReturnValue(1000);
+    const fetchFn = vi.fn().mockResolvedValue({ ok: true });
+    function Dashboard(): ReactElement {
+      return (
+        <StrictMode>
+          <FetchProvider fetchFn={fetchFn}>
+            <UsageMetricsProvider project="project" dashboard="dashboard">
+              <DataQueriesProvider definitions={definitions} queryOptions={queryOptions} />
+            </UsageMetricsProvider>
+          </FetchProvider>
+        </StrictMode>
+      );
+    }
+
+    mocks.timeSeriesResults = [done, loading];
+    const { rerender } = render(<Dashboard />);
+    expect(fetchFn).not.toHaveBeenCalled();
+
+    mocks.timeSeriesResults = [done, { ...done, data: undefined, error: new Error('boom') }];
+    rerender(<Dashboard />);
+    expect(fetchFn).toHaveBeenCalledTimes(1);
+    expect(fetchFn).toHaveBeenCalledWith(
+      '/api/v1/view',
+      expect.objectContaining({
+        body: JSON.stringify({ project: 'project', dashboard: 'dashboard', render_time: 0, render_errors: 1 }),
+      }),
+    );
+  });
+
+  it('does not record queries from a render that is never committed', () => {
+    markQuery.mockClear();
+    mocks.timeSeriesResults = [done, loading];
+    render(
+      <UsageMetricsContext.Provider value={usageMetrics}>
+        <Suspense fallback={null}>
+          <DataQueriesProvider definitions={definitions} queryOptions={queryOptions}>
+            <Suspend />
+          </DataQueriesProvider>
+        </Suspense>
+      </UsageMetricsContext.Provider>,
+    );
+
+    expect(markQuery).not.toHaveBeenCalled();
+  });
+});

plugin-system/src/runtime/UsageMetricsProvider.test.tsx

@@ -0,0 +1,108 @@
+// Copyright The Perses Authors
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import { FetchProvider } from '@perses-dev/client';
+import type { QueryDefinition } from '@perses-dev/spec';
+import { act, renderHook } from '@testing-library/react';
+import type { PropsWithChildren, ReactElement } from 'react';
+import { StrictMode, useMemo } from 'react';
+
+import { UsageMetricsContext, UsageMetricsProvider, useUsageMetrics } from './UsageMetricsProvider';
+
+const query: QueryDefinition = { kind: 'TimeSeriesQuery', spec: { plugin: { kind: 'TestQuery', spec: {} } } };
+const secondQuery: QueryDefinition = { ...query, spec: { plugin: { kind: 'OtherQuery', spec: {} } } };
+
+describe('UsageMetricsProvider', () => {
+  it('keeps pending queries across rerenders and submits only once, including zero-duration renders', () => {
+    const fetchFn = vi.fn().mockResolvedValue({ ok: true });
+    vi.spyOn(Date, 'now').mockReturnValue(1000);
+    function Wrapper({ children }: PropsWithChildren): ReactElement {
+      return (
+        <StrictMode>
+          <FetchProvider fetchFn={fetchFn}>
+            <UsageMetricsProvider project="project" dashboard="dashboard" apiPrefix="/prefix">
+              {children}
+            </UsageMetricsProvider>
+          </FetchProvider>
+        </StrictMode>
+      );
+    }
+    try {
+      const { result, rerender } = renderHook(() => useUsageMetrics(), { wrapper: Wrapper });
+      expect(fetchFn).not.toHaveBeenCalled();
+      act(() => {
+        result.current.markQuery(query, 'pending');
+        result.current.markQuery(secondQuery, 'pending');
+        result.current.markQuery(query, 'success');
+      });
+      rerender();
+      expect(fetchFn).not.toHaveBeenCalled();
+      act(() => result.current.markQuery(secondQuery, 'error'));
+      expect(fetchFn).toHaveBeenCalledTimes(1);
+      expect(fetchFn).toHaveBeenCalledWith(
+        '/prefix/api/v1/view',
+        expect.objectContaining({
+          method: 'POST',
+          body: JSON.stringify({ project: 'project', dashboard: 'dashboard', render_time: 0, render_errors: 1 }),
+        }),
+      );
+      act(() => {
+        result.current.markQuery(secondQuery, 'pending');
+        result.current.markQuery(secondQuery, 'success');
+      });
+      expect(fetchFn).toHaveBeenCalledTimes(1);
+    } finally {
+      vi.restoreAllMocks();
+    }
+  });
+
+  it('logs a warning when the metrics cannot be submitted', async () => {
+    const error = new Error('network down');
+    const fetchFn = vi.fn().mockRejectedValue(error);
+    const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined);
+    function Wrapper({ children }: PropsWithChildren): ReactElement {
+      return (
+        <FetchProvider fetchFn={fetchFn}>
+          <UsageMetricsProvider project="project" dashboard="dashboard">
+            {children}
+          </UsageMetricsProvider>
+        </FetchProvider>
+      );
+    }
+    try {
+      const { result } = renderHook(() => useUsageMetrics(), { wrapper: Wrapper });
+      await act(async () => result.current.markQuery(query, 'success'));
+      expect(fetchFn).toHaveBeenCalledTimes(1);
+      expect(warn).toHaveBeenCalledWith('Failed to submit dashboard usage metrics', error);
+    } finally {
+      vi.restoreAllMocks();
+    }
+  });
+
+  it('supports directly supplied context values', () => {
+    const markQuery = vi.fn();
+    const stats = { project: 'project', dashboard: 'dashboard', markQuery };
+    function Wrapper({ children }: PropsWithChildren): ReactElement {
+      const value = useMemo(() => stats, []);
+      return <UsageMetricsContext.Provider value={value}>{children}</UsageMetricsContext.Provider>;
+    }
+    const { result } = renderHook(() => useUsageMetrics(), { wrapper: Wrapper });
+    act(() => result.current.markQuery(query, 'success'));
+    expect(markQuery).toHaveBeenCalledWith(query, 'success');
+  });
+
+  it('is safe without a provider', () => {
+    const { result } = renderHook(() => useUsageMetrics());
+    expect(() => result.current.markQuery(query, 'success')).not.toThrow();
+  });
+});