[Discover] Refactor totalHits$ loading state handling to omit race conditions (#196114)

Fix loading state management in `use_discover_histogram.ts`

Moving the loading state for `totalHits$` to the `fetchAll` function, which is executed before the hook. This ensures that the loading state is set at a higher level, preventing a situation where the overall data fetching is in a `loading` state, but the histogram is marked as `complete` while receiving new properties (like a new data view ID) without access to refreshed data views.

src/plugins/discover/public/application/main/components/layout/use_discover_histogram.ts

@@ -61,7 +61,7 @@ export const useDiscoverHistogram = ({
   hideChart,
 }: UseDiscoverHistogramProps) => {
   const services = useDiscoverServices();
-  const savedSearchData$ = stateContainer.dataState.data$;
+  const { main$, documents$, totalHits$ } = stateContainer.dataState.data$;
   const savedSearchState = useSavedSearch();
   const isEsqlMode = useIsEsqlMode();
 
@@ -153,10 +153,7 @@ export const useDiscoverHistogram = ({
    * Total hits
    */
 
-  const setTotalHitsError = useMemo(
-    () => sendErrorTo(savedSearchData$.totalHits$),
-    [savedSearchData$.totalHits$]
-  );
+  const setTotalHitsError = useMemo(() => sendErrorTo(totalHits$), [totalHits$]);
 
   useEffect(() => {
     const subscription = createTotalHitsObservable(unifiedHistogram?.state$)?.subscribe(
@@ -172,7 +169,7 @@ export const useDiscoverHistogram = ({
           return;
         }
 
-        const { result: totalHitsResult } = savedSearchData$.totalHits$.getValue();
+        const { result: totalHitsResult } = totalHits$.getValue();
 
         if (
           (status === UnifiedHistogramFetchStatus.loading ||
@@ -184,18 +181,22 @@ export const useDiscoverHistogram = ({
           return;
         }
 
-        // Sync the totalHits$ observable with the unified histogram state
-        savedSearchData$.totalHits$.next({
-          fetchStatus: status.toString() as FetchStatus,
-          result,
-        });
+        const fetchStatus = status.toString() as FetchStatus;
+
+        // Do not sync the loading state since it's already handled by fetchAll
+        if (fetchStatus !== FetchStatus.LOADING) {
+          totalHits$.next({
+            fetchStatus,
+            result,
+          });
+        }
 
         if (status !== UnifiedHistogramFetchStatus.complete || typeof result !== 'number') {
           return;
         }
 
         // Check the hits count to set a partial or no results state
-        checkHitCount(savedSearchData$.main$, result);
+        checkHitCount(main$, result);
       }
     );
 
@@ -204,8 +205,8 @@ export const useDiscoverHistogram = ({
     };
   }, [
     isEsqlMode,
-    savedSearchData$.main$,
-    savedSearchData$.totalHits$,
+    main$,
+    totalHits$,
     setTotalHitsError,
     stateContainer.appState,
     unifiedHistogram?.state$,
@@ -234,7 +235,7 @@ export const useDiscoverHistogram = ({
 
   const [initialEsqlProps] = useState(() =>
     getUnifiedHistogramPropsForEsql({
-      documentsValue: savedSearchData$.documents$.getValue(),
+      documentsValue: documents$.getValue(),
       savedSearch: stateContainer.savedSearchState.getState(),
     })
   );

src/plugins/discover/public/application/main/data_fetching/fetch_all.ts

@@ -92,13 +92,9 @@ export function fetchAll(
     // Mark all subjects as loading
     sendLoadingMsg(dataSubjects.main$);
     sendLoadingMsg(dataSubjects.documents$, { query });
-
-    // histogram for data view mode will send `loading` for totalHits$
-    if (isEsqlQuery) {
-      sendLoadingMsg(dataSubjects.totalHits$, {
-        result: dataSubjects.totalHits$.getValue().result,
-      });
-    }
+    sendLoadingMsg(dataSubjects.totalHits$, {
+      result: dataSubjects.totalHits$.getValue().result,
+    });
 
     // Start fetching all required requests
     const response = isEsqlQuery

src/plugins/discover/public/application/main/state_management/discover_data_state_container.test.ts

@@ -159,7 +159,6 @@ describe('test getDataStateContainer', () => {
         expect(
           stateContainer.searchSessionManager.getCurrentSearchSessionId as jest.Mock
         ).toHaveBeenCalled();
-
         unsubscribe();
         done();
       }
@@ -169,21 +168,24 @@ describe('test getDataStateContainer', () => {
   });
 
   it('should update app state from default profile state', async () => {
+    mockFetchDocuments.mockResolvedValue({ records: [] });
     const stateContainer = getDiscoverStateMock({ isTimeBased: true });
     const dataState = stateContainer.dataState;
     const dataUnsub = dataState.subscribe();
     const appUnsub = stateContainer.appState.initAndSync();
-    discoverServiceMock.profilesManager.resolveDataSourceProfile({});
+    await discoverServiceMock.profilesManager.resolveDataSourceProfile({});
     stateContainer.actions.setDataView(dataViewMock);
     stateContainer.internalState.transitions.setResetDefaultProfileState({
       columns: true,
       rowHeight: true,
     });
+
     dataState.data$.totalHits$.next({
       fetchStatus: FetchStatus.COMPLETE,
       result: 0,
     });
     dataState.refetch$.next(undefined);
+
     await waitFor(() => {
       expect(dataState.data$.main$.value.fetchStatus).toBe(FetchStatus.COMPLETE);
     });
@@ -202,7 +204,7 @@ describe('test getDataStateContainer', () => {
     const dataState = stateContainer.dataState;
     const dataUnsub = dataState.subscribe();
     const appUnsub = stateContainer.appState.initAndSync();
-    discoverServiceMock.profilesManager.resolveDataSourceProfile({});
+    await discoverServiceMock.profilesManager.resolveDataSourceProfile({});
     stateContainer.actions.setDataView(dataViewMock);
     stateContainer.internalState.transitions.setResetDefaultProfileState({
       columns: false,

src/plugins/discover/public/application/main/state_management/discover_data_state_container.ts

@@ -185,6 +185,12 @@ export function getDataStateContainer({
     documents$: new BehaviorSubject<DataDocumentsMsg>(initialState),
     totalHits$: new BehaviorSubject<DataTotalHitsMsg>(initialState),
   };
+  // This is debugging code, helping you to understand which messages are sent to the data observables
+  // Adding a debugger in the functions can be helpful to understand what triggers a message
+  // dataSubjects.main$.subscribe((msg) => addLog('dataSubjects.main$', msg));
+  // dataSubjects.documents$.subscribe((msg) => addLog('dataSubjects.documents$', msg));
+  // dataSubjects.totalHits$.subscribe((msg) => addLog('dataSubjects.totalHits$', msg););
+  // Add window.ELASTIC_DISCOVER_LOGGER = 'debug' to see messages in console
 
   let autoRefreshDone: AutoRefreshDoneFn | undefined | null = null;
   /**

test/functional/apps/discover/group3/_lens_vis.ts

@@ -110,8 +110,7 @@ export default function ({ getService, getPageObjects }: FtrProviderContext) {
     return seriesType;
   }
 
-  // Failing: See https://github.com/elastic/kibana/issues/184600
-  describe.skip('discover lens vis', function () {
+  describe('discover lens vis', function () {
     before(async () => {
       await security.testUser.setRoles(['kibana_admin', 'test_logstash_reader']);
       await esArchiver.loadIfNeeded('test/functional/fixtures/es_archiver/logstash_functional');