fix: DEV-3617: Memory leaks in Taxonomy + Repeater cases (#935)

* Reduce memory footprint by leveraging single-source store

* Taxonomy lazy loading

* Cleanup

* Combine dynamic loading with lazyload

* Remove unused code

* Remove unused code

* Add feature description

* Properly detach/attach shared stores on reset

* Cleanup

* Add large taxonomy example

* Tweaks

Author: nick-skriabin

examples/taxonomy_large/annotations/0.json

@@ -0,0 +1,24 @@
+{
+    "annotations": [
+        {
+            "result": [
+            ]
+        }
+    ],
+    "data": {
+        "asin": "1384719342",
+        "helpful": [
+            13,
+            14
+        ],
+        "overall": 5,
+        "reviewText": "What a waste of money and time!",
+        "reviewTime": "03 16, 2013",
+        "reviewerID": "A14VAT5EAX3D9S",
+        "reviewerName": "Jake",
+        "summary": "Jake",
+        "unixReviewTime": 1363392000
+    },
+    "id": 1,
+    "task_path": "../examples/sentiment_analysis/tasks.json"
+}

examples/taxonomy_large/config.xml

@@ -0,0 +1,28 @@
+<View>
+  <Header>Document ID: $document_id</Header>
+  <Repeater indexFlag="{{idx}}" on="$images" mode="pagination">
+    <View style="display: flex; margin-bottom: 10px;">
+      <View style="margin-right: 10px; width: 20em">
+        <Header value="Labels">Header</Header>
+        <RectangleLabels name="labels_{{idx}}" toName="page_{{idx}}">
+          <Label alias="title" value="Title"/>
+          <Label alias="reference" value="Document reference"/>
+        </RectangleLabels>
+        <RectangleLabels name="taxonomy_{{idx}}" toName="page_{{idx}}">
+          <Label alias="taxonomy" background="#B2CDFE" selectedColor="#74A7FE" value="Primary taxonomy">
+          </Label>
+          <Label alias="taxonomy_other" background="#F9AEC8" selectedColor="#FF7FAC" value="Other taxonomies">
+          </Label>
+        </RectangleLabels>
+        <View>
+          <Taxonomy name="categories_{{idx}}" perRegion="true" toName="page_{{idx}}" visibleWhen="region-selected" whenLabelValue="taxonomy" value="$categories" sharedStore="categories"/>
+          <Taxonomy name="taxonomy_other_{{idx}}" perRegion="true" toName="page_{{idx}}" visibleWhen="region-selected" whenLabelValue="taxonomy_other" value="$other" sharedStore="taxonomies"/>
+        </View>
+      </View>
+      <View style="width: 2048px">
+        <Image maxHeight="auto" name="page_{{idx}}" value="$images[{{idx}}].url">
+        </Image>
+      </View>
+    </View>
+  </Repeater>
+</View>

examples/taxonomy_large/index.js

@@ -0,0 +1,5 @@
+import config from "./config.xml";
+import tasks from "./tasks.json";
+import annotation from "./annotations/0.json";
+
+export const TaxonomyLarge = { config, tasks, annotation };

examples/taxonomy_large/tasks.json

@@ -11,6 +11,7 @@ import { isDefined } from "./utils/utilities";
 import { Hotkey } from "./core/Hotkey";
 import defaultOptions from './defaultOptions';
 import { destroy } from "mobx-state-tree";
+import { destroy as destroySharedStore } from './mixins/SharedChoiceStore/mixin';
 
 configure({
   isolateGlobalState: true,
@@ -70,6 +71,7 @@ export class LabelStudio {
 
     const destructor = () => {
       unmountComponentAtNode(rootElement);
+      destroySharedStore();
       destroy(this.store);
     };
 

src/LabelStudio.js

@@ -39,7 +39,7 @@ function unionTag(arr) {
 function tagsTypes(arr) {
   const type = types.frozen(arr.map(val => val.toLowerCase()));
 
-  type.describe = ()=>`(${arr.join("|")})`;
+  type.describe = () => `(${arr.join("|")})`;
   type.value = arr;
   return type;
 }
@@ -109,4 +109,16 @@ function getParentTagOfTypeString(node, str) {
 const oneOfTools = _oneOf(Registry.getTool, "Not expecting tool: ");
 const toolsArray = _mixedArray(oneOfTools);
 
-export default { unionArray, allModelsTypes, unionTag, tagsTypes, isType, getParentOfTypeString, getParentTagOfTypeString, tagsArray, toolsArray };
+const Types = {
+  unionArray,
+  allModelsTypes,
+  unionTag,
+  tagsTypes,
+  isType,
+  getParentOfTypeString,
+  getParentTagOfTypeString,
+  tagsArray,
+  toolsArray,
+};
+
+export default Types;

src/core/Types.js

@@ -43,6 +43,7 @@ import { ImageTools } from "../examples/image_tools";
  */
 import { HTMLDocument } from "../examples/html_document";
 import { Taxonomy } from "../examples/taxonomy";
+import { TaxonomyLarge } from "../examples/taxonomy_large";
 
 /**
  * RichText (HTML or plain text)
@@ -68,7 +69,7 @@ import { TimeSeriesSingle } from "../examples/timeseries_single";
  */
 // import { AllTypes } from "../examples/all_types";
 
-const data = VideoRectangles;
+const data = TaxonomyLarge;
 
 function getData(task) {
   if (task && task.data) {

src/env/development.js

@@ -1,23 +1,23 @@
-import { flow, types } from "mobx-state-tree";
+import { getRoot, types } from "mobx-state-tree";
 import ProcessAttrsMixin from "./ProcessAttrs";
 import { parseValue } from "../utils/data";
 
+
 const DynamicChildrenMixin = types.model({
 })
-  .views(self => {
-    return {
-      get defaultChildType() {
-        console.error("DynamicChildrenMixin needs to implement defaultChildType getter in views");
-        return undefined;
-      },
-    };
-  })
+  .views(() => ({
+    get defaultChildType() {
+      console.error("DynamicChildrenMixin needs to implement defaultChildType getter in views");
+      return undefined;
+    },
+  }))
   .actions(self => {
     const prepareDynamicChildrenData = (data) => data?.map?.(obj => ({
       type: self.defaultChildType,
       ...obj,
       children: prepareDynamicChildrenData(obj.children),
     }));
+
     const postprocessDynamicChildren = (children, store) => {
       children?.forEach(item => {
         postprocessDynamicChildren(item.children, store);
@@ -26,29 +26,45 @@ const DynamicChildrenMixin = types.model({
     };
 
     return {
+      updateWithDynamicChildren(data) {
+        const list = (self.children ?? []).concat(prepareDynamicChildrenData(data));
+
+        self.children = list;
+      },
+
       updateValue(store) {
         // If we want to use resolveValue or another asynchronous method here
         // we may need to rewrite this, initRoot and the other related methods
         // (actually a lot of them) to work asynchronously as well
 
-        const valueFromTask = parseValue(self.value, store.task.dataObj);
+        setTimeout(() => {
+          self.updateDynamicChildren(store);
+        });
+      },
+
+      updateDynamicChildren(store) {
+        if (self.locked !== true) {
+          const valueFromTask = parseValue(self.value, store.task.dataObj);
 
-        if (!valueFromTask) return;
+          if (!valueFromTask) return;
 
-        self.generateDynamicChildren(valueFromTask, store);
-        if (self.annotation) self.needsUpdate?.();
+          self.updateWithDynamicChildren(valueFromTask);
+          // self.generateDynamicChildren(valueFromTask, store);
+          if (self.annotation) self.needsUpdate?.();
+        }
       },
 
       generateDynamicChildren(data, store) {
-        if (!self.children) {
-          self.children = prepareDynamicChildrenData(data);
-        } else {
-          self.children.push(...prepareDynamicChildrenData(data));
-        }
+        if (self.children) {
+          const children = self.children;
+          const len = children.length;
+          const start = len - data.length;
+          const slice = children.slice(start, len);
 
-        if (self.children) postprocessDynamicChildren(self.children.slice(self.children.length - data.length,self.children.length), store);
+          postprocessDynamicChildren(slice, store);
+        }
       },
     };
   });
 
-export default types.compose(ProcessAttrsMixin, DynamicChildrenMixin);
+export default types.compose(ProcessAttrsMixin, DynamicChildrenMixin);

src/mixins/DynamicChildrenMixin.js

@@ -0,0 +1,35 @@
+import { destroy, detach, types } from "mobx-state-tree";
+import { SharedStoreModel } from "./model";
+import { Stores } from "./mixin";
+
+/**
+ * StoreExtender injects into the AnnotationStore and holds every created SharedStore.
+ *
+ * Underlying tags that use SharedStoreMixin have access to methods of this mixin to add
+ * their SharedStore instances.
+ */
+export const StoreExtender = types.model("StoreExtender", {
+  sharedStores: types.optional(types.map(SharedStoreModel), {}),
+}).actions((self) => ({
+  addSharedStore(store) {
+    self.sharedStores.set(store.id, store);
+  },
+  beforeReset() {
+    self.sharedStores.forEach((store) => {
+      detach(store);
+    });
+    self.sharedStores.clear();
+  },
+  afterReset() {
+    Stores.forEach((store) => {
+      self.addSharedStore(store);
+    });
+  },
+  beforeDestroy() {
+    self.sharedStores.forEach((store) => {
+      detach(store);
+      destroy(store);
+    });
+    self.sharedStores.clear();
+  },
+}));

src/mixins/SharedChoiceStore/extender.js

@@ -0,0 +1,105 @@
+import { types } from "mobx-state-tree";
+import Types from "../../core/Types";
+import { SharedStoreModel } from "./model";
+
+/**
+ * StoreIds and Stores act as a cache.
+ *
+ * The reason behind those is that we're creating a new store on the `preProcessSnapshot` when there's no
+ * access to the State Tree. When the store is created, it's put into the cache and retrieved back in the
+ * `afterCreate` hook of the model.
+ *
+ * StoreIds is just a map of existing store IDs to reference to during the `preProcessSnapshot`.
+ */
+export const Stores = new Map();
+const StoreIds = new Set();
+
+/**
+ * Defines the ID to group SharedStores by.
+ */
+const SharedStoreID = types.optional(types.maybeNull(types.string), null);
+
+/**
+ * Defines the Store model referenced from the Annotation Store
+ */
+const Store = types.optional(types.maybeNull(types.late(() => types.reference(SharedStoreModel))), null);
+
+/**
+ * SharedStoreMixin, when injected into the model, provides an AnnotationStore level shared storages to
+ * reduce the memory footprint and computation time.
+ *
+ * It was specifically designed to be used with Repeater tag where the memory issues are the most sound.
+ *
+ * This mixin provedes a `sharedStore` property to the model which is a reference to the shared store.
+ *
+ * The concept behind it is that whenever a model is parsing a snapshot, children are subtracted from the
+ * initial snapshot, and put into the newly created SharedStore.
+ *
+ * The store is then put into the cache and attached to the model in the `afterCreate` hook. Any subsequent
+ * models lookup the store in the cache first and use its id instead of creating a new one.
+ *
+ * When the store is fullfilled with children, it's locked and cannot be modified anymore. The allows the model
+ * not to process children anymore and just use the store.
+ *
+ * Shared Stores live on the AnnotationStore level meaning that even if the user switches between annotations or
+ * create new ones, they will all use the same shared store decreasing the memory footprint and computation time.
+ */
+export const SharedStoreMixin = types.model("SharedStoreMixin", {
+  sharedstore: SharedStoreID,
+  store: Store,
+})
+  .views((self) => ({
+    get children() {
+      return self.sharedChildren;
+    },
+
+    get locked() {
+      return self.store?.locked ?? false;
+    },
+
+    set children(val) {
+      self.store?.lock();
+      self.store.setChildren(val);
+    },
+
+    get sharedChildren() {
+      return self.store.children ?? [];
+    },
+
+    get storeId() {
+      return self.sharedstore ?? self.name;
+    },
+  }))
+  .actions(self => ({
+    afterCreate() {
+      if (!self.store) {
+        const store = Stores.get(self.storeId);
+        const annotationStore = Types.getParentOfTypeString(self, "AnnotationStore");
+
+        annotationStore.addSharedStore(store);
+        StoreIds.add(self.storeId);
+        self.store = self.storeId;
+      }
+    },
+  }))
+  .preProcessSnapshot((sn) => {
+    const storeId = sn.sharedstore ?? sn.name;
+
+    if (StoreIds.has(storeId)) {
+      sn.store = storeId;
+    } else {
+      Stores.set(storeId, SharedStoreModel.create({
+        id: storeId,
+        children: sn._children ?? sn.children ?? [],
+      }));
+    }
+
+    return sn;
+  });
+
+export const destroy = () => {
+  console.log("destroying");
+  Stores.clear();
+  StoreIds.clear();
+};
+