Improve volume and segment index performance (#8460)

* WIP: Improve segment index performance

* implementation with triple cast on every element, not faster than scala

* measure what else is slow

* wip cleanup, measure other stuff

* wip: use volume bucket buffer, fossil multi-put

* logging

* WIP rewrite segmentIndexBuffer to include caching, use Set

* use new segment index buffer functionality

* WIP fossil multi-get for buckets

* fix setting volumeBucketDataHasChanged

* do not parallelize prefill

* use batched multi-get and multi-put

* connect to temporaryStore, compress if needed

* load buckets from temporarystore, remote fallback layer

* load buckets from fallback layer in one request

* some cleanup

* readerOnly segmentIndexBuffer

* cleanup

* cleanup

* format, remove logging

* debug wrong volume values for editable mapping (id seems to be mapped with outdated mapping?)

* request correct version of editable mapping data

* changelog, migration

* cleanup

* make cpp types more explicit

* clean up cpp code

* const, try/catch

* clang-format

* size_t for index

* Do not cache EditableMappingBucketProvider across versions

* pr feedback part 1; consistent volumeLayer naming

* use fossil multi-get when requesting multiple segmentIds from segment index

* use map for more efficient lookups during gathering segmentIndex values

* add one more conversion from set to seq for less map overhead

* Update webknossos-tracingstore/app/com/scalableminds/webknossos/tracingstore/tracings/volume/VolumeSegmentIndexBuffer.scala

Co-authored-by: MichaelBuessemeyer <[email protected]>

* Update webknossos-tracingstore/app/com/scalableminds/webknossos/tracingstore/tracings/volume/VolumeSegmentIndexBuffer.scala

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>

---------

Co-authored-by: MichaelBuessemeyer <[email protected]>
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>

Author: 3 people

.editorconfig

@@ -12,6 +12,9 @@ insert_final_newline = true
 [*.py]
 indent_size = 4
 
+[*.cpp]
+indent_size = 4
+
 [*.md]
 trim_trailing_whitespace = false
 

CHANGELOG.unreleased.md

@@ -13,6 +13,7 @@ For upgrade instructions, please check the [migration guide](MIGRATIONS.released
 ### Added
 - Added a credit system making payment for long running jobs possible. For now it is in testing phase. [#8352](https://github.com/scalableminds/webknossos/pull/8352)
 - The maximum available storage of an organization is now enforced during upload. [#8385](https://github.com/scalableminds/webknossos/pull/8385)
+- Performance improvements for volume annotation save requests. [#8460](https://github.com/scalableminds/webknossos/pull/8460)
 
 ### Changed
 
@@ -22,7 +23,8 @@ For upgrade instructions, please check the [migration guide](MIGRATIONS.released
 - Re-enabled jobs planned to be paid with credits for organizations without a paid plan. [#8478](https://github.com/scalableminds/webknossos/pull/8478)
 - Fixed that the dataset extent tooltip in the right details bar in the dashboard did not properly update when switching datasets. [#8477](https://github.com/scalableminds/webknossos/pull/8477)
 - Fixed a bug where task creation with volume zip as input would fail. [#8468](https://github.com/scalableminds/webknossos/pull/8468)
-- Fixed that a warning message about a newer version of an annotation was shown multiple times. [#8486](https://github.com/scalableminds/webknossos/pull/8486)
+- Fixed a bug where segment statistics would sometimes be wrong in case of an on-disk segmentation fallback layer with segment index file. [#8460](https://github.com/scalableminds/webknossos/pull/8460)
+- Fixed a bug where sometimes outdated segment statistics would be displayed. [#8460](https://github.com/scalableminds/webknossos/pull/8460)
 
 ### Removed
 

MIGRATIONS.unreleased.md

@@ -8,6 +8,8 @@ User-facing changes are documented in the [changelog](CHANGELOG.released.md).
 ## Unreleased
 [Commits](https://github.com/scalableminds/webknossos/compare/25.03.0...HEAD)
 
+ - New FossilDB version `0.1.37` (`master__525:` on dockerhub) is required. [#8460](https://github.com/scalableminds/webknossos/pull/8460)
+
 ### Postgres Evolutions:
 - [129-credit-transactions.sql](conf/evolutions/129-credit-transactions.sql)
 - [130-replace-text-types.sql](conf/evolutions/130-replace-text-types.sql)

docker-compose.yml

@@ -273,7 +273,7 @@ services:
 
   # FossilDB
   fossildb:
-    image: scalableminds/fossildb:master__510
+    image: scalableminds/fossildb:master__525
     command:
       - fossildb
       - -c

fossildb/version

@@ -1 +1 @@
-0.1.34
+0.1.37

tools/proxy/proxy.js

@@ -96,7 +96,7 @@ proxy.on("error", (err, req, res) => {
       <head>
         <title>503 Service Unavailable</title>
         <script type="text/javascript">
-          let countdown = 5;
+          let countdown = 3;
           function updateCountdown() {
             document.getElementById('countdown').textContent = countdown;
             countdown--;
@@ -114,7 +114,7 @@ proxy.on("error", (err, req, res) => {
       <body>
         <h1>Bad gateway</h1>
         <p>The server might still be starting up, please try again in a few seconds or check console output.</p>
-        <p>Reloading in <span id="countdown">5</span> seconds...</p>
+        <p>Reloading in <span id="countdown">3</span> seconds...</p>
       </body>
     </html>
   `);

util/src/main/scala/com/scalableminds/util/tools/BoxImplicits.scala

@@ -23,6 +23,18 @@ trait BoxImplicits {
     val failures = boxes.collect { case f: Failure => f }
     if (failures.isEmpty) None else Some(failures.map(_.msg))
   }
+
+  def assertNoFailure(boxes: Seq[Box[_]]): Box[Unit] = {
+    val firstFailure = boxes.find {
+      case _: Failure => true
+      case _          => false
+    }
+    firstFailure match {
+      case Some(failure) => Failure(s"At least one failure contained in list of ${boxes.length} boxes: $failure")
+      case None          => Full(())
+    }
+  }
+
 }
 
 object BoxImplicits extends BoxImplicits

webknossos-datastore/app/com/scalableminds/webknossos/datastore/helpers/NativeBucketScanner.scala

@@ -0,0 +1,8 @@
+package com.scalableminds.webknossos.datastore.helpers
+
+import com.github.sbt.jni.nativeLoader
+
+@nativeLoader("webknossosJni0")
+class NativeBucketScanner() {
+  @native def collectSegmentIds(bucketBytes: Array[Byte], bytesPerElement: Int, isSigned: Boolean): Array[Long]
+}

webknossos-datastore/app/com/scalableminds/webknossos/datastore/helpers/SegmentStatistics.scala

@@ -2,7 +2,7 @@ package com.scalableminds.webknossos.datastore.helpers
 
 import com.scalableminds.util.geometry.{BoundingBox, Vec3Int}
 import com.scalableminds.util.tools.{Fox, FoxImplicits}
-import com.scalableminds.webknossos.datastore.geometry.ListOfVec3IntProto
+import com.scalableminds.webknossos.datastore.geometry.Vec3IntProto
 import com.scalableminds.webknossos.datastore.models.datasource.DataLayer
 import com.scalableminds.webknossos.datastore.models.{AdditionalCoordinate, SegmentInteger}
 import play.api.libs.json.{Json, OFormat}
@@ -22,15 +22,15 @@ trait SegmentStatistics extends ProtoGeometryImplicits with FoxImplicits {
   def calculateSegmentVolume(segmentId: Long,
                              mag: Vec3Int,
                              additionalCoordinates: Option[Seq[AdditionalCoordinate]],
-                             getBucketPositions: (Long, Vec3Int) => Fox[ListOfVec3IntProto],
+                             getBucketPositions: (Long, Vec3Int) => Fox[Set[Vec3IntProto]],
                              getTypedDataForBucketPosition: (
                                  Vec3Int,
                                  Vec3Int,
                                  Option[Seq[AdditionalCoordinate]]) => Fox[Array[SegmentInteger]])(
       implicit ec: ExecutionContext): Fox[Long] =
     for {
-      bucketPositionsProtos: ListOfVec3IntProto <- getBucketPositions(segmentId, mag)
-      bucketPositionsInMag = bucketPositionsProtos.values.map(vec3IntFromProto)
+      bucketPositionsProtos: Set[Vec3IntProto] <- getBucketPositions(segmentId, mag)
+      bucketPositionsInMag = bucketPositionsProtos.map(vec3IntFromProto)
       volumeBoxes <- Fox.serialSequence(bucketPositionsInMag.toList)(bucketPosition =>
         for {
           dataTyped: Array[SegmentInteger] <- getTypedDataForBucketPosition(bucketPosition, mag, additionalCoordinates)
@@ -43,15 +43,15 @@ trait SegmentStatistics extends ProtoGeometryImplicits with FoxImplicits {
   def calculateSegmentBoundingBox(segmentId: Long,
                                   mag: Vec3Int,
                                   additionalCoordinates: Option[Seq[AdditionalCoordinate]],
-                                  getBucketPositions: (Long, Vec3Int) => Fox[ListOfVec3IntProto],
+                                  getBucketPositions: (Long, Vec3Int) => Fox[Set[Vec3IntProto]],
                                   getTypedDataForBucketPosition: (
                                       Vec3Int,
                                       Vec3Int,
                                       Option[Seq[AdditionalCoordinate]]) => Fox[Array[SegmentInteger]])(
       implicit ec: ExecutionContext): Fox[BoundingBox] =
     for {
-      allBucketPositions: ListOfVec3IntProto <- getBucketPositions(segmentId, mag)
-      relevantBucketPositions = filterOutInnerBucketPositions(allBucketPositions)
+      allBucketPositions: Set[Vec3IntProto] <- getBucketPositions(segmentId, mag)
+      relevantBucketPositions = filterOutInnerBucketPositions(allBucketPositions.toSeq)
       boundingBoxMutable = scala.collection.mutable.ListBuffer[Int](Int.MaxValue,
                                                                     Int.MaxValue,
                                                                     Int.MaxValue,
@@ -78,16 +78,16 @@ trait SegmentStatistics extends ProtoGeometryImplicits with FoxImplicits {
         )
 
   // The buckets that form the outer walls of the bounding box are relevant (in each of those the real min/max voxel positions could occur)
-  private def filterOutInnerBucketPositions(bucketPositions: ListOfVec3IntProto): Seq[Vec3Int] =
-    if (bucketPositions.values.isEmpty) List.empty
+  private def filterOutInnerBucketPositions(bucketPositions: Seq[Vec3IntProto]): Seq[Vec3Int] =
+    if (bucketPositions.isEmpty) Seq.empty
     else {
-      val minX = bucketPositions.values.map(_.x).min
-      val minY = bucketPositions.values.map(_.y).min
-      val minZ = bucketPositions.values.map(_.z).min
-      val maxX = bucketPositions.values.map(_.x).max
-      val maxY = bucketPositions.values.map(_.y).max
-      val maxZ = bucketPositions.values.map(_.z).max
-      bucketPositions.values
+      val minX = bucketPositions.map(_.x).min
+      val minY = bucketPositions.map(_.y).min
+      val minZ = bucketPositions.map(_.z).min
+      val maxX = bucketPositions.map(_.x).max
+      val maxY = bucketPositions.map(_.y).max
+      val maxZ = bucketPositions.map(_.z).max
+      bucketPositions
         .filter(pos =>
           pos.x == minX || pos.x == maxX || pos.y == minY || pos.y == maxY || pos.z == minZ || pos.z == maxZ)
         .map(vec3IntFromProto)

webknossos-datastore/app/com/scalableminds/webknossos/datastore/services/SegmentIndexFileService.scala

@@ -5,7 +5,7 @@ import com.scalableminds.util.geometry.{BoundingBox, Vec3Int}
 import com.scalableminds.util.io.PathUtils
 import com.scalableminds.util.tools.{Fox, FoxImplicits}
 import com.scalableminds.webknossos.datastore.DataStoreConfig
-import com.scalableminds.webknossos.datastore.geometry.ListOfVec3IntProto
+import com.scalableminds.webknossos.datastore.geometry.Vec3IntProto
 import com.scalableminds.webknossos.datastore.helpers.SegmentStatistics
 import com.scalableminds.webknossos.datastore.models.datasource.DataLayer
 import com.scalableminds.webknossos.datastore.models.requests.{
@@ -176,7 +176,7 @@ class SegmentIndexFileService @Inject()(config: DataStoreConfig,
       organizationId: String,
       datasetDirectoryName: String,
       dataLayerName: String,
-      mappingName: Option[String])(segmentOrAgglomerateId: Long, mag: Vec3Int): Fox[ListOfVec3IntProto] =
+      mappingName: Option[String])(segmentOrAgglomerateId: Long, mag: Vec3Int): Fox[Set[Vec3IntProto]] =
     for {
       segmentIds <- getSegmentIdsForAgglomerateIdIfNeeded(organizationId,
                                                           datasetDirectoryName,
@@ -185,8 +185,8 @@ class SegmentIndexFileService @Inject()(config: DataStoreConfig,
                                                           mappingName)
       positionsPerSegment <- Fox.serialCombined(segmentIds)(segmentId =>
         getBucketPositions(organizationId, datasetDirectoryName, dataLayerName, segmentId, mag))
-      positionsCollected = positionsPerSegment.flatten.distinct
-    } yield ListOfVec3IntProto.of(positionsCollected.map(vec3IntToProto))
+      positionsCollected = positionsPerSegment.flatten.toSet.map(vec3IntToProto)
+    } yield positionsCollected
 
   private def getBucketPositions(organizationId: String,
                                  datasetDirectoryName: String,