performance: parallelize diffgraph application, store and load (#288)

Author: bbrehm

README.md

@@ -6,6 +6,21 @@ Code formatting is maintained via
 sbt scalafmt Test/scalafmt
 ```
 
+## Diverse notes
+By default, diffgraph application, deserialization from storage, and serialization to storage are all multi-threaded.
+
+This can be globally disabled via `flatgraph.misc.Misc.force_singlethreaded()`, for easier debugging. 
+
+In order to quickly glance the input of flatgraph files, you can extract the manifest json with `tail`, e.g. `tail someGraph.fg | jless`:
+Our output writer always places the manifest at the end, with a bunch of preceding newlines, such that this will not contain binary garbage.
+
+This is suitable for quick command-line debugging. However, that approach will fail if e.g. somebody appended two flatgraph files -- deserialization will
+read the file from the beginning, and find the offset of the true manifest from the header, and ignore trailing garbage like an appended fake manifest.
+So don't dare to do security checks with that!
+
+
+
+
 ## Core Features
 - [x] Access nodes and neighbors
 - [x] Add nodes and edges

core/src/main/scala/flatgraph/DiffGraphApplier.scala

@@ -0,0 +1,27 @@
+package flatgraph.misc
+
+import java.util.concurrent
+
+object Misc {
+
+  @volatile var _overrideExecutor: Option[concurrent.ExecutorService] => concurrent.ExecutorService = defaultExecutorProvider
+
+  def force_singlethreaded(): Unit = {
+    // This one is magic -- it can get garbage collected, no manual shutdown required!
+    // Unfortunately this behavior is apparently not documented officially :(
+    // But that is the behavior on java8 and java23 and, presumably, everywhere we care about:
+    // https://github.com/openjdk/jdk8u/blob/1a6e3a5ea32d5c671cb46a590046f16426089921/jdk/src/share/classes/java/util/concurrent/Executors.java#L170
+    // https://github.com/openjdk/jdk23u/blob/9101cc14972ce6bdeb966e67bcacc8b693c37d0a/src/java.base/share/classes/java/util/concurrent/Executors.java#L192
+    this._overrideExecutor = (something: Option[concurrent.ExecutorService]) => concurrent.Executors.newSingleThreadExecutor()
+  }
+
+  def defaultExecutorProvider(requested: Option[concurrent.ExecutorService]): concurrent.ExecutorService = requested.getOrElse {
+    java.lang.Thread.currentThread() match {
+      case fjt: concurrent.ForkJoinWorkerThread => fjt.getPool
+      case _                                    => concurrent.ForkJoinPool.commonPool()
+    }
+  }
+
+  def maybeOverrideExecutor(requested: Option[concurrent.ExecutorService]): concurrent.ExecutorService =
+    this._overrideExecutor.apply(requested)
+}

core/src/main/scala/flatgraph/misc/Misc.scala

@@ -1,8 +1,8 @@
 package flatgraph.storage
 
-import com.github.luben.zstd.Zstd
-import flatgraph.*
+import flatgraph.{AccessHelpers, FreeSchema, GNode, Graph, Schema}
 import flatgraph.Edge.Direction
+import flatgraph.misc.Misc
 import flatgraph.storage.Manifest.{GraphItem, OutlineStorage}
 
 import java.nio.channels.FileChannel
@@ -11,29 +11,42 @@ import java.nio.file.Path
 import java.nio.{ByteBuffer, ByteOrder}
 import java.util.Arrays
 import scala.collection.mutable
+import java.util.concurrent
 
 object Deserialization {
 
-  def readGraph(storagePath: Path, schemaMaybe: Option[Schema], persistOnClose: Boolean = true): Graph = {
+  def readGraph(
+    storagePath: Path,
+    schemaMaybe: Option[Schema],
+    persistOnClose: Boolean = true,
+    requestedExecutor: Option[concurrent.ExecutorService] = None
+  ): Graph = {
+    val executor    = Misc.maybeOverrideExecutor(requestedExecutor)
     val fileChannel = new java.io.RandomAccessFile(storagePath.toAbsolutePath.toFile, "r").getChannel
+    val queue       = mutable.ArrayBuffer[concurrent.Future[Any]]()
+    val zstdCtx     = new ZstdWrapper.ZstdCtx
+    def submitJob[T](block: => T): concurrent.Future[T] = {
+      val res = executor.submit((() => block))
+      queue.addOne(res.asInstanceOf[concurrent.Future[Any]])
+      res
+    }
+
     try {
       // fixme: Use convenience methods from schema to translate string->id. Fix after we get strict schema checking.
       val manifest = GraphItem.read(readManifest(fileChannel))
-      val pool     = readPool(manifest, fileChannel)
+      val pool     = submitJob { readPool(manifest, fileChannel, zstdCtx) }
       val schema   = schemaMaybe.getOrElse(freeSchemaFromManifest(manifest))
       val storagePathMaybe =
         if (persistOnClose) Option(storagePath)
         else None
       val g         = new Graph(schema, storagePathMaybe)
       val nodekinds = mutable.HashMap[String, Short]()
       for (nodeKind <- g.schema.nodeKinds) nodekinds(g.schema.getNodeLabel(nodeKind)) = nodeKind.toShort
-      val kindRemapper = Array.fill(manifest.nodes.size)(-1.toShort)
       val nodeRemapper = new Array[Array[GNode]](manifest.nodes.length)
       for {
         (nodeItem, idx) <- manifest.nodes.zipWithIndex
         nodeKind        <- nodekinds.get(nodeItem.nodeLabel)
       } {
-        kindRemapper(idx) = nodeKind
         val nodes = new Array[GNode](nodeItem.nnodes)
         for (seq <- Range(0, nodes.length)) nodes(seq) = g.schema.makeNode(g, nodeKind, seq)
         g.nodesArray(nodeKind) = nodes
@@ -66,11 +79,17 @@ object Deserialization {
         val direction = Direction.fromOrdinal(edgeItem.inout)
         if (nodeKind.isDefined && edgeKind.isDefined) {
           val pos = g.schema.neighborOffsetArrayIndex(nodeKind.get, direction, edgeKind.get)
-          g.neighbors(pos) = deltaDecode(readArray(fileChannel, edgeItem.qty, nodeRemapper, pool).asInstanceOf[Array[Int]])
-          g.neighbors(pos + 1) = readArray(fileChannel, edgeItem.neighbors, nodeRemapper, pool)
-          val property = readArray(fileChannel, edgeItem.property, nodeRemapper, pool)
-          if (property != null)
-            g.neighbors(pos + 2) = property
+          submitJob {
+            g.neighbors(pos) = deltaDecode(readArray(fileChannel, edgeItem.qty, nodeRemapper, pool, zstdCtx).asInstanceOf[Array[Int]])
+          }
+          submitJob {
+            g.neighbors(pos + 1) = readArray(fileChannel, edgeItem.neighbors, nodeRemapper, pool, zstdCtx)
+          }
+          submitJob {
+            val property = readArray(fileChannel, edgeItem.property, nodeRemapper, pool, zstdCtx)
+            if (property != null)
+              g.neighbors(pos + 2) = property
+          }
         }
       }
 
@@ -91,12 +110,18 @@ object Deserialization {
         val propertyKind = propertykinds.get((property.nodeLabel, property.propertyLabel))
         if (nodeKind.isDefined && propertyKind.isDefined) {
           val pos = g.schema.propertyOffsetArrayIndex(nodeKind.get, propertyKind.get)
-          g.properties(pos) = deltaDecode(readArray(fileChannel, property.qty, nodeRemapper, pool).asInstanceOf[Array[Int]])
-          g.properties(pos + 1) = readArray(fileChannel, property.property, nodeRemapper, pool)
+          submitJob {
+            g.properties(pos) = deltaDecode(readArray(fileChannel, property.qty, nodeRemapper, pool, zstdCtx).asInstanceOf[Array[Int]])
+          }
+          submitJob { g.properties(pos + 1) = readArray(fileChannel, property.property, nodeRemapper, pool, zstdCtx) }
         }
       }
+      queue.foreach { _.get() }
       g
-    } finally fileChannel.close()
+    } catch {
+      case ex: java.util.concurrent.ExecutionException =>
+        throw ex.getCause()
+    } finally { fileChannel.close(); zstdCtx.close(); }
   }
 
   private def freeSchemaFromManifest(manifest: Manifest.GraphItem): FreeSchema = {
@@ -171,23 +196,17 @@ object Deserialization {
 
   }
 
-  private def readPool(manifest: GraphItem, fileChannel: FileChannel): Array[String] = {
-    val stringPoolLength = ZstdWrapper(
-      Zstd
-        .decompress(
-          fileChannel.map(FileChannel.MapMode.READ_ONLY, manifest.stringPoolLength.startOffset, manifest.stringPoolLength.compressedLength),
-          manifest.stringPoolLength.decompressedLength
-        )
-        .order(ByteOrder.LITTLE_ENDIAN)
-    )
-    val stringPoolBytes = ZstdWrapper(
-      Zstd
-        .decompress(
-          fileChannel.map(FileChannel.MapMode.READ_ONLY, manifest.stringPoolBytes.startOffset, manifest.stringPoolBytes.compressedLength),
-          manifest.stringPoolBytes.decompressedLength
-        )
-        .order(ByteOrder.LITTLE_ENDIAN)
-    )
+  private def readPool(manifest: GraphItem, fileChannel: FileChannel, zstdCtx: ZstdWrapper.ZstdCtx): Array[String] = {
+    val stringPoolLength = zstdCtx
+      .decompress(
+        fileChannel.map(FileChannel.MapMode.READ_ONLY, manifest.stringPoolLength.startOffset, manifest.stringPoolLength.compressedLength),
+        manifest.stringPoolLength.decompressedLength
+      )
+    val stringPoolBytes = zstdCtx
+      .decompress(
+        fileChannel.map(FileChannel.MapMode.READ_ONLY, manifest.stringPoolBytes.startOffset, manifest.stringPoolBytes.compressedLength),
+        manifest.stringPoolBytes.decompressedLength
+      )
     val poolBytes = new Array[Byte](manifest.stringPoolBytes.decompressedLength)
     stringPoolBytes.get(poolBytes)
     val pool    = new Array[String](manifest.stringPoolLength.decompressedLength >> 2)
@@ -215,11 +234,18 @@ object Deserialization {
     a
   }
 
-  private def readArray(channel: FileChannel, ptr: OutlineStorage, nodes: Array[Array[GNode]], stringPool: Array[String]): Array[?] = {
+  private def readArray(
+    channel: FileChannel,
+    ptr: OutlineStorage,
+    nodes: Array[Array[GNode]],
+    stringPoolFuture: concurrent.Future[Array[String]],
+    zstdCtx: ZstdWrapper.ZstdCtx
+  ): Array[?] = {
     if (ptr == null) return null
-    val dec = ZstdWrapper(
-      Zstd.decompress(channel.map(FileChannel.MapMode.READ_ONLY, ptr.startOffset, ptr.compressedLength), ptr.decompressedLength)
-    ).order(ByteOrder.LITTLE_ENDIAN)
+    if (ptr.typ == StorageType.String) stringPoolFuture.get()
+
+    val dec =
+      zstdCtx.decompress(channel.map(FileChannel.MapMode.READ_ONLY, ptr.startOffset, ptr.compressedLength), ptr.decompressedLength)
     ptr.typ match {
       case StorageType.Bool =>
         val bytes = new Array[Byte](dec.limit())
@@ -253,9 +279,10 @@ object Deserialization {
         dec.asDoubleBuffer().get(res)
         res
       case StorageType.String =>
-        val res    = new Array[String](dec.limit() >> 2)
-        val intbuf = dec.asIntBuffer()
-        var idx    = 0
+        val stringPool = stringPoolFuture.get()
+        val res        = new Array[String](dec.limit() >> 2)
+        val intbuf     = dec.asIntBuffer()
+        var idx        = 0
         while (idx < res.length) {
           val offset = intbuf.get(idx)
           if (offset >= 0) res(idx) = stringPool(offset)

core/src/main/scala/flatgraph/storage/Deserialization.scala

@@ -34,8 +34,8 @@ object Manifest {
     var nodes: Array[NodeItem],
     var edges: Array[EdgeItem],
     var properties: Array[PropertyItem],
-    val stringPoolLength: OutlineStorage,
-    val stringPoolBytes: OutlineStorage
+    val stringPoolLength: OutlineStorage = new OutlineStorage(StorageType.Int),
+    val stringPoolBytes: OutlineStorage = new OutlineStorage(StorageType.Byte)
   ) {
     var version = 0
   }
@@ -96,9 +96,9 @@ object Manifest {
     val nodeLabel: String,
     val edgeLabel: String,
     val inout: Byte, // 0: Incoming, 1: Outgoing; see Edge.Direction enum
-    var qty: OutlineStorage,
-    var neighbors: OutlineStorage,
-    var property: OutlineStorage
+    var qty: OutlineStorage = new OutlineStorage,
+    var neighbors: OutlineStorage = new OutlineStorage,
+    var property: OutlineStorage = new OutlineStorage
   ) {
     Edge.Direction.verifyEncodingRange(inout)
   }
@@ -122,11 +122,20 @@ object Manifest {
     }
   }
 
-  class PropertyItem(val nodeLabel: String, val propertyLabel: String, var qty: OutlineStorage, var property: OutlineStorage)
+  class PropertyItem(
+    val nodeLabel: String,
+    val propertyLabel: String,
+    var qty: OutlineStorage = new OutlineStorage,
+    var property: OutlineStorage = new OutlineStorage
+  )
 
   object OutlineStorage {
     def write(item: OutlineStorage): ujson.Value = {
       if (item == null) return ujson.Null
+      if (item.typ == null) {
+        assert(item.startOffset == -1L && item.compressedLength == -1 && item.decompressedLength == -1, s"bad OutlineStorage ${item}")
+        return ujson.Null
+      }
       val res = ujson.Obj()
       res(Keys.Type) = item.typ
       res(Keys.StartOffset) = ujson.Num(item.startOffset.toDouble)
@@ -143,17 +152,25 @@ object Manifest {
 
     def read(item: ujson.Value): OutlineStorage = {
       if (item.isNull) return null
-      val res = new OutlineStorage(item.obj(Keys.Type).str)
+      val res = new OutlineStorage
+      res.typ = item.obj(Keys.Type).str
       res.startOffset = item.obj(Keys.StartOffset).num.toLong
       res.compressedLength = item.obj(Keys.CompressedLength).num.toInt
       res.decompressedLength = item.obj(Keys.DecompressedLength).num.toInt
       res
     }
   }
 
-  class OutlineStorage(var typ: String) {
+  class OutlineStorage {
+    var typ: String             = null
     var startOffset: Long       = -1L
     var compressedLength: Int   = -1
     var decompressedLength: Int = -1
+    def this(_typ: String) = {
+      this()
+      this.typ = _typ
+    }
+
+    override def toString: String = super.toString + s"($typ, $startOffset, $compressedLength, $decompressedLength)"
   }
 }