Dynamic bucket allocation & miscellaneous volume improvements

CHANGELOG.unreleased.md

@@ -14,6 +14,7 @@ For upgrade instructions, please check the [migration guide](MIGRATIONS.released
 - The visible meshes are now included in the link copied from the "Share" modal or the "Share" button next to the dataset position. They are automatically loaded for users that open the shared link. [#5993](https://github.com/scalableminds/webknossos/pull/5993)
 
 ### Changed
+- Improved stability and speed of volume annotations when annotating large areas. [#6055](https://github.com/scalableminds/webknossos/pull/6055)
 
 ### Fixed
 - Fixed a bug where deactivated users would still be listed as allowed to access the datasets of their team. [#6070](https://github.com/scalableminds/webknossos/pull/6070)

frontend/javascripts/libs/task_pool.js

@@ -1,6 +1,11 @@
 // @flow
 import { type Saga, type Task, join, call, fork } from "oxalis/model/sagas/effect-generators";
 
+/*
+  Given an array of async tasks, processTaskWithPool
+  allows to execute at most ${poolSize} tasks concurrently.
+ */
+
 export default function* processTaskWithPool(
   tasks: Array<() => Saga<void>>,
   poolSize: number,

frontend/javascripts/libs/worker_pool.js

@@ -0,0 +1,29 @@
+// @flow
+
+import _ from "lodash";
+
+export default class WorkerPool<P, R> {
+  // This class can be used to instantiate multiple web workers
+  // which are then used for computation in a simple round-robin manner.
+  //
+  // Example:
+  // const compressionPool = new WorkerPool(
+  //   () => createWorker(ByteArrayToLz4Base64Worker),
+  //   COMPRESSION_WORKER_COUNT,
+  // );
+  // const promise1 = compressionPool.submit(data1);
+  // const promise2 = compressionPool.submit(data2);
+
+  workers: Array<(...args: Array<P>) => R>;
+  currentWorkerIdx: number;
+
+  constructor(workerFn: () => (...args: Array<P>) => R, count: number) {
+    this.workers = _.range(0, count).map(_idx => workerFn());
+    this.currentWorkerIdx = 0;
+  }
+
+  submit(...args: Array<P>): R {
+    this.currentWorkerIdx = (this.currentWorkerIdx + 1) % this.workers.length;
+    return this.workers[this.currentWorkerIdx](...args);
+  }
+}

frontend/javascripts/oxalis/model.js

@@ -6,6 +6,7 @@ import {} from "oxalis/model/actions/settings_actions";
 
 import _ from "lodash";
 
+import { COMPRESSING_BATCH_SIZE } from "oxalis/model/bucket_data_handling/pushqueue";
 import { type Vector3 } from "oxalis/constants";
 import type { Versions } from "oxalis/view/version_view";
 import { getActiveSegmentationTracingLayer } from "oxalis/model/accessors/volumetracing_accessor";
@@ -259,6 +260,22 @@ export class OxalisModel {
     return storeStateSaved && pushQueuesSaved;
   }
 
+  getPushQueueStats() {
+    const compressingBucketCount = _.sum(
+      Utils.values(this.dataLayers).map(
+        dataLayer => dataLayer.pushQueue.compressionTaskQueue.tasks.length * COMPRESSING_BATCH_SIZE,
+      ),
+    );
+    const waitingForCompressionBucketCount = _.sum(
+      Utils.values(this.dataLayers).map(dataLayer => dataLayer.pushQueue.pendingQueue.size),
+    );
+
+    return {
+      compressingBucketCount,
+      waitingForCompressionBucketCount,
+    };
+  }
+
   forceSave = () => {
     // In contrast to the save function, this method will trigger exactly one saveNowAction
     // regardless of what the current save state is

frontend/javascripts/oxalis/model/bucket_data_handling/bucket.js

@@ -129,11 +129,18 @@ export class DataBucket {
   elementClass: ElementClass;
   visualizedMesh: ?Object;
   visualizationColor: number;
+  // If dirty, the bucket's data was potentially edited and needs to be
+  // saved to the server.
+  dirty: boolean;
+  // `dirtyCount` reflects how many pending snapshots of the bucket exist.
+  // A pending snapshot is a snapshot which was either
+  // - not yet saved (after successful saving the dirtyCount is decremented) or
+  // - not yet created by the PushQueue, since the PushQueue creates the snapshots
+  //   in a debounced manner
   dirtyCount: number = 0;
   pendingOperations: Array<(BucketDataArray) => void> = [];
 
   state: BucketStateEnumType;
-  dirty: boolean;
   accessed: boolean;
   data: ?BucketDataArray;
   temporalBucketManager: TemporalBucketManager;
@@ -228,6 +235,7 @@ export class DataBucket {
     // so that at least the big memory hog is tamed (unfortunately,
     // this doesn't help against references which point directly to this.data)
     this.data = null;
+    this.trigger("bucketCollected");
     // Remove all event handlers (see https://github.com/ai/nanoevents#remove-all-listeners)
     this.emitter.events = {};
   }
@@ -359,10 +367,11 @@ export class DataBucket {
     return data;
   }
 
-  setData(newData: BucketDataArray) {
+  setData(newData: BucketDataArray, newPendingOperations: Array<(BucketDataArray) => void>) {
     this.data = newData;
+    this.pendingOperations = newPendingOperations;
     this.dirty = true;
-    this.trigger("bucketLabeled");
+    this.endDataMutation();
   }
 
   uint8ToTypedBuffer(arrayBuffer: ?Uint8Array) {

frontend/javascripts/oxalis/model/bucket_data_handling/data_cube.js

@@ -6,7 +6,6 @@
 import BackboneEvents from "backbone-events-standalone";
 import _ from "lodash";
 
-import ErrorHandling from "libs/error_handling";
 import {
   type Bucket,
   DataBucket,
@@ -15,33 +14,42 @@ import {
   NullBucket,
   type BucketDataArray,
 } from "oxalis/model/bucket_data_handling/bucket";
+import { type ElementClass } from "types/api_flow_types";
+import { type ProgressCallback } from "libs/progress_callback";
+import { V3 } from "libs/mjs";
 import { VoxelNeighborQueue2D, VoxelNeighborQueue3D } from "oxalis/model/volumetracing/volumelayer";
+import { areBoundingBoxesOverlappingOrTouching } from "libs/utils";
 import {
   getResolutions,
   ResolutionInfo,
   getMappingInfo,
 } from "oxalis/model/accessors/dataset_accessor";
 import { getSomeTracing } from "oxalis/model/accessors/tracing_accessor";
-import { V3 } from "libs/mjs";
 import { globalPositionToBucketPosition } from "oxalis/model/helpers/position_converter";
 import { listenToStoreProperty } from "oxalis/model/helpers/listener_helpers";
 import ArbitraryCubeAdapter from "oxalis/model/bucket_data_handling/arbitrary_cube_adapter";
 import BoundingBox from "oxalis/model/bucket_data_handling/bounding_box";
+import Dimensions, { type DimensionMap } from "oxalis/model/dimensions";
+import ErrorHandling from "libs/error_handling";
 import PullQueue from "oxalis/model/bucket_data_handling/pullqueue";
 import PushQueue from "oxalis/model/bucket_data_handling/pushqueue";
 import Store, { type Mapping } from "oxalis/store";
 import TemporalBucketManager from "oxalis/model/bucket_data_handling/temporal_bucket_manager";
-import Dimensions, { type DimensionMap } from "oxalis/model/dimensions";
+import Toast from "libs/toast";
 import constants, {
   type Vector3,
   type Vector4,
   type BoundingBoxType,
   type LabelMasksByBucketAndW,
   MappingStatusEnum,
 } from "oxalis/constants";
-import { type ElementClass } from "types/api_flow_types";
-import { areBoundingBoxesOverlappingOrTouching } from "libs/utils";
-import { type ProgressCallback } from "libs/progress_callback";
+
+const warnAboutTooManyAllocations = _.once(() => {
+  const msg =
+    "webKnossos needed to allocate an unusually large amount of image data. It is advised to save your work and reload the page.";
+  ErrorHandling.notify(new Error(msg));
+  Toast.warning(msg, { sticky: true });
+});
 
 class CubeEntry {
   data: Map<number, Bucket>;
@@ -65,12 +73,11 @@ const FLOODFILL_VOXEL_THRESHOLD = 5 * 1000000;
 const USE_FLOODFILL_VOXEL_THRESHOLD = false;
 
 class DataCube {
-  MAXIMUM_BUCKET_COUNT = constants.MAXIMUM_BUCKET_COUNT_PER_LAYER;
+  BUCKET_COUNT_SOFT_LIMIT = constants.MAXIMUM_BUCKET_COUNT_PER_LAYER;
   arbitraryCube: ArbitraryCubeAdapter;
   upperBoundary: Vector3;
   buckets: Array<DataBucket>;
   bucketIterator: number = 0;
-  bucketCount: number = 0;
   cubes: Array<CubeEntry>;
   boundingBox: BoundingBox;
   pullQueue: PullQueue;
@@ -81,11 +88,16 @@ class DataCube {
   resolutionInfo: ResolutionInfo;
   layerName: string;
 
+  dirtyBucketsCount: number;
+
   // The cube stores the buckets in a separate array for each zoomStep. For each
   // zoomStep the cube-array contains the boundaries and an array holding the buckets.
-  // The bucket-arrays are initialized large enough to hold the whole cube. Thus no
-  // expanding is necessary. bucketCount keeps track of how many buckets are currently
-  // in the cube.
+  // The bucket array is initialized as an empty array and grows dynamically. If the
+  // length exceeds BUCKET_COUNT_SOFT_LIMIT, it is tried to garbage-collect an older
+  // bucket when placing a new one. If this does not succeed (happens if all buckets
+  // in a volume annotation layer are dirty), the array grows further.
+  // If the array grows beyond 2 * BUCKET_COUNT_SOFT_LIMIT, the user is warned about
+  // this.
   //
   // Each bucket consists of an access-value, the zoomStep and the actual data.
   // The access-values are used for garbage collection. When a bucket is accessed, its
@@ -108,13 +120,42 @@ class DataCube {
     this.resolutionInfo = resolutionInfo;
     this.layerName = layerName;
 
+    this.dirtyBucketsCount = 0;
+
+    if (isSegmentation && process.env.BABEL_ENV !== "test") {
+      setInterval(() => {
+        console.log("#######");
+
+        let dirtyBucketsCount = 0;
+        for (let i = 0; i < this.buckets.length; i++) {
+          if (this.buckets[i].dirty) {
+            dirtyBucketsCount++;
+          }
+        }
+
+        let weirdDirtyState = 0;
+        for (let i = 0; i < this.buckets.length; i++) {
+          if (this.buckets[i].dirty && this.buckets[i].dirtyCount === 0) {
+            weirdDirtyState++;
+          }
+        }
+        if (weirdDirtyState > 0) {
+          console.log("weirdDirtyState", weirdDirtyState);
+        }
+
+        console.log("dirtyBucketsCount", dirtyBucketsCount);
+        console.log("this.buckets.length", this.buckets.length);
+      }, 2000);
+    }
+
     _.extend(this, BackboneEvents);
 
     this.cubes = [];
-    if (isSegmentation) {
-      this.MAXIMUM_BUCKET_COUNT *= 2;
+    if (isSegmentation && process.env.BABEL_ENV !== "test") {
+      // undo
+      this.BUCKET_COUNT_SOFT_LIMIT = Math.floor(this.BUCKET_COUNT_SOFT_LIMIT / 10);
     }
-    this.buckets = new Array(this.MAXIMUM_BUCKET_COUNT);
+    this.buckets = [];
 
     // Initializing the cube-arrays with boundaries
     const cubeBoundary = [
@@ -280,57 +321,59 @@ class DataCube {
   }
 
   markBucketsAsUnneeded(): void {
-    for (let i = 0; i < this.bucketCount; i++) {
+    for (let i = 0; i < this.buckets.length; i++) {
       this.buckets[i].markAsUnneeded();
     }
   }
 
   addBucketToGarbageCollection(bucket: DataBucket): void {
-    if (this.bucketCount >= this.MAXIMUM_BUCKET_COUNT) {
-      for (let i = 0; i < this.bucketCount; i++) {
-        this.bucketIterator = ++this.bucketIterator % this.MAXIMUM_BUCKET_COUNT;
+    if (this.buckets.length >= this.BUCKET_COUNT_SOFT_LIMIT) {
+      let foundCollectibleBucket = false;
+
+      for (let i = 0; i < this.buckets.length; i++) {
+        this.bucketIterator = (this.bucketIterator + 1) % this.buckets.length;
         if (this.buckets[this.bucketIterator].shouldCollect()) {
+          foundCollectibleBucket = true;
           break;
         }
       }
 
-      if (!this.buckets[this.bucketIterator].shouldCollect()) {
-        if (process.env.BABEL_ENV === "test") {
-          throw new Error("Bucket was forcefully evicted/garbage-collected.");
+      if (foundCollectibleBucket) {
+        this.collectBucket(this.buckets[this.bucketIterator]);
+      } else {
+        const warnMessage = `More than ${this.buckets.length} buckets needed to be allocated.`;
+        if (this.buckets.length % 100 === 0) {
+          console.warn(warnMessage);
+          ErrorHandling.notify(new Error(warnMessage), {
+            elementClass: this.elementClass,
+            isSegmentation: this.isSegmentation,
+            resolutionInfo: this.resolutionInfo,
+          });
         }
-        const errorMessage =
-          "A bucket was forcefully garbage-collected. This indicates that too many buckets are currently in RAM.";
-        console.error(errorMessage);
-        ErrorHandling.notify(new Error(errorMessage), {
-          elementClass: this.elementClass,
-          isSegmentation: this.isSegmentation,
-          resolutionInfo: this.resolutionInfo,
-        });
-      }
 
-      this.collectBucket(this.buckets[this.bucketIterator]);
-      this.bucketCount--;
-    }
+        if (this.buckets.length > (2 * constants.MAXIMUM_BUCKET_COUNT_PER_LAYER) / 10) {
+          warnAboutTooManyAllocations();
+        }
 
-    this.bucketCount++;
-    if (this.buckets[this.bucketIterator]) {
-      this.buckets[this.bucketIterator].trigger("bucketCollected");
+        // Effectively, push to `this.buckets` by setting the iterator to
+        // a new index.
+        this.bucketIterator = this.buckets.length;
+      }
     }
+
     this.buckets[this.bucketIterator] = bucket;
-    this.bucketIterator = ++this.bucketIterator % this.MAXIMUM_BUCKET_COUNT;
+    // Note that bucketIterator is allowed to point to the next free
+    // slot (i.e., bucketIterator == this.buckets.length).
+    this.bucketIterator = (this.bucketIterator + 1) % (this.buckets.length + 1);
   }
 
   collectAllBuckets(): void {
     this.pullQueue.clear();
     this.pullQueue.abortRequests();
     for (const bucket of this.buckets) {
-      if (bucket != null) {
-        this.collectBucket(bucket);
-        bucket.trigger("bucketCollected");
-      }
+      this.collectBucket(bucket);
     }
     this.buckets = [];
-    this.bucketCount = 0;
     this.bucketIterator = 0;
   }
 
@@ -648,10 +691,7 @@ class DataCube {
     if (bucket.type === "null") {
       return;
     }
-    bucket.setData(data);
-    bucket.pendingOperations = newPendingOperations;
-
-    this.pushQueue.insert(bucket);
+    bucket.setData(data, newPendingOperations);
   }
 
   triggerPushQueue() {

frontend/javascripts/oxalis/model/bucket_data_handling/layer_rendering_manager.js

@@ -89,6 +89,8 @@ function consumeBucketsFromArrayBuffer(
     const bucket = cube.getOrCreateBucket(bucketAddress);
 
     if (bucket.type !== "null") {
+      // This tells the bucket collection, that the buckets are necessary for rendering
+      bucket.markAsNeeded();
       bucketsWithPriorities.push({ bucket, priority });
     }
 
@@ -237,16 +239,14 @@ export default class LayerRenderingManager {
 
       pickingPromise.then(
         buffer => {
+          this.cube.markBucketsAsUnneeded();
           const bucketsWithPriorities = consumeBucketsFromArrayBuffer(
             buffer,
             this.cube,
             this.textureBucketManager.maximumCapacity,
           );
 
           const buckets = bucketsWithPriorities.map(({ bucket }) => bucket);
-          this.cube.markBucketsAsUnneeded();
-          // This tells the bucket collection, that the buckets are necessary for rendering
-          buckets.forEach(b => b.markAsNeeded());
 
           this.textureBucketManager.setActiveBuckets(
             buckets,