[ML] Explain log rate spikes: Fix race conditions when loading document count stats.

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;
 }