From dfa8d7593a616ffa58afe7331f84266ad07e4bb2 Mon Sep 17 00:00:00 2001
From: Marco polo <salazarm@users.noreply.github.com>
Date: Tue, 17 Sep 2024 15:26:22 -0400
Subject: [PATCH] [StaticSetFilter] Remove potential layoutEffect infinite loop
 (#24542)

## Summary & Motivation
Fix potential infinite loop caused by useLayoutEffect and
isFirstRenderRef. The layout and regular effects were both gated by
isFirstRenderRef, but since the effect is responsible for setting this
flag, React's behavior allows useLayoutEffect to loop. This PR separates
the refs for layout and regular effects to prevent the looping issue.


## How I Tested These Changes

jest tests

## Changelog

NOCHANGELOG
---
 .../src/hooks/useQueryPersistedState.tsx      |  5 +-
 .../src/ui/BaseFilters/useStaticSetFilter.tsx | 51 ++++++++++---------
 2 files changed, 31 insertions(+), 25 deletions(-)

diff --git a/js_modules/dagster-ui/packages/ui-core/src/hooks/useQueryPersistedState.tsx b/js_modules/dagster-ui/packages/ui-core/src/hooks/useQueryPersistedState.tsx
index 95e5b89acdd19..f48fd1cfac1dc 100644
--- a/js_modules/dagster-ui/packages/ui-core/src/hooks/useQueryPersistedState.tsx
+++ b/js_modules/dagster-ui/packages/ui-core/src/hooks/useQueryPersistedState.tsx
@@ -109,8 +109,9 @@ export function useQueryPersistedState<T extends QueryPersistedDataType>(
       }
 
       // Check if the query has changed. If so, perform a replace. Otherwise, do nothing
-      // to ensure that we don't end up in a `replace` loop.
-      if (!areQueriesEqual(currentQueryString, next)) {
+      // to ensure that we don't end up in a `replace` loop. If we're not in prod, always run
+      // the `replace` so that we surface any unwanted loops during development.
+      if (process.env.NODE_ENV !== 'production' || !areQueriesEqual(currentQueryString, next)) {
         currentQueryString = next;
         history.replace(`${location.pathname}?${qs.stringify(next, {arrayFormat: 'brackets'})}`);
       }
diff --git a/js_modules/dagster-ui/packages/ui-core/src/ui/BaseFilters/useStaticSetFilter.tsx b/js_modules/dagster-ui/packages/ui-core/src/ui/BaseFilters/useStaticSetFilter.tsx
index 8e43551b01f90..f2b4994dad935 100644
--- a/js_modules/dagster-ui/packages/ui-core/src/ui/BaseFilters/useStaticSetFilter.tsx
+++ b/js_modules/dagster-ui/packages/ui-core/src/ui/BaseFilters/useStaticSetFilter.tsx
@@ -3,7 +3,6 @@ import {
   ComponentProps,
   Fragment,
   useContext,
-  useEffect,
   useLayoutEffect,
   useMemo,
   useRef,
@@ -69,7 +68,7 @@ export function useStaticSetFilter<TValue>({
   renderActiveStateLabel,
   filterBarTagLabel,
   isLoadingFilters,
-  state,
+  state: controlledState,
   getStringValue,
   getTooltipText,
   onStateChanged,
@@ -91,23 +90,25 @@ export function useStaticSetFilter<TValue>({
   }, [StaticFilterSorter, name, _unsortedValues]);
 
   // This filter can be used as both a controlled and an uncontrolled component necessitating an innerState for the uncontrolled case.
-  const [innerState, setState] = useState(() => new Set(state || []));
+  const [innerStateRef, setInnerState] = useState(() => ({
+    current: new Set(controlledState || []),
+  }));
 
-  const isFirstRenderRef = useRef(true);
+  const isFirstLayoutEffectRender = useRef(true);
 
   useLayoutEffect(() => {
-    if (!isFirstRenderRef.current) {
-      onStateChanged?.(innerState);
+    if (!isFirstLayoutEffectRender.current) {
+      onStateChanged?.(innerStateRef.current);
     }
+    isFirstLayoutEffectRender.current = false;
     // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [innerState]);
+  }, [innerStateRef.current]);
 
-  useEffect(() => {
-    if (!isFirstRenderRef.current) {
-      setState(state ? new Set(state) : new Set());
-    }
-    isFirstRenderRef.current = false;
-  }, [state]);
+  useMemo(() => {
+    innerStateRef.current = new Set(controlledState);
+
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [controlledState]);
 
   const currentQueryRef = useRef<string>('');
 
@@ -115,8 +116,8 @@ export function useStaticSetFilter<TValue>({
     () => ({
       name,
       icon,
-      state: innerState,
-      isActive: innerState.size > 0,
+      state: innerStateRef.current,
+      isActive: innerStateRef.current.size > 0,
       isLoadingFilters,
       getResults: (query) => {
         currentQueryRef.current = query;
@@ -160,7 +161,9 @@ export function useStaticSetFilter<TValue>({
                   value={selectAllSymbol}
                   renderLabel={() => <>{selectAllText ?? 'Select all'}</>}
                   isForcedActive={
-                    (state instanceof Set ? state.size : state?.length) === allValues.length
+                    (controlledState instanceof Set
+                      ? controlledState.size
+                      : controlledState?.length) === allValues.length
                   }
                   filter={filterObjRef.current}
                   allowMultipleSelections={allowMultipleSelections}
@@ -196,13 +199,15 @@ export function useStaticSetFilter<TValue>({
             return true;
           });
           if (hasAnyUnselected) {
-            setState(new Set([...Array.from(state), ...selectResults.map(({value}) => value)]));
+            setInnerState({
+              current: new Set([...Array.from(state), ...selectResults.map(({value}) => value)]),
+            });
           } else {
             const stateCopy = new Set(state);
             selectResults.forEach(({value}) => {
               stateCopy.delete(value);
             });
-            setState(stateCopy);
+            setInnerState({current: stateCopy});
           }
           return;
         }
@@ -216,7 +221,7 @@ export function useStaticSetFilter<TValue>({
             newState.add(value);
           }
         }
-        setState(newState);
+        setInnerState({current: newState});
         if (closeOnSelect) {
           close();
         }
@@ -233,26 +238,26 @@ export function useStaticSetFilter<TValue>({
       activeJSX: (
         <SetFilterActiveState
           name={name}
-          state={innerState}
+          state={innerStateRef.current}
           getStringValue={getStringValue}
           getTooltipText={getTooltipText}
           renderLabel={renderActiveStateLabel || renderLabel}
           onRemove={() => {
-            setState(new Set());
+            setInnerState({current: new Set()});
           }}
           icon={icon}
           matchType={matchType}
           filterBarTagLabel={filterBarTagLabel}
         />
       ),
-      setState,
+      setState: (val) => setInnerState({current: val}),
       menuWidth,
     }),
     // eslint-disable-next-line react-hooks/exhaustive-deps
     [
       name,
       icon,
-      innerState,
+      innerStateRef.current,
       getStringValue,
       renderActiveStateLabel,
       renderLabel,