feat(memory): free memory automatically

Kitware · Feb 27, 2024 · 1efa6b4 · 1efa6b4
1 parent 27c80e1
commit 1efa6b4
Show file tree

Hide file tree

Showing 3 changed files with 60 additions and 17 deletions.
diff --git a/examples/vtk/glyph.py b/examples/vtk/glyph.py
@@ -105,7 +105,7 @@ def setup_vtk():
     iren.SetInteractorStyle(istyle)
     iren.SetRenderWindow(renWin)
     ren.ResetCamera()
-    return renWin
+    return renWin, cs2, ss
 
 
 # -----------------------------------------------------------------------------
@@ -116,7 +116,7 @@ def setup_vtk():
 class App:
     def __init__(self, server=None):
         self.server = get_server(server, client_type="vue3")
-        self.render_window = setup_vtk()
+        self.render_window, self.cone, self.sphere = setup_vtk()
         self.ui = self._build_ui()
 
     @property
@@ -126,6 +126,14 @@ def ctrl(self):
     def _build_ui(self):
         with SinglePageLayout(self.server) as layout:
             layout.icon.click = self.ctrl.view_reset_camera
+            with layout.toolbar:
+                vuetify.VSpacer()
+                vuetify.VSlider(
+                    v_model=("resolution", 6),
+                    dense=True,
+                    hide_details=True,
+                )
+                vuetify.VBtn("Update", click=self.ctrl.view_update)
 
             with layout.content:
                 with vuetify.VContainer(
@@ -135,7 +143,8 @@ def _build_ui(self):
                     with vuetify.VContainer(
                         fluid=True, classes="pa-0 fill-height", style="width: 50%;"
                     ):
-                        vtklocal.LocalView(self.render_window)
+                        view = vtklocal.LocalView(self.render_window)
+                        self.ctrl.view_update = view.update
                     with vuetify.VContainer(
                         fluid=True, classes="pa-0 fill-height", style="width: 50%;"
                     ):

diff --git a/trame_vtklocal/widgets/vtklocal.py b/trame_vtklocal/widgets/vtklocal.py
@@ -30,7 +30,7 @@ def __init__(self, render_window, **kwargs):
 
         # Must trigger update after registration
         self._window_id = self.object_manager.RegisterObject(render_window)
-        self.api.update()
+        self.object_manager.Update()
 
         self._attributes["rw_id"] = f'render-window="{self._window_id}"'
         self._attributes["ref"] = f'ref="{self.__ref}"'
@@ -47,12 +47,4 @@ def object_manager(self):
 
     def update(self):
         self.api.update()
-        # TODO broadcast modified [(id,mtime), ...]
-        print(f"update({self.api.active_ids=})")
-        for vtk_id in self.api.active_ids:
-            vtk_obj = self.object_manager.GetObjectAtId(vtk_id)
-            print(f" - {vtk_id}:{vtk_obj.GetClassName()}")
-        self.server.js_call(
-            self.__ref,
-            "update",
-        )
+        self.server.js_call(self.__ref, "update")
diff --git a/vue-components/src/components/VtkLocal.js b/vue-components/src/components/VtkLocal.js
@@ -2,11 +2,15 @@ import { inject, ref, unref, onMounted, onBeforeUnmount, nextTick } from "vue";
 import { createModule } from "../utils";
 
 export default {
-  emits: ["updated"],
+  emits: ["updated", "memory"],
   props: {
     renderWindow: {
       type: String,
     },
+    cacheSize: {
+      type: Number,
+      default: 100000000,
+    },
     wsClient: {
       type: Object,
     },
@@ -19,7 +23,8 @@ export default {
     const canvas = ref(null);
     const client = props.wsClient || trame?.client;
     const stateMTimes = {};
-    const hashesAvailable = new Set();
+    const hashesMTime = {};
+    const currentMTime = ref(1);
     let objectManager = null;
 
     function resize() {
@@ -56,7 +61,7 @@ export default {
       console.log(`vtkLocal::hash(${hash})`);
       const array = new Uint8Array(await blob.arrayBuffer());
       objectManager.registerBlob(hash, array);
-      hashesAvailable.add(hash);
+      hashesMTime[hash] = unref(currentMTime);
       return blob;
     }
 
@@ -80,9 +85,10 @@ export default {
         objectManager.unRegisterState(vtkId);
       });
       serverStatus.hashes.forEach((hash) => {
-        if (!hashesAvailable.has(hash)) {
+        if (!hashesMTime[hash]) {
           pendingRequests.push(fetchHash(hash));
         }
+        hashesMTime[hash] = unref(currentMTime);
       });
       await Promise.all(pendingRequests);
       // Shows memory, feel free to remove.
@@ -96,6 +102,42 @@ export default {
       objectManager.update(startEventLoop);
       resize();
       emit("updated");
+
+      // Memory management
+      currentMTime.value++;
+      const threshold =
+        props.cacheSize + objectManager.getTotalVTKDataObjectMemoryUsage();
+      if (objectManager.getTotalBlobMemoryUsage() > threshold) {
+        // Need to remove old blobs
+        const threshold =
+          props.cacheSize + objectManager.getTotalVTKDataObjectMemoryUsage();
+        const tsMap = {};
+        let mtimeToFree = unref(currentMTime);
+        Object.entries(hashesMTime).forEach(([hash, mtime]) => {
+          if (mtime < mtimeToFree) {
+            mtimeToFree = mtime;
+          }
+          const sMtime = mtime.toString();
+          if (tsMap[sMtime]) {
+            tsMap[sMtime].push(hash);
+          } else {
+            tsMap[sMtime] = [hash];
+          }
+        });
+
+        // Remove blobs starting by the old ones
+        while (objectManager.getTotalBlobMemoryUsage() > threshold) {
+          const hashesToDelete = tsMap[mtimeToFree];
+          for (let i = 0; i < hashesToDelete.length; i++) {
+            objectManager.unRegisterBlob(hashesToDelete[i]);
+          }
+          mtimeToFree++;
+        }
+      }
+      emit("memory", {
+        blobs: objectManager.getTotalBlobMemoryUsage(),
+        scene: objectManager.getTotalVTKDataObjectMemoryUsage(),
+      });
     }
 
     onMounted(async () => {