[ML] Improvements for urlState hook. (#70576) (#71156)

Makes two improvements to the urlState hook (also known as appState in some places):

- There was always a risk to run into a race condition because setUrlState could refer to a stale version of the state to act upon, for example if two calls were done in parallel. This is now fixed by using a local state copy of what we get from useLocation(). This allows us to use the callback version of useState's set function so we can make sure we always modify the latest state.
- Calls to history.push() are now gated by a check if the change actually referred to the corresponding instance of urlState (either _g or _a), this should reduce the updates resulting re-renders.

The two changes should make the use of setUrlState more safe against the pitfalls (race conditions/stale updates/lots of rerenders) we previously faced.

Co-authored-by: Elastic Machine <elasticmachine@users.noreply.github.com>

x-pack/plugins/ml/public/application/components/controls/select_interval/select_interval.test.tsx

@@ -11,13 +11,17 @@ import { mount } from 'enzyme';
 
 import { EuiSelect } from '@elastic/eui';
 
+import { UrlStateProvider } from '../../../util/url_state';
+
 import { SelectInterval } from './select_interval';
 
 describe('SelectInterval', () => {
   test('creates correct initial selected value', () => {
     const wrapper = mount(
       <MemoryRouter>
-        <SelectInterval />
+        <UrlStateProvider>
+          <SelectInterval />
+        </UrlStateProvider>
       </MemoryRouter>
     );
     const select = wrapper.find(EuiSelect);
@@ -29,7 +33,9 @@ describe('SelectInterval', () => {
   test('currently selected value is updated correctly on click', (done) => {
     const wrapper = mount(
       <MemoryRouter>
-        <SelectInterval />
+        <UrlStateProvider>
+          <SelectInterval />
+        </UrlStateProvider>
       </MemoryRouter>
     );
     const select = wrapper.find(EuiSelect).first();

x-pack/plugins/ml/public/application/components/controls/select_severity/select_severity.test.tsx

@@ -11,13 +11,17 @@ import { mount } from 'enzyme';
 
 import { EuiSuperSelect } from '@elastic/eui';
 
+import { UrlStateProvider } from '../../../util/url_state';
+
 import { SelectSeverity } from './select_severity';
 
 describe('SelectSeverity', () => {
   test('creates correct severity options and initial selected value', () => {
     const wrapper = mount(
       <MemoryRouter>
-        <SelectSeverity />
+        <UrlStateProvider>
+          <SelectSeverity />
+        </UrlStateProvider>
       </MemoryRouter>
     );
     const select = wrapper.find(EuiSuperSelect);
@@ -65,7 +69,9 @@ describe('SelectSeverity', () => {
   test('state for currently selected value is updated correctly on click', (done) => {
     const wrapper = mount(
       <MemoryRouter>
-        <SelectSeverity />
+        <UrlStateProvider>
+          <SelectSeverity />
+        </UrlStateProvider>
       </MemoryRouter>
     );
 

x-pack/plugins/ml/public/application/explorer/hooks/use_selected_cells.ts

@@ -9,12 +9,10 @@ import { useUrlState } from '../../util/url_state';
 import { SWIMLANE_TYPE } from '../explorer_constants';
 import { AppStateSelectedCells } from '../explorer_utils';
 
-export const useSelectedCells = (): [
-  AppStateSelectedCells | undefined,
-  (swimlaneSelectedCells: AppStateSelectedCells) => void
-] => {
-  const [appState, setAppState] = useUrlState('_a');
-
+export const useSelectedCells = (
+  appState: any,
+  setAppState: ReturnType<typeof useUrlState>[1]
+): [AppStateSelectedCells | undefined, (swimlaneSelectedCells: AppStateSelectedCells) => void] => {
   // keep swimlane selection, restore selectedCells from AppState
   const selectedCells = useMemo(() => {
     return appState?.mlExplorerSwimlane?.selectedType !== undefined

x-pack/plugins/ml/public/application/routing/router.tsx

@@ -12,6 +12,7 @@ import { IUiSettingsClient, ChromeStart } from 'kibana/public';
 import { ChromeBreadcrumb } from 'kibana/public';
 import { IndexPatternsContract } from 'src/plugins/data/public';
 import { MlContext, MlContextValue } from '../contexts/ml';
+import { UrlStateProvider } from '../util/url_state';
 
 import * as routes from './routes';
 
@@ -48,21 +49,23 @@ export const MlRouter: FC<{ pageDeps: PageDependencies }> = ({ pageDeps }) => {
 
   return (
     <HashRouter>
-      <div className="ml-app">
-        {Object.entries(routes).map(([name, route]) => (
-          <Route
-            key={name}
-            path={route.path}
-            exact
-            render={(props) => {
-              window.setTimeout(() => {
-                setBreadcrumbs(route.breadcrumbs);
-              });
-              return route.render(props, pageDeps);
-            }}
-          />
-        ))}
-      </div>
+      <UrlStateProvider>
+        <div className="ml-app">
+          {Object.entries(routes).map(([name, route]) => (
+            <Route
+              key={name}
+              path={route.path}
+              exact
+              render={(props) => {
+                window.setTimeout(() => {
+                  setBreadcrumbs(route.breadcrumbs);
+                });
+                return route.render(props, pageDeps);
+              }}
+            />
+          ))}
+        </div>
+      </UrlStateProvider>
     </HashRouter>
   );
 };

x-pack/plugins/ml/public/application/routing/routes/explorer.tsx

@@ -152,7 +152,7 @@ const ExplorerUrlStateManager: FC<ExplorerUrlStateManagerProps> = ({ jobsWithTim
   const [tableInterval] = useTableInterval();
   const [tableSeverity] = useTableSeverity();
 
-  const [selectedCells, setSelectedCells] = useSelectedCells();
+  const [selectedCells, setSelectedCells] = useSelectedCells(appState, setAppState);
   useEffect(() => {
     explorerService.setSelectedCells(selectedCells);
   }, [JSON.stringify(selectedCells)]);

x-pack/plugins/ml/public/application/util/url_state.test.ts → x-pack/plugins/ml/public/application/util/url_state.test.tsx

@@ -4,8 +4,9 @@
  * you may not use this file except in compliance with the Elastic License.
  */
 
-import { renderHook, act } from '@testing-library/react-hooks';
-import { getUrlState, useUrlState } from './url_state';
+import React, { FC } from 'react';
+import { render, act } from '@testing-library/react';
+import { parseUrlState, useUrlState, UrlStateProvider } from './url_state';
 
 const mockHistoryPush = jest.fn();
 
@@ -22,7 +23,7 @@ jest.mock('react-router-dom', () => ({
 describe('getUrlState', () => {
   test('properly decode url with _g and _a', () => {
     expect(
-      getUrlState(
+      parseUrlState(
         "?_a=(mlExplorerFilter:(),mlExplorerSwimlane:(viewByFieldName:action),query:(query_string:(analyze_wildcard:!t,query:'*')))&_g=(ml:(jobIds:!(dec-2)),refreshInterval:(display:Off,pause:!f,value:0),time:(from:'2019-01-01T00:03:40.000Z',mode:absolute,to:'2019-08-30T11:55:07.000Z'))&savedSearchId=571aaf70-4c88-11e8-b3d7-01146121b73d"
       )
     ).toEqual({
@@ -64,13 +65,19 @@ describe('useUrlState', () => {
   });
 
   test('pushes a properly encoded search string to history', () => {
-    const { result } = renderHook(() => useUrlState('_a'));
+    const TestComponent: FC = () => {
+      const [, setUrlState] = useUrlState('_a');
+      return <button onClick={() => setUrlState({ query: {} })}>ButtonText</button>;
+    };
+
+    const { getByText } = render(
+      <UrlStateProvider>
+        <TestComponent />
+      </UrlStateProvider>
+    );
 
     act(() => {
-      const [, setUrlState] = result.current;
-      setUrlState({
-        query: {},
-      });
+      getByText('ButtonText').click();
     });
 
     expect(mockHistoryPush).toHaveBeenCalledWith({

x-pack/plugins/ml/public/application/util/url_state.ts → x-pack/plugins/ml/public/application/util/url_state.tsx

@@ -5,7 +5,7 @@
  */
 
 import { parse, stringify } from 'query-string';
-import { useCallback } from 'react';
+import React, { createContext, useCallback, useContext, useMemo, FC } from 'react';
 import { isEqual } from 'lodash';
 import { decode, encode } from 'rison-node';
 import { useHistory, useLocation } from 'react-router-dom';
@@ -14,8 +14,16 @@ import { Dictionary } from '../../../common/types/common';
 
 import { getNestedProperty } from './object_utils';
 
-export type SetUrlState = (attribute: string | Dictionary<any>, value?: any) => void;
-export type UrlState = [Dictionary<any>, SetUrlState];
+type Accessor = '_a' | '_g';
+export type SetUrlState = (
+  accessor: Accessor,
+  attribute: string | Dictionary<any>,
+  value?: any
+) => void;
+export interface UrlState {
+  searchString: string;
+  setUrlState: SetUrlState;
+}
 
 /**
  * Set of URL query parameters that require the rison serialization.
@@ -30,7 +38,7 @@ function isRisonSerializationRequired(queryParam: string): boolean {
   return risonSerializedParams.has(queryParam);
 }
 
-export function getUrlState(search: string): Dictionary<any> {
+export function parseUrlState(search: string): Dictionary<any> {
   const urlState: Dictionary<any> = {};
   const parsedQueryString = parse(search, { sort: false });
 
@@ -56,22 +64,31 @@ export function getUrlState(search: string): Dictionary<any> {
 // - `history.push()` is the successor of `save`.
 // - The exposed state and set call make use of the above and make sure that
 //   different urlStates(e.g. `_a` / `_g`) don't overwrite each other.
-export const useUrlState = (accessor: string): UrlState => {
+// This uses a context to be able to maintain only one instance
+// of the url state. It gets passed down with `UrlStateProvider`
+// and can be used via `useUrlState`.
+export const urlStateStore = createContext<UrlState>({
+  searchString: '',
+  setUrlState: () => {},
+});
+const { Provider } = urlStateStore;
+export const UrlStateProvider: FC = ({ children }) => {
   const history = useHistory();
-  const { search } = useLocation();
+  const { search: searchString } = useLocation();
 
-  const setUrlState = useCallback(
-    (attribute: string | Dictionary<any>, value?: any) => {
-      const urlState = getUrlState(search);
-      const parsedQueryString = parse(search, { sort: false });
+  const setUrlState: SetUrlState = useCallback(
+    (accessor: Accessor, attribute: string | Dictionary<any>, value?: any) => {
+      const prevSearchString = searchString;
+      const urlState = parseUrlState(prevSearchString);
+      const parsedQueryString = parse(prevSearchString, { sort: false });
 
       if (!Object.prototype.hasOwnProperty.call(urlState, accessor)) {
         urlState[accessor] = {};
       }
 
       if (typeof attribute === 'string') {
         if (isEqual(getNestedProperty(urlState, `${accessor}.${attribute}`), value)) {
-          return;
+          return prevSearchString;
         }
 
         urlState[accessor][attribute] = value;
@@ -83,7 +100,10 @@ export const useUrlState = (accessor: string): UrlState => {
       }
 
       try {
-        const oldLocationSearch = stringify(parsedQueryString, { sort: false, encode: false });
+        const oldLocationSearchString = stringify(parsedQueryString, {
+          sort: false,
+          encode: false,
+        });
 
         Object.keys(urlState).forEach((a) => {
           if (isRisonSerializationRequired(a)) {
@@ -92,20 +112,41 @@ export const useUrlState = (accessor: string): UrlState => {
             parsedQueryString[a] = urlState[a];
           }
         });
-        const newLocationSearch = stringify(parsedQueryString, { sort: false, encode: false });
+        const newLocationSearchString = stringify(parsedQueryString, {
+          sort: false,
+          encode: false,
+        });
 
-        if (oldLocationSearch !== newLocationSearch) {
-          history.push({
-            search: stringify(parsedQueryString, { sort: false }),
-          });
+        if (oldLocationSearchString !== newLocationSearchString) {
+          const newSearchString = stringify(parsedQueryString, { sort: false });
+          history.push({ search: newSearchString });
         }
       } catch (error) {
         // eslint-disable-next-line no-console
         console.error('Could not save url state', error);
       }
     },
-    [search]
+    [searchString]
   );
 
-  return [getUrlState(search)[accessor], setUrlState];
+  return <Provider value={{ searchString, setUrlState }}>{children}</Provider>;
+};
+
+export const useUrlState = (accessor: Accessor) => {
+  const { searchString, setUrlState: setUrlStateContext } = useContext(urlStateStore);
+
+  const urlState = useMemo(() => {
+    const fullUrlState = parseUrlState(searchString);
+    if (typeof fullUrlState === 'object') {
+      return fullUrlState[accessor];
+    }
+    return undefined;
+  }, [searchString]);
+
+  const setUrlState = useCallback(
+    (attribute: string | Dictionary<any>, value?: any) =>
+      setUrlStateContext(accessor, attribute, value),
+    [accessor, setUrlStateContext]
+  );
+  return [urlState, setUrlState];
 };