[ML] Fix race condition when updating data grid row count (#149518)

This fixes an edges case where handling the row count information (the
actual count and the relation (`eq/gte`)) could trigger two independent
state updates. Because of those multiple state updates the check whether
to show the mini histograms by default could fail. For example updating
from `count: 0, relation: eq` to `count: 10000, relation: gte` could
cause an intermediate state `count: 10000, relation, eq` which would
pass the check to show the mini charts. The fix here is to combine the
two settings to use only one `useState()`.

x-pack/plugins/ml/public/application/components/data_grid/common.ts

@@ -46,7 +46,7 @@ import { RuntimeMappings } from '../../../../common/types/fields';
 import { isRuntimeMappings } from '../../../../common/util/runtime_field_utils';
 
 export const INIT_MAX_COLUMNS = 10;
-export const COLUMN_CHART_DEFAULT_VISIBILITY_ROWS_THRESHOLED = 10000;
+export const COLUMN_CHART_DEFAULT_VISIBILITY_ROWS_THRESHOLD = 10000;
 
 export const euiDataGridStyle: EuiDataGridStyle = {
   border: 'all',

x-pack/plugins/ml/public/application/components/data_grid/types.ts

@@ -29,6 +29,11 @@ export type DataGridItem = Record<string, any>;
 export type ChartsVisible = boolean | undefined;
 export type RowCountRelation = estypes.SearchTotalHitsRelation | undefined;
 
+export interface RowCountInfo {
+  rowCount: number;
+  rowCountRelation: RowCountRelation;
+}
+
 export type IndexPagination = Pick<EuiDataGridPaginationProps, 'pageIndex' | 'pageSize'>;
 
 export type OnChangeItemsPerPage = (pageSize: any) => void;
@@ -105,8 +110,7 @@ export interface UseDataGridReturnType {
   setErrorMessage: Dispatch<SetStateAction<string>>;
   setNoDataMessage: Dispatch<SetStateAction<string>>;
   setPagination: Dispatch<SetStateAction<IndexPagination>>;
-  setRowCount: Dispatch<SetStateAction<number>>;
-  setRowCountRelation: Dispatch<SetStateAction<RowCountRelation>>;
+  setRowCountInfo: Dispatch<SetStateAction<RowCountInfo>>;
   setSortingColumns: Dispatch<SetStateAction<EuiDataGridSorting['columns']>>;
   setStatus: Dispatch<SetStateAction<INDEX_STATUS>>;
   setTableItems: Dispatch<SetStateAction<DataGridItem[]>>;

x-pack/plugins/ml/public/application/components/data_grid/use_data_grid.tsx

@@ -15,7 +15,7 @@ import { ChartData } from '../../../../common/types/field_histograms';
 import { INDEX_STATUS } from '../../data_frame_analytics/common';
 
 import { ColumnChart } from './column_chart';
-import { COLUMN_CHART_DEFAULT_VISIBILITY_ROWS_THRESHOLED, INIT_MAX_COLUMNS } from './common';
+import { COLUMN_CHART_DEFAULT_VISIBILITY_ROWS_THRESHOLD, INIT_MAX_COLUMNS } from './common';
 import {
   ChartsVisible,
   ColumnId,
@@ -24,10 +24,15 @@ import {
   OnChangeItemsPerPage,
   OnChangePage,
   OnSort,
-  RowCountRelation,
+  RowCountInfo,
   UseDataGridReturnType,
 } from './types';
 
+const rowCountDefault: RowCountInfo = {
+  rowCount: 0,
+  rowCountRelation: undefined,
+};
+
 export const useDataGrid = (
   columns: EuiDataGridColumn[],
   defaultPageSize = 5,
@@ -40,14 +45,15 @@ export const useDataGrid = (
   const [noDataMessage, setNoDataMessage] = useState('');
   const [errorMessage, setErrorMessage] = useState('');
   const [status, setStatus] = useState(INDEX_STATUS.UNUSED);
-  const [rowCount, setRowCount] = useState(0);
-  const [rowCountRelation, setRowCountRelation] = useState<RowCountRelation>(undefined);
+  const [rowCountInfo, setRowCountInfo] = useState<RowCountInfo>(rowCountDefault);
   const [columnCharts, setColumnCharts] = useState<ChartData[]>([]);
   const [tableItems, setTableItems] = useState<DataGridItem[]>([]);
   const [pagination, setPagination] = useState(defaultPagination);
   const [sortingColumns, setSortingColumns] = useState<EuiDataGridSorting['columns']>([]);
   const [chartsVisible, setChartsVisible] = useState<ChartsVisible>(undefined);
 
+  const { rowCount, rowCountRelation } = rowCountInfo;
+
   const toggleChartVisibility = () => {
     if (chartsVisible !== undefined) {
       setChartsVisible(!chartsVisible);
@@ -148,7 +154,7 @@ export const useDataGrid = (
   useEffect(() => {
     if (chartsVisible === undefined && rowCount > 0 && rowCountRelation !== undefined) {
       setChartsVisible(
-        rowCount <= COLUMN_CHART_DEFAULT_VISIBILITY_ROWS_THRESHOLED &&
+        rowCount <= COLUMN_CHART_DEFAULT_VISIBILITY_ROWS_THRESHOLD &&
           rowCountRelation !== ES_CLIENT_TOTAL_HITS_RELATION.GTE
       );
     }
@@ -174,8 +180,7 @@ export const useDataGrid = (
     setErrorMessage,
     setNoDataMessage,
     setPagination,
-    setRowCount,
-    setRowCountRelation,
+    setRowCountInfo,
     setSortingColumns,
     setStatus,
     setTableItems,

x-pack/plugins/ml/public/application/data_frame_analytics/common/get_index_data.ts

@@ -27,8 +27,7 @@ export const getIndexData = async (
     const {
       pagination,
       setErrorMessage,
-      setRowCount,
-      setRowCountRelation,
+      setRowCountInfo,
       setStatus,
       setTableItems,
       sortingColumns,
@@ -64,12 +63,13 @@ export const getIndexData = async (
       });
 
       if (!options.didCancel) {
-        setRowCount(typeof resp.hits.total === 'number' ? resp.hits.total : resp.hits.total!.value);
-        setRowCountRelation(
-          typeof resp.hits.total === 'number'
-            ? ('eq' as estypes.SearchTotalHitsRelation)
-            : resp.hits.total!.relation
-        );
+        setRowCountInfo({
+          rowCount: typeof resp.hits.total === 'number' ? resp.hits.total : resp.hits.total!.value,
+          rowCountRelation:
+            typeof resp.hits.total === 'number'
+              ? ('eq' as estypes.SearchTotalHitsRelation)
+              : resp.hits.total!.relation,
+        });
         setTableItems(
           resp.hits.hits.map((d) =>
             getProcessedFields(

x-pack/plugins/ml/public/application/data_frame_analytics/pages/analytics_creation/hooks/use_index_data.ts

@@ -152,8 +152,7 @@ export const useIndexData = (
     pagination,
     resetPagination,
     setErrorMessage,
-    setRowCount,
-    setRowCountRelation,
+    setRowCountInfo,
     setStatus,
     setTableItems,
     sortingColumns,
@@ -199,12 +198,13 @@ export const useIndexData = (
         const resp: IndexSearchResponse = await ml.esSearch(esSearchRequest);
         const docs = resp.hits.hits.map((d) => getProcessedFields(d.fields ?? {}));
 
-        setRowCount(typeof resp.hits.total === 'number' ? resp.hits.total : resp.hits.total!.value);
-        setRowCountRelation(
-          typeof resp.hits.total === 'number'
-            ? ('eq' as estypes.SearchTotalHitsRelation)
-            : resp.hits.total!.relation
-        );
+        setRowCountInfo({
+          rowCount: typeof resp.hits.total === 'number' ? resp.hits.total : resp.hits.total!.value,
+          rowCountRelation:
+            typeof resp.hits.total === 'number'
+              ? ('eq' as estypes.SearchTotalHitsRelation)
+              : resp.hits.total!.relation,
+        });
         setTableItems(docs);
         setStatus(INDEX_STATUS.LOADED);
       } catch (e) {

x-pack/plugins/transform/public/app/hooks/use_index_data.ts

@@ -154,8 +154,7 @@ export const useIndexData = (
     setColumnCharts,
     setCcsWarning,
     setErrorMessage,
-    setRowCount,
-    setRowCountRelation,
+    setRowCountInfo,
     setStatus,
     setTableItems,
     sortingColumns,
@@ -206,12 +205,13 @@ export const useIndexData = (
     const docs = resp.hits.hits.map((d) => getProcessedFields(d.fields ?? {}));
 
     setCcsWarning(isCrossClusterSearch && isMissingFields);
-    setRowCount(typeof resp.hits.total === 'number' ? resp.hits.total : resp.hits.total!.value);
-    setRowCountRelation(
-      typeof resp.hits.total === 'number'
-        ? ('eq' as estypes.SearchTotalHitsRelation)
-        : resp.hits.total!.relation
-    );
+    setRowCountInfo({
+      rowCount: typeof resp.hits.total === 'number' ? resp.hits.total : resp.hits.total!.value,
+      rowCountRelation:
+        typeof resp.hits.total === 'number'
+          ? ('eq' as estypes.SearchTotalHitsRelation)
+          : resp.hits.total!.relation,
+    });
     setTableItems(docs);
     setStatus(INDEX_STATUS.LOADED);
   };

x-pack/plugins/transform/public/app/hooks/use_pivot_data.ts

@@ -143,8 +143,7 @@ export const usePivotData = (
     resetPagination,
     setErrorMessage,
     setNoDataMessage,
-    setRowCount,
-    setRowCountRelation,
+    setRowCountInfo,
     setStatus,
     setTableItems,
     sortingColumns,
@@ -154,8 +153,10 @@ export const usePivotData = (
   const getPreviewData = async () => {
     if (!validationStatus.isValid) {
       setTableItems([]);
-      setRowCount(0);
-      setRowCountRelation(ES_CLIENT_TOTAL_HITS_RELATION.EQ);
+      setRowCountInfo({
+        rowCount: 0,
+        rowCountRelation: ES_CLIENT_TOTAL_HITS_RELATION.EQ,
+      });
       setNoDataMessage(validationStatus.errorMessage!);
       return;
     }
@@ -175,8 +176,10 @@ export const usePivotData = (
     if (!isPostTransformsPreviewResponseSchema(resp)) {
       setErrorMessage(getErrorMessage(resp));
       setTableItems([]);
-      setRowCount(0);
-      setRowCountRelation(ES_CLIENT_TOTAL_HITS_RELATION.EQ);
+      setRowCountInfo({
+        rowCount: 0,
+        rowCountRelation: ES_CLIENT_TOTAL_HITS_RELATION.EQ,
+      });
       setPreviewMappingsProperties({});
       setStatus(INDEX_STATUS.ERROR);
       return;
@@ -208,8 +211,10 @@ export const usePivotData = (
     populatedProperties = getCombinedProperties(populatedProperties, docs);
 
     setTableItems(docs);
-    setRowCount(docs.length);
-    setRowCountRelation(ES_CLIENT_TOTAL_HITS_RELATION.EQ);
+    setRowCountInfo({
+      rowCount: docs.length,
+      rowCountRelation: ES_CLIENT_TOTAL_HITS_RELATION.EQ,
+    });
     setPreviewMappingsProperties(populatedProperties);
     setStatus(INDEX_STATUS.LOADED);