[ML] Explain log rate spikes: Fix race conditions when loading docume…

…nt count stats. (#137715) (#137817)

- Fixes a race condition when loading document count stats for the histogram charts. Previously, data for the splitted histogram were loaded with to custom hooks in parallel which meant data in the chart could end up out of sync when transitioning from one state to the other.
- Fixes another race condition where stale data fetched via hovering would populate the chart late. This was fixed by adding abort signals to the data fetching hooks.
- Adds caching to the data fetched via hovering the analysis table rows.

(cherry picked from commit 24de690)

Co-authored-by: Walter Rafelsberger <[email protected]>

x-pack/plugins/aiops/public/components/explain_log_rate_spikes/explain_log_rate_spikes_page.tsx

@@ -106,8 +106,7 @@ export const ExplainLogRateSpikesPage: FC<ExplainLogRateSpikesPageProps> = ({
   }, [pinnedChangePoint, selectedChangePoint]);
 
   const {
-    overallDocStats,
-    selectedDocStats,
+    documentStats,
     timefilter,
     earliest,
     latest,
@@ -121,9 +120,7 @@ export const ExplainLogRateSpikesPage: FC<ExplainLogRateSpikesPageProps> = ({
     currentSelectedChangePoint
   );
 
-  const totalCount = currentSelectedChangePoint
-    ? overallDocStats.totalCount + selectedDocStats.totalCount
-    : overallDocStats.totalCount;
+  const { totalCount, documentCountStats, documentCountStatsCompare } = documentStats;
 
   useEffect(
     // TODO: Consolidate this hook/function with with Data visualizer's
@@ -221,15 +218,15 @@ export const ExplainLogRateSpikesPage: FC<ExplainLogRateSpikesPageProps> = ({
               setSearchParams={setSearchParams}
             />
           </EuiFlexItem>
-          {overallDocStats?.totalCount !== undefined && (
+          {documentCountStats !== undefined && (
             <EuiFlexItem>
               <EuiPanel paddingSize="m">
                 <DocumentCountContent
                   brushSelectionUpdateHandler={setWindowParameters}
                   clearSelectionHandler={clearSelection}
-                  documentCountStats={overallDocStats.documentCountStats}
+                  documentCountStats={documentCountStats}
                   documentCountStatsSplit={
-                    currentSelectedChangePoint ? selectedDocStats.documentCountStats : undefined
+                    currentSelectedChangePoint ? documentCountStatsCompare : undefined
                   }
                   totalCount={totalCount}
                   changePoint={currentSelectedChangePoint}

x-pack/plugins/aiops/public/hooks/use_data.ts

@@ -108,13 +108,13 @@ export const useData = (
   }, [fieldStatsRequest, selectedChangePoint]);
 
   const selectedChangePointStatsRequest = useMemo(() => {
-    return fieldStatsRequest
+    return fieldStatsRequest && selectedChangePoint
       ? { ...fieldStatsRequest, selectedChangePoint, includeSelectedChangePoint: true }
       : undefined;
   }, [fieldStatsRequest, selectedChangePoint]);
 
-  const { docStats: overallDocStats } = useDocumentCountStats(overallStatsRequest, lastRefresh);
-  const { docStats: selectedDocStats } = useDocumentCountStats(
+  const documentStats = useDocumentCountStats(
+    overallStatsRequest,
     selectedChangePointStatsRequest,
     lastRefresh
   );
@@ -177,8 +177,7 @@ export const useData = (
   }, [searchString, JSON.stringify(searchQuery)]);
 
   return {
-    overallDocStats,
-    selectedDocStats,
+    documentStats,
     timefilter,
     /** Start timestamp filter */
     earliest: fieldStatsRequest?.earliest,

x-pack/plugins/aiops/public/hooks/use_document_count_stats.ts

@@ -5,10 +5,13 @@
  * 2.0.
  */
 
-import { useCallback, useEffect, useState, useMemo } from 'react';
+import { useCallback, useEffect, useRef, useState } from 'react';
 import { lastValueFrom } from 'rxjs';
+
 import { i18n } from '@kbn/i18n';
 import type { ToastsStart } from '@kbn/core/public';
+import { stringHash } from '@kbn/ml-string-hash';
+
 import { useAiOpsKibana } from '../kibana_context';
 import { extractErrorProperties } from '../application/utils/error_utils';
 import {
@@ -21,6 +24,7 @@ import {
 export interface DocumentStats {
   totalCount: number;
   documentCountStats?: DocumentCountStats;
+  documentCountStatsCompare?: DocumentCountStats;
 }
 
 function displayError(toastNotifications: ToastsStart, index: string, err: any) {
@@ -51,52 +55,101 @@ function displayError(toastNotifications: ToastsStart, index: string, err: any)
 
 export function useDocumentCountStats<TParams extends DocumentStatsSearchStrategyParams>(
   searchParams: TParams | undefined,
+  searchParamsCompare: TParams | undefined,
   lastRefresh: number
-): {
-  docStats: DocumentStats;
-} {
+): DocumentStats {
   const {
     services: {
       data,
       notifications: { toasts },
     },
   } = useAiOpsKibana();
 
-  const [stats, setStats] = useState<DocumentStats>({
+  const abortCtrl = useRef(new AbortController());
+
+  const [documentStats, setDocumentStats] = useState<DocumentStats>({
     totalCount: 0,
   });
 
+  const [documentStatsCache, setDocumentStatsCache] = useState<Record<string, DocumentStats>>({});
+
   const fetchDocumentCountData = useCallback(async () => {
     if (!searchParams) return;
 
+    const cacheKey = stringHash(
+      `${JSON.stringify(searchParams)}_${JSON.stringify(searchParamsCompare)}`
+    );
+
+    if (documentStatsCache[cacheKey]) {
+      setDocumentStats(documentStatsCache[cacheKey]);
+      return;
+    }
+
     try {
+      abortCtrl.current = new AbortController();
+
       const resp = await lastValueFrom(
-        data.search.search({
-          params: getDocumentCountStatsRequest(searchParams),
-        })
+        data.search.search(
+          {
+            params: getDocumentCountStatsRequest(searchParams),
+          },
+          { abortSignal: abortCtrl.current.signal }
+        )
       );
+
       const documentCountStats = processDocumentCountStats(resp?.rawResponse, searchParams);
       const totalCount = documentCountStats?.totalCount ?? 0;
-      setStats({
+
+      const newStats: DocumentStats = {
         documentCountStats,
         totalCount,
+      };
+
+      if (searchParamsCompare) {
+        const respCompare = await lastValueFrom(
+          data.search.search(
+            {
+              params: getDocumentCountStatsRequest(searchParamsCompare),
+            },
+            { abortSignal: abortCtrl.current.signal }
+          )
+        );
+
+        const documentCountStatsCompare = processDocumentCountStats(
+          respCompare?.rawResponse,
+          searchParamsCompare
+        );
+        const totalCountCompare = documentCountStatsCompare?.totalCount ?? 0;
+
+        newStats.documentCountStatsCompare = documentCountStatsCompare;
+        newStats.totalCount = totalCount + totalCountCompare;
+      }
+
+      setDocumentStats(newStats);
+      setDocumentStatsCache({
+        ...documentStatsCache,
+        [cacheKey]: newStats,
       });
     } catch (error) {
-      displayError(toasts, searchParams!.index, extractErrorProperties(error));
+      // An `AbortError` gets triggered when a user cancels a request by navigating away, we need to ignore these errors.
+      if (error.name !== 'AbortError') {
+        displayError(toasts, searchParams!.index, extractErrorProperties(error));
+      }
     }
-  }, [data?.search, searchParams, toasts]);
+  }, [data?.search, documentStatsCache, searchParams, searchParamsCompare, toasts]);
 
   useEffect(
     function getDocumentCountData() {
       fetchDocumentCountData();
+      return () => abortCtrl.current.abort();
     },
-    [fetchDocumentCountData]
+    [fetchDocumentCountData, lastRefresh]
   );
 
-  return useMemo(
-    () => ({
-      docStats: stats,
-    }),
-    [stats]
-  );
+  // Clear the document count stats cache when the outer page (date picker/search bar) triggers a refresh.
+  useEffect(() => {
+    setDocumentStatsCache({});
+  }, [lastRefresh]);
+
+  return documentStats;
 }