Dynamic bucket allocation & miscellaneous volume improvements

CHANGELOG.unreleased.md

@@ -31,6 +31,7 @@ For upgrade instructions, please check the [migration guide](MIGRATIONS.released
 - Changed a number of API routes from GET to POST to avoid unwanted side effects. [#6023](https://github.com/scalableminds/webknossos/pull/6023)
 - Removed unused datastore route `checkInbox` (use `checkInboxBlocking` instead). [#6023](https://github.com/scalableminds/webknossos/pull/6023)
 - Migrated to Google Analytics 4. [#6031](https://github.com/scalableminds/webknossos/pull/6031)
+- Improved stability and speed of volume annotations when annotating large areas. [#6055](https://github.com/scalableminds/webknossos/pull/6055)
 
 ### Fixed
 - Fixed volume-related bugs which could corrupt the volume data in certain scenarios. [#5955](https://github.com/scalableminds/webknossos/pull/5955)

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;
@@ -230,6 +237,7 @@ export class DataBucket {
     this.data = null;
     // Remove all event handlers (see https://github.com/ai/nanoevents#remove-all-listeners)
     this.emitter.events = {};
+    this.trigger("bucketCollected");
   }
 
   needsRequest(): boolean {
@@ -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,44 +321,50 @@ 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} 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 {
@@ -326,11 +373,9 @@ class DataCube {
     for (const bucket of this.buckets) {
       if (bucket != null) {
         this.collectBucket(bucket);
-        bucket.trigger("bucketCollected");
       }
     }
     this.buckets = [];
-    this.bucketCount = 0;
     this.bucketIterator = 0;
   }
 
@@ -648,10 +693,7 @@ class DataCube {
     if (bucket.type === "null") {
       return;
     }
-    bucket.setData(data);
-    bucket.pendingOperations = newPendingOperations;
-
-    this.pushQueue.insert(bucket);
+    bucket.setData(data, newPendingOperations);
   }
 
   triggerPushQueue() {