[Components]: Fix race condition in tree reload handling

.changeset/tame-stamps-dig.md

@@ -0,0 +1,5 @@
+---
+"@itwin/presentation-components": patch
+---
+
+`usePresentationTreeState`: Fix a race condition where iModel and ruleset change notifications could be missed while the tree was reloading.

apps/full-stack-tests/src/components/tree/Update.test.tsx

@@ -496,7 +496,7 @@ describe("Tree update", () => {
             async () => {
               await expectTree(result.current!.nodeLoader, expectedUpdatedTree);
             },
-            { timeout: 9999999 },
+            { timeout: 30000 },
           );
         }
       })();

packages/components/src/presentation-components/tree/controlled/UsePresentationTreeState.ts

@@ -173,11 +173,11 @@ export function usePresentationTreeState<TEventHandler extends TreeEventHandler
     renderedItems.current = items;
   }, []);
 
+  const modelSourceRef = useLatest(state?.nodeLoader.modelSource);
   useTreeReload({
     pageSize: dataProviderProps.pagingSize,
-    modelSource: state?.nodeLoader.modelSource,
+    modelSource: modelSourceRef,
     dataProviderProps: treeStateProps,
-    ruleset: dataProviderProps.ruleset,
     onReload,
     renderedItems,
   });

packages/components/src/presentation-components/tree/controlled/UseTreeReload.ts

@@ -7,7 +7,6 @@
 import { MutableRefObject, useEffect } from "react";
 import { RenderedItemsRange, Subscription, TreeModelSource } from "@itwin/components-react";
 import { IModelApp } from "@itwin/core-frontend";
-import { Ruleset } from "@itwin/presentation-common";
 import { IModelHierarchyChangeEventArgs, Presentation } from "@itwin/presentation-frontend";
 import { getRulesetId } from "../../common/Utils.js";
 import { PresentationTreeDataProvider, PresentationTreeDataProviderProps } from "../DataProvider.js";
@@ -23,8 +22,7 @@ export interface ReloadedTree {
 export interface TreeReloadParams {
   dataProviderProps: PresentationTreeDataProviderProps;
   pageSize: number;
-  ruleset: string | Ruleset;
-  modelSource?: TreeModelSource;
+  modelSource: MutableRefObject<TreeModelSource | undefined>;
   onReload: (params: ReloadedTree) => void;
   renderedItems: MutableRefObject<RenderedItemsRange | undefined>;
 }
@@ -39,19 +37,29 @@ export function useTreeReload(params: TreeReloadParams) {
 }
 
 function useModelSourceUpdateOnBriefcaseUpdate(params: TreeReloadParams): void {
-  const { dataProviderProps, ruleset, pageSize, modelSource, onReload, renderedItems } = params;
+  const { dataProviderProps, pageSize, modelSource, onReload, renderedItems } = params;
 
   useEffect(() => {
-    if (!modelSource || !dataProviderProps.imodel.isBriefcaseConnection()) {
+    if (!dataProviderProps.imodel.isBriefcaseConnection()) {
       return;
     }
 
     let subscription: Subscription | undefined;
-
     const reload = () => {
+      const currentModelSource = modelSource?.current;
+      /* v8 ignore if -- @preserve */
+      if (!currentModelSource) {
+        return;
+      }
       /* v8 ignore next -- @preserve */
       subscription?.unsubscribe();
-      subscription = startTreeReload({ dataProviderProps, ruleset, pageSize, modelSource, renderedItems, onReload });
+      subscription = startTreeReload({
+        dataProviderProps,
+        pageSize,
+        modelSource: currentModelSource,
+        renderedItems: renderedItems.current,
+        onReload,
+      });
     };
 
     const removePullListener = dataProviderProps.imodel.txns.onChangesPulled.addListener(reload);
@@ -62,60 +70,72 @@ function useModelSourceUpdateOnBriefcaseUpdate(params: TreeReloadParams): void {
       removePushListener();
       subscription?.unsubscribe();
     };
-  }, [modelSource, pageSize, dataProviderProps, ruleset, onReload, renderedItems]);
+  }, [modelSource, pageSize, dataProviderProps, onReload, renderedItems]);
 }
 
 function useModelSourceUpdateOnIModelHierarchyUpdate(params: TreeReloadParams): void {
-  const { dataProviderProps, ruleset, pageSize, modelSource, onReload, renderedItems } = params;
+  const { dataProviderProps, pageSize, modelSource, onReload, renderedItems } = params;
 
   useEffect(() => {
-    if (!modelSource) {
-      return;
-    }
-
     let subscription: Subscription | undefined;
     const removeListener = Presentation.presentation.onIModelHierarchyChanged.addListener(
       (args: IModelHierarchyChangeEventArgs) => {
-        if (args.rulesetId !== getRulesetId(ruleset) || args.imodelKey !== dataProviderProps.imodel.key) {
+        if (
+          args.rulesetId !== getRulesetId(dataProviderProps.ruleset) ||
+          args.imodelKey !== dataProviderProps.imodel.key
+        ) {
+          return;
+        }
+
+        const currentModelSource = modelSource?.current;
+        /* v8 ignore if -- @preserve */
+        if (!currentModelSource) {
           return;
         }
 
         /* v8 ignore next -- @preserve */
         subscription?.unsubscribe();
-        subscription = startTreeReload({ dataProviderProps, ruleset, pageSize, modelSource, renderedItems, onReload });
+        subscription = startTreeReload({
+          dataProviderProps,
+          pageSize,
+          modelSource: currentModelSource,
+          renderedItems: renderedItems.current,
+          onReload,
+        });
       },
     );
 
     return () => {
       removeListener();
       subscription?.unsubscribe();
     };
-  }, [modelSource, pageSize, dataProviderProps, ruleset, onReload, renderedItems]);
+  }, [modelSource, pageSize, dataProviderProps, onReload, renderedItems]);
 }
 
 function useModelSourceUpdateOnRulesetModification(params: TreeReloadParams): void {
-  const { dataProviderProps, ruleset, pageSize, modelSource, onReload, renderedItems } = params;
+  const { dataProviderProps, pageSize, modelSource, onReload, renderedItems } = params;
 
   useEffect(() => {
-    if (!modelSource) {
-      return;
-    }
-
     let subscription: Subscription | undefined;
     const removeListener = Presentation.presentation.rulesets().onRulesetModified.addListener((modifiedRuleset) => {
-      if (modifiedRuleset.id !== getRulesetId(ruleset)) {
+      if (modifiedRuleset.id !== getRulesetId(dataProviderProps.ruleset)) {
+        return;
+      }
+
+      const currentModelSource = modelSource?.current;
+      /* v8 ignore if -- @preserve */
+      if (!currentModelSource) {
         return;
       }
 
       // use ruleset id as only registered rulesets can be modified.
       /* v8 ignore next -- @preserve */
       subscription?.unsubscribe();
       subscription = startTreeReload({
-        dataProviderProps,
-        ruleset: modifiedRuleset.id,
+        dataProviderProps: { ...dataProviderProps, ruleset: modifiedRuleset.id },
         pageSize,
-        modelSource,
-        renderedItems,
+        modelSource: currentModelSource,
+        renderedItems: renderedItems.current,
         onReload,
       });
     });
@@ -124,64 +144,84 @@ function useModelSourceUpdateOnRulesetModification(params: TreeReloadParams): vo
       removeListener();
       subscription?.unsubscribe();
     };
-  }, [dataProviderProps, ruleset, modelSource, pageSize, onReload, renderedItems]);
+  }, [dataProviderProps, modelSource, pageSize, onReload, renderedItems]);
 }
 
 function useModelSourceUpdateOnRulesetVariablesChange(params: TreeReloadParams): void {
-  const { dataProviderProps, pageSize, ruleset, modelSource, onReload, renderedItems } = params;
+  const { dataProviderProps, pageSize, modelSource, onReload, renderedItems } = params;
 
   useEffect(() => {
-    if (!modelSource) {
-      return;
-    }
-
     let subscription: Subscription | undefined;
-    const removeListener = Presentation.presentation.vars(getRulesetId(ruleset)).onVariableChanged.addListener(() => {
-      // note: we should probably debounce these events while accumulating changed variables in case multiple vars are changed
-      /* v8 ignore next -- @preserve */
-      subscription?.unsubscribe();
-      subscription = startTreeReload({ dataProviderProps, ruleset, pageSize, modelSource, renderedItems, onReload });
-    });
+    const removeListener = Presentation.presentation
+      .vars(getRulesetId(dataProviderProps.ruleset))
+      .onVariableChanged.addListener(() => {
+        const currentModelSource = modelSource?.current;
+        /* v8 ignore if -- @preserve */
+        if (!currentModelSource) {
+          return;
+        }
+
+        // note: we should probably debounce these events while accumulating changed variables in case multiple vars are changed
+        /* v8 ignore next -- @preserve */
+        subscription?.unsubscribe();
+        subscription = startTreeReload({
+          dataProviderProps,
+          pageSize,
+          modelSource: currentModelSource,
+          renderedItems: renderedItems.current,
+          onReload,
+        });
+      });
 
     return () => {
       removeListener();
       subscription?.unsubscribe();
     };
-  }, [dataProviderProps, modelSource, pageSize, ruleset, onReload, renderedItems]);
+  }, [dataProviderProps, modelSource, pageSize, onReload, renderedItems]);
 }
 
 function useModelSourceUpdateOnUnitSystemChange(params: TreeReloadParams): void {
-  const { dataProviderProps, pageSize, ruleset, modelSource, onReload, renderedItems } = params;
+  const { dataProviderProps, pageSize, modelSource, onReload, renderedItems } = params;
 
   useEffect(() => {
-    if (!modelSource) {
-      return;
-    }
-
     let subscription: Subscription | undefined;
     const removeListener = IModelApp.quantityFormatter.onActiveFormattingUnitSystemChanged.addListener(() => {
+      const currentModelSource = modelSource?.current;
+      /* v8 ignore if -- @preserve */
+      if (!currentModelSource) {
+        return;
+      }
+
       /* v8 ignore next -- @preserve */
       subscription?.unsubscribe();
-      subscription = startTreeReload({ dataProviderProps, ruleset, pageSize, modelSource, renderedItems, onReload });
+      subscription = startTreeReload({
+        dataProviderProps,
+        pageSize,
+        modelSource: currentModelSource,
+        renderedItems: renderedItems.current,
+        onReload,
+      });
     });
 
     return () => {
       removeListener();
       subscription?.unsubscribe();
     };
-  }, [dataProviderProps, modelSource, pageSize, ruleset, onReload, renderedItems]);
+  }, [dataProviderProps, modelSource, pageSize, onReload, renderedItems]);
 }
 
 function startTreeReload({
   dataProviderProps,
-  ruleset,
   modelSource,
   pageSize,
   renderedItems,
   onReload,
-}: Required<TreeReloadParams>): Subscription {
-  const dataProvider = new PresentationTreeDataProvider({ ...dataProviderProps, ruleset });
-  return reloadTree(modelSource.getModel(), dataProvider, pageSize, renderedItems.current).subscribe({
+}: Omit<Required<TreeReloadParams>, "modelSource" | "renderedItems"> & {
+  modelSource: TreeModelSource;
+  renderedItems: RenderedItemsRange | undefined;
+}): Subscription {
+  const dataProvider = new PresentationTreeDataProvider(dataProviderProps);
+  return reloadTree(modelSource.getModel(), dataProvider, pageSize, renderedItems).subscribe({
     next: (newModelSource) => onReload({ modelSource: newModelSource, dataProvider }),
   });
 }