From f236041b5c0c3f91d54e4801f8c26328c24e5781 Mon Sep 17 00:00:00 2001
From: Brent Yi <yibrenth@gmail.com>
Date: Mon, 22 Jul 2024 15:54:10 +0900
Subject: [PATCH] Fix broken cleanup, skinned mesh race condition (#250)

* Fix broken cleanup, skinned mesh edge case

* Remove log
---
 examples/23_smpl_visualizer_skinned.py  |  2 --
 src/viser/client/src/MessageHandler.tsx |  3 ---
 src/viser/client/src/SceneTree.tsx      | 11 -----------
 src/viser/client/src/SceneTreeState.tsx | 16 ++++++++++++++--
 4 files changed, 14 insertions(+), 18 deletions(-)

diff --git a/examples/23_smpl_visualizer_skinned.py b/examples/23_smpl_visualizer_skinned.py
index 7fc26100..2fb364c4 100644
--- a/examples/23_smpl_visualizer_skinned.py
+++ b/examples/23_smpl_visualizer_skinned.py
@@ -134,8 +134,6 @@ def main(model_path: Path) -> None:
         # Match transform control gizmos to joint positions.
         for i, control in enumerate(gui_elements.transform_controls):
             control.position = smpl_outputs.T_parent_joint[i, :3, 3]
-            print(control.position)
-
             skinned_handle.bones[i].wxyz = tf.SO3.from_matrix(
                 smpl_outputs.T_world_joint[i, :3, :3]
             ).wxyz
diff --git a/src/viser/client/src/MessageHandler.tsx b/src/viser/client/src/MessageHandler.tsx
index 9b8f26bf..1a133b0c 100644
--- a/src/viser/client/src/MessageHandler.tsx
+++ b/src/viser/client/src/MessageHandler.tsx
@@ -476,9 +476,6 @@ function useMessageHandler() {
                 );
               },
               () => {
-                bones.forEach((bone) => {
-                  bone.remove();
-                });
                 skeleton.dispose();
                 cleanupMesh();
               },
diff --git a/src/viser/client/src/SceneTree.tsx b/src/viser/client/src/SceneTree.tsx
index d3981a07..b7dcbbc6 100644
--- a/src/viser/client/src/SceneTree.tsx
+++ b/src/viser/client/src/SceneTree.tsx
@@ -131,9 +131,6 @@ export function SceneNodeThreeObject(props: {
   const makeObject = viewer.useSceneTree(
     (state) => state.nodeFromName[props.name]?.makeObject,
   );
-  const cleanup = viewer.useSceneTree(
-    (state) => state.nodeFromName[props.name]?.cleanup,
-  );
   const unmountWhenInvisible = viewer.useSceneTree(
     (state) => state.nodeFromName[props.name]?.unmountWhenInvisible,
   );
@@ -256,14 +253,6 @@ export function SceneNodeThreeObject(props: {
     }
   });
 
-  // Clean up when done.
-  React.useEffect(() => {
-    return () => {
-      cleanup && cleanup();
-      delete viewer.nodeRefFromName.current[props.name];
-    };
-  });
-
   // Clicking logic.
   const sendClicksThrottled = useThrottledMessageSender(50);
   const [hovered, setHovered] = React.useState(false);
diff --git a/src/viser/client/src/SceneTreeState.tsx b/src/viser/client/src/SceneTreeState.tsx
index ad5245c7..d542cee3 100644
--- a/src/viser/client/src/SceneTreeState.tsx
+++ b/src/viser/client/src/SceneTreeState.tsx
@@ -1,10 +1,11 @@
-import React from "react";
+import React, { useContext } from "react";
 import { MakeObject, SceneNode } from "./SceneTree";
 import { CoordinateFrame } from "./ThreeAssets";
 import * as THREE from "three";
 import { create } from "zustand";
 import { subscribeWithSelector } from "zustand/middleware";
 import { immer } from "zustand/middleware/immer";
+import { ViewerContext } from "./App";
 
 interface SceneTreeState {
   nodeFromName: { [name: string]: undefined | SceneNode };
@@ -37,6 +38,7 @@ rootNodeTemplate.children.push("/WorldAxes");
 /** Declare a scene state, and return a hook for accessing it. Note that we put
 effort into avoiding a global state! */
 export function useSceneTreeState() {
+  const viewer = useContext(ViewerContext)!;
   return React.useState(() =>
     create(
       subscribeWithSelector(
@@ -51,8 +53,10 @@ export function useSceneTreeState() {
           addSceneNode: (node) =>
             set((state) => {
               const existingNode = state.nodeFromName[node.name];
-              if (existingNode) {
+              if (existingNode !== undefined) {
                 // Node already exists.
+                delete viewer.nodeRefFromName.current[node.name];
+                existingNode.cleanup && existingNode.cleanup(); // Free resources.
                 state.nodeFromName[node.name] = {
                   ...node,
                   children: existingNode.children,
@@ -83,7 +87,10 @@ export function useSceneTreeState() {
               findChildrenRecursive(name);
 
               removeNames.forEach((removeName) => {
+                const node = state.nodeFromName[removeName]!;
+                node.cleanup && node.cleanup(); // Free resources.
                 delete state.nodeFromName[removeName];
+                delete viewer.nodeRefFromName.current[removeName];
               });
 
               // Remove node from parent's children list.
@@ -99,6 +106,11 @@ export function useSceneTreeState() {
                 if (key !== "" && key !== "/WorldAxes")
                   delete state.nodeFromName[key];
               }
+              Object.values(state.nodeFromName).forEach((node) => {
+                // Free resources.
+                if (node === undefined || node.cleanup === undefined) return;
+                node.cleanup();
+              });
               state.nodeFromName[""] = rootNodeTemplate;
               state.nodeFromName["/WorldAxes"] = rootAxesNode;
             }),