From 379eed5713e9bddec79b195e91e0c39414c1e468 Mon Sep 17 00:00:00 2001
From: Rhys Howell <rhys.howell@mongodb.com>
Date: Mon, 29 Sep 2025 14:23:13 +0200
Subject: [PATCH 1/2] fix(data-modeling): throttle diagramming prop updates to
 avoid issues with interactions involving frequent rendering COMPASS-9738

---
 .../src/components/diagram-editor.tsx         | 109 +++++++++++++++---
 1 file changed, 93 insertions(+), 16 deletions(-)

diff --git a/packages/compass-data-modeling/src/components/diagram-editor.tsx b/packages/compass-data-modeling/src/components/diagram-editor.tsx
index bc9607d8ca2..4a64fcc26e1 100644
--- a/packages/compass-data-modeling/src/components/diagram-editor.tsx
+++ b/packages/compass-data-modeling/src/components/diagram-editor.tsx
@@ -121,6 +121,8 @@ const ZOOM_OPTIONS = {
   minZoom: 0.25,
 };
 
+const REFRESH_DIAGRAM_RATE_MS = 250;
+
 type SelectedItems = NonNullable<DiagramState>['selectedItems'];
 
 const DiagramContent: React.FunctionComponent<{
@@ -319,6 +321,96 @@ const DiagramContent: React.FunctionComponent<{
     [handleNodesConnect]
   );
 
+  // Throttling mechanism for diagram content updates
+  const lastDiagramUpdateMS = useRef(0);
+  const pendingUpdate = useRef<NodeJS.Timeout | null>(null);
+  const [throttledDiagramProps, setThrottledDiagramProps] = useState({
+    isDarkMode,
+    diagramLabel,
+    edges,
+    nodes,
+    onNodeClick,
+    onPaneClick,
+    onEdgeClick,
+    onFieldClick,
+    onNodeDragStop,
+    onConnect,
+  });
+
+  // Throttle diagram props updating. This ensures we don't run
+  // into broken state bugs with ReactFlow like COMPASS-9738.
+  useEffect(() => {
+    const now = Date.now();
+    const timeSinceLastUpdate = now - lastDiagramUpdateMS.current;
+
+    const updateProps = () => {
+      lastDiagramUpdateMS.current = Date.now();
+      setThrottledDiagramProps({
+        isDarkMode,
+        diagramLabel,
+        edges,
+        nodes,
+        onNodeClick,
+        onPaneClick,
+        onEdgeClick,
+        onFieldClick,
+        onNodeDragStop,
+        onConnect,
+      });
+    };
+
+    if (timeSinceLastUpdate >= REFRESH_DIAGRAM_RATE_MS) {
+      updateProps();
+    } else {
+      if (pendingUpdate.current) {
+        clearTimeout(pendingUpdate.current);
+      }
+
+      // Schedule update for the remaining time.
+      const remainingTime = REFRESH_DIAGRAM_RATE_MS - timeSinceLastUpdate;
+      pendingUpdate.current = setTimeout(updateProps, remainingTime);
+    }
+
+    return () => {
+      if (pendingUpdate.current) {
+        clearTimeout(pendingUpdate.current);
+        pendingUpdate.current = null;
+      }
+    };
+  }, [
+    isDarkMode,
+    diagramLabel,
+    edges,
+    nodes,
+    onNodeClick,
+    onPaneClick,
+    onEdgeClick,
+    onFieldClick,
+    onNodeDragStop,
+    onConnect,
+  ]);
+
+  const diagramContent = useMemo(() => {
+    return (
+      <Diagram
+        isDarkMode={throttledDiagramProps.isDarkMode}
+        title={throttledDiagramProps.diagramLabel}
+        edges={throttledDiagramProps.edges}
+        nodes={throttledDiagramProps.nodes}
+        // With threshold too low clicking sometimes gets confused with
+        // dragging
+        nodeDragThreshold={5}
+        onNodeClick={throttledDiagramProps.onNodeClick}
+        onPaneClick={throttledDiagramProps.onPaneClick}
+        onEdgeClick={throttledDiagramProps.onEdgeClick}
+        onFieldClick={throttledDiagramProps.onFieldClick}
+        fitViewOptions={ZOOM_OPTIONS}
+        onNodeDragStop={throttledDiagramProps.onNodeDragStop}
+        onConnect={throttledDiagramProps.onConnect}
+      />
+    );
+  }, [throttledDiagramProps]);
+
   return (
     <div
       ref={setDiagramContainerRef}
@@ -340,22 +432,7 @@ const DiagramContent: React.FunctionComponent<{
             impact your data
           </Banner>
         )}
-        <Diagram
-          isDarkMode={isDarkMode}
-          title={diagramLabel}
-          edges={edges}
-          nodes={nodes}
-          // With threshold too low clicking sometimes gets confused with
-          // dragging
-          nodeDragThreshold={5}
-          onNodeClick={onNodeClick}
-          onPaneClick={onPaneClick}
-          onEdgeClick={onEdgeClick}
-          onFieldClick={onFieldClick}
-          fitViewOptions={ZOOM_OPTIONS}
-          onNodeDragStop={onNodeDragStop}
-          onConnect={onConnect}
-        />
+        {diagramContent}
       </div>
     </div>
   );

From 8326a1ef154000a7af5c9cf9c0dd25cd96540169 Mon Sep 17 00:00:00 2001
From: Rhys Howell <rhys.howell@mongodb.com>
Date: Mon, 29 Sep 2025 17:39:51 +0200
Subject: [PATCH 2/2] fixup: move throttling props to generic hook in Compass
 components

---
 .../src/hooks/use-throttled-props.tsx         |  46 +++++++
 packages/compass-components/src/index.ts      |   1 +
 .../src/components/diagram-editor.tsx         | 126 +++++-------------
 3 files changed, 82 insertions(+), 91 deletions(-)
 create mode 100644 packages/compass-components/src/hooks/use-throttled-props.tsx

diff --git a/packages/compass-components/src/hooks/use-throttled-props.tsx b/packages/compass-components/src/hooks/use-throttled-props.tsx
new file mode 100644
index 00000000000..4dab1e44797
--- /dev/null
+++ b/packages/compass-components/src/hooks/use-throttled-props.tsx
@@ -0,0 +1,46 @@
+import { useEffect, useRef, useState } from 'react';
+
+const DEFAULT_REFRESH_RATE_MS = 250;
+
+export const useThrottledProps = <T extends Record<string, unknown>>(
+  props: T,
+  refreshRate: number = DEFAULT_REFRESH_RATE_MS
+): T => {
+  // Throttling mechanism for Component content updates.
+  const lastUpdateMS = useRef(0);
+  const pendingUpdate = useRef<NodeJS.Timeout | null>(null);
+  const [throttledProps, setThrottledProps] = useState<T>(props);
+
+  // Throttle props updating. This ensures we don't run
+  // into broken state bugs with ReactFlow like COMPASS-9738.
+  useEffect(() => {
+    const now = Date.now();
+    const timeSinceLastUpdate = now - lastUpdateMS.current;
+
+    const updateProps = () => {
+      lastUpdateMS.current = Date.now();
+      setThrottledProps(props);
+    };
+
+    if (timeSinceLastUpdate >= refreshRate) {
+      updateProps();
+    } else {
+      if (pendingUpdate.current) {
+        clearTimeout(pendingUpdate.current);
+      }
+
+      // Schedule update for the remaining time.
+      const remainingTime = refreshRate - timeSinceLastUpdate;
+      pendingUpdate.current = setTimeout(updateProps, remainingTime);
+    }
+
+    return () => {
+      if (pendingUpdate.current) {
+        clearTimeout(pendingUpdate.current);
+        pendingUpdate.current = null;
+      }
+    };
+  }, [props, refreshRate]);
+
+  return throttledProps;
+};
diff --git a/packages/compass-components/src/index.ts b/packages/compass-components/src/index.ts
index ffacf0f8226..83d02c33e66 100644
--- a/packages/compass-components/src/index.ts
+++ b/packages/compass-components/src/index.ts
@@ -152,6 +152,7 @@ export {
   Theme,
   ThemeProvider,
 } from './hooks/use-theme';
+export { useThrottledProps } from './hooks/use-throttled-props';
 export {
   ContentWithFallback,
   FadeInPlaceholder,
diff --git a/packages/compass-data-modeling/src/components/diagram-editor.tsx b/packages/compass-data-modeling/src/components/diagram-editor.tsx
index 4a64fcc26e1..649317209eb 100644
--- a/packages/compass-data-modeling/src/components/diagram-editor.tsx
+++ b/packages/compass-data-modeling/src/components/diagram-editor.tsx
@@ -29,6 +29,7 @@ import {
   useDarkMode,
   useDrawerActions,
   useDrawerState,
+  useThrottledProps,
   rafraf,
 } from '@mongodb-js/compass-components';
 import { cancelAnalysis, retryAnalysis } from '../store/analysis-process';
@@ -121,8 +122,6 @@ const ZOOM_OPTIONS = {
   minZoom: 0.25,
 };
 
-const REFRESH_DIAGRAM_RATE_MS = 250;
-
 type SelectedItems = NonNullable<DiagramState>['selectedItems'];
 
 const DiagramContent: React.FunctionComponent<{
@@ -321,95 +320,34 @@ const DiagramContent: React.FunctionComponent<{
     [handleNodesConnect]
   );
 
-  // Throttling mechanism for diagram content updates
-  const lastDiagramUpdateMS = useRef(0);
-  const pendingUpdate = useRef<NodeJS.Timeout | null>(null);
-  const [throttledDiagramProps, setThrottledDiagramProps] = useState({
-    isDarkMode,
-    diagramLabel,
-    edges,
-    nodes,
-    onNodeClick,
-    onPaneClick,
-    onEdgeClick,
-    onFieldClick,
-    onNodeDragStop,
-    onConnect,
-  });
-
-  // Throttle diagram props updating. This ensures we don't run
-  // into broken state bugs with ReactFlow like COMPASS-9738.
-  useEffect(() => {
-    const now = Date.now();
-    const timeSinceLastUpdate = now - lastDiagramUpdateMS.current;
-
-    const updateProps = () => {
-      lastDiagramUpdateMS.current = Date.now();
-      setThrottledDiagramProps({
-        isDarkMode,
-        diagramLabel,
-        edges,
-        nodes,
-        onNodeClick,
-        onPaneClick,
-        onEdgeClick,
-        onFieldClick,
-        onNodeDragStop,
-        onConnect,
-      });
-    };
-
-    if (timeSinceLastUpdate >= REFRESH_DIAGRAM_RATE_MS) {
-      updateProps();
-    } else {
-      if (pendingUpdate.current) {
-        clearTimeout(pendingUpdate.current);
-      }
-
-      // Schedule update for the remaining time.
-      const remainingTime = REFRESH_DIAGRAM_RATE_MS - timeSinceLastUpdate;
-      pendingUpdate.current = setTimeout(updateProps, remainingTime);
-    }
-
-    return () => {
-      if (pendingUpdate.current) {
-        clearTimeout(pendingUpdate.current);
-        pendingUpdate.current = null;
-      }
-    };
-  }, [
-    isDarkMode,
-    diagramLabel,
-    edges,
-    nodes,
-    onNodeClick,
-    onPaneClick,
-    onEdgeClick,
-    onFieldClick,
-    onNodeDragStop,
-    onConnect,
-  ]);
+  const diagramProps = useMemo(
+    () => ({
+      isDarkMode,
+      title: diagramLabel,
+      edges,
+      nodes,
+      onNodeClick,
+      onPaneClick,
+      onEdgeClick,
+      onFieldClick,
+      onNodeDragStop,
+      onConnect,
+    }),
+    [
+      isDarkMode,
+      diagramLabel,
+      edges,
+      nodes,
+      onNodeClick,
+      onPaneClick,
+      onEdgeClick,
+      onFieldClick,
+      onNodeDragStop,
+      onConnect,
+    ]
+  );
 
-  const diagramContent = useMemo(() => {
-    return (
-      <Diagram
-        isDarkMode={throttledDiagramProps.isDarkMode}
-        title={throttledDiagramProps.diagramLabel}
-        edges={throttledDiagramProps.edges}
-        nodes={throttledDiagramProps.nodes}
-        // With threshold too low clicking sometimes gets confused with
-        // dragging
-        nodeDragThreshold={5}
-        onNodeClick={throttledDiagramProps.onNodeClick}
-        onPaneClick={throttledDiagramProps.onPaneClick}
-        onEdgeClick={throttledDiagramProps.onEdgeClick}
-        onFieldClick={throttledDiagramProps.onFieldClick}
-        fitViewOptions={ZOOM_OPTIONS}
-        onNodeDragStop={throttledDiagramProps.onNodeDragStop}
-        onConnect={throttledDiagramProps.onConnect}
-      />
-    );
-  }, [throttledDiagramProps]);
+  const throttledDiagramProps = useThrottledProps(diagramProps);
 
   return (
     <div
@@ -432,7 +370,13 @@ const DiagramContent: React.FunctionComponent<{
             impact your data
           </Banner>
         )}
-        {diagramContent}
+        <Diagram
+          {...throttledDiagramProps}
+          // With threshold too low clicking sometimes gets confused with
+          // dragging.
+          nodeDragThreshold={5}
+          fitViewOptions={ZOOM_OPTIONS}
+        />
       </div>
     </div>
   );