perf(tokenize): opt-in tokenizer worker pool for step 2

.github/workflows/ci.yml

@@ -74,6 +74,9 @@ jobs:
       - name: Build & test Rust tokenizer (rust_tokenizer)
         run: devenv shell -- bash -c 'cd tokenize/rustTokenizer && make && make test'
 
+      - name: "Tokenizer worker: same output as tokenBySha.pl"
+        run: devenv shell -- bash tests/test_token_worker.sh
+
       - name: "Pipeline runner: artifact staleness guard"
         run: devenv shell -- bash test_ensure_artifacts.sh
 

README.md

@@ -51,8 +51,8 @@ HTML views land in the sibling directory `../cregit-files/html`. See
 ## Preliminaries
 
 - Code is written in Scala, C++, Rust and Perl.
-- Platform: Linux x86_64 or macOS arm64 — the pinned `srcml` 1.1.0 parser is a
-  prebuilt binary available only for those platforms.
+- Platform: Linux x86_64 or macOS arm64 — `devenv` builds the pinned `srcml`
+  1.1.0 parser from its release tag (`nix/srcml.nix`).
 
 ## Prerequisites
 
@@ -88,7 +88,7 @@ This builds, in dependency order:
 
 | artifact                                                      | module              | toolchain                 |
 | ------------------------------------------------------------- | ------------------- | ------------------------- |
-| `tokenize/srcMLtoken/srcml2token`                              | C++ transcoder      | gcc + xerces-c            |
+| `tokenize/srcMLtoken/srcml2token`                              | srcML tokenizer     | gcc + xerces-c + libsrcml |
 | `tokenize/rustTokenizer` binary                                | Rust tokenizer      | cargo                     |
 | `blobExec/target/scala-2.13/blobExec-0.1.0-assembly.jar`       | tokenization driver | sbt, JDK 21               |
 | `{slickGitLog,persons,remapCommits}/target/scala-2.10/*-one-jar.jar` | history / persons / remap tools | sbt 0.13, JDK 8 |

blobExec/shard_build.sh

@@ -88,9 +88,9 @@ if [ -z "$THREADS" ]; then
 fi
 
 # Default tokenizer dispatcher (srcML for .c/.h, rustTokenizer for .rs, …); the
-# --srcml/--ctags paths are forwarded only to the srcML parser.
+# --srcml2token/--ctags paths are forwarded only to the srcML parser.
 if [ -z "$TOK_CMD" ]; then
-  TOK_CMD="$CREGIT/tokenize/tokenize.pl --srcml2token=$CREGIT/tokenize/srcMLtoken/srcml2token --srcml=$(command -v srcml) --ctags=$(command -v ctags)"
+  TOK_CMD="$CREGIT/tokenize/tokenize.pl --srcml2token=$CREGIT/tokenize/srcMLtoken/srcml2token --ctags=$(command -v ctags)"
 fi
 export BFG_TOKENIZE_CMD="$TOK_CMD"
 JAVA="$(command -v java)"

blobExec/src/main/scala/cregit/blobexec/BlobExec.scala

@@ -115,9 +115,13 @@ object BlobExec {
       inserter: ObjectInserter,
       timeoutSeconds: Int = DefaultTimeoutSeconds,
       onTimeout: () => Unit = () => (),
-      onParserCrash: () => Unit = () => ()
+      onParserCrash: () => Unit = () => (),
+      workerPool: Option[TokenizerWorkerPool] = None
   ): Outcome = {
-    val (exitCode, stdout, stderr) = invoke(bytes, origSha, filename, fullPath, command, timeoutSeconds)
+    val (exitCode, stdout, stderr) = workerPool match {
+      case Some(pool) => pool.invoke(bytes, filename, timeoutSeconds)
+      case None       => invoke(bytes, origSha, filename, fullPath, command, timeoutSeconds)
+    }
 
     if (exitCode == TimeoutExitCode) {
       // A timed-out child skips exactly one blob. Deliberately *not* routed

blobExec/src/main/scala/cregit/blobexec/Main.scala

@@ -92,7 +92,7 @@ object Main {
   // `raw` (not `s`): the mask example below contains a regex backslash, which a
   // processed-escape interpolator rejects. `$$` therefore renders a literal `$`.
   private val Usage =
-    raw"""Usage: blobExec [--abort-on-error] [--pipeline | --pipeline-trees | --shard=K/N] [--warm=<db>] [--mask-widened] [--tokenizer-identity=<ext>=<value>,...] [--retokenize=<ext>,...] [--memo-dir=<dir>] [--blob-timeout=<seconds>] [--stall-timeout=<seconds>] <src.git> <dst.git> <db.sqlite> <command> <fileMaskRegex>
+    raw"""Usage: blobExec [--abort-on-error] [--pipeline | --pipeline-trees | --shard=K/N] [--warm=<db>] [--mask-widened] [--tokenizer-identity=<ext>=<value>,...] [--retokenize=<ext>,...] [--memo-dir=<dir>] [--tokenizer-worker=<path>] [--blob-timeout=<seconds>] [--stall-timeout=<seconds>] <src.git> <dst.git> <db.sqlite> <command> <fileMaskRegex>
       |
       |  --abort-on-error  exit immediately (status 2) on the first non-zero
       |                    exit from <command>, instead of skipping that blob
@@ -188,6 +188,11 @@ object Main {
       |  --memo-dir=<dir>  the memo directory ($$BFG_MEMO_DIR) whose entries
       |                    --retokenize must purge. Only read with --retokenize;
       |                    passing it alone is an error rather than a no-op.
+      |  --tokenizer-worker=<path>
+      |                    send each blob to one of a pool of persistent
+      |                    tokenizer processes (tokenizeByBlobId/tokenWorker.pl)
+      |                    instead of running <command> once per blob. Its
+      |                    output is the same. Needs --pipeline or --pipeline-trees.
       |  --blob-timeout=<seconds>
       |                    wall-clock budget for one <command> invocation
       |                    (default ${BlobExec.DefaultTimeoutSeconds}). A child that exceeds it is
@@ -265,6 +270,7 @@ object Main {
     var tokenizerIdentity = TokenizerIdentity.empty
     var retokenizeExtensions: Set[String] = Set.empty
     var memoDir: Option[java.nio.file.Path] = None
+    var tokenizerWorker: Option[java.nio.file.Path] = None
     flags.foreach {
       case "--abort-on-error" => abortOnError = true
       case "--pipeline"       => pipeline = true
@@ -316,6 +322,8 @@ object Main {
           sys.exit(1)
         }
         warmPath = Some(p)
+      case t if t.startsWith("--tokenizer-worker=") =>
+        tokenizerWorker = Some(Paths.get(t.stripPrefix("--tokenizer-worker=")))
       case t if t.startsWith("--blob-timeout=") =>
         val spec = t.stripPrefix("--blob-timeout=")
         parsePositiveSeconds(spec) match {
@@ -426,6 +434,17 @@ object Main {
       sys.exit(1)
     }
 
+    tokenizerWorker.foreach { path =>
+      if (!Files.isRegularFile(path) || !Files.isExecutable(path)) {
+        System.err.println(s"Error: --tokenizer-worker [$path] is not an executable file")
+        sys.exit(1)
+      }
+      if (!(pipeline || pipelineTrees)) {
+        System.err.println("Error: --tokenizer-worker needs --pipeline or --pipeline-trees")
+        sys.exit(1)
+      }
+    }
+
     if (positional.length != 5) {
       System.err.println(Usage)
       sys.exit(1)
@@ -476,7 +495,8 @@ object Main {
         s"maskWidened=$maskWidened denylistEntries=${denylist.size} " +
         s"tokenizerIdentity=${if (tokenizerIdentity.isEmpty) "none" else tokenizerIdentity.render} " +
         s"retokenize=${if (retokenizeExtensions.isEmpty) "none" else retokenizeExtensions.toVector.sorted.mkString(",")} " +
-        s"memoDir=${memoDir.map(_.toString).getOrElse("none")}"
+        s"memoDir=${memoDir.map(_.toString).getOrElse("none")} " +
+        s"tokenizerWorker=${tokenizerWorker.map(_.toString).getOrElse("none")}"
     )
 
     val src: FileRepository = openSrc(srcPath)
@@ -565,15 +585,18 @@ object Main {
 
     val stats = try {
       val parallelism = math.max(1, Runtime.getRuntime.availableProcessors)
+      val workerPool = tokenizerWorker.map(path => new TokenizerWorkerPool(path.toString, parallelism))
       val walker = new Walker(
         src, dst, mapping, mask.r, command, abortOnError, parallelism,
         pipeline, pipelineTrees, shard,
         destinationMayContainObjects = incremental,
         blobTimeoutSeconds = blobTimeoutSeconds,
         stallTimeoutSeconds = stallTimeoutSeconds,
+        workerPool = workerPool,
         denylist = denylist
       )
-      walker.run()
+      try walker.run()
+      finally workerPool.foreach(_.close())
     } finally {
       mapping.close()
       dst.close()

blobExec/src/main/scala/cregit/blobexec/TokenizerWorkerPool.scala

@@ -0,0 +1,142 @@
+package cregit.blobexec
+
+import java.io.{BufferedOutputStream, DataInputStream, EOFException, IOException}
+import java.nio.charset.StandardCharsets.UTF_8
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
+import java.util.concurrent.{ArrayBlockingQueue, ConcurrentHashMap, CountDownLatch, TimeUnit}
+import scala.jdk.CollectionConverters._
+
+/** Persistent tokenizer processes that speak tokenizeByBlobId/WORKER_PROTOCOL.md,
+  * one request in flight per worker. */
+final class TokenizerWorkerPool(workerPath: String, size: Int) extends AutoCloseable {
+  import TokenizerWorkerPool._
+
+  private final class Worker(val index: Int, val process: Process) {
+    val stdin  = new BufferedOutputStream(process.getOutputStream)
+    val stdout = new DataInputStream(process.getInputStream)
+  }
+
+  private val idle    = new ArrayBlockingQueue[Worker](size)
+  private val workers = ConcurrentHashMap.newKeySet[Worker]()
+  private val closed  = new AtomicBoolean(false)
+
+  try (0 until size).foreach(addWorker)
+  catch { case failure: Throwable => close(); throw failure }
+
+  /** Same result as [[BlobExec.invoke]]: (exit, stdout, stderr), and
+    * [[BlobExec.TimeoutExitCode]] for a blob that ran out of time. */
+  def invoke(
+      bytes: Array[Byte],
+      filename: String,
+      timeoutSeconds: Int
+  ): (Int, Array[Byte], String) = {
+    val worker = idle.take()
+    val backstop = timeoutSeconds + BackstopSlackSeconds
+    within(backstop)(request(worker, bytes, filename, timeoutSeconds)) match {
+      case None =>
+        System.err.println(
+          s"blobExec: tokenizer worker ${worker.index} gave no response within ${backstop}s; replacing it")
+        replace(worker)
+        Timeout
+      case Some(Left(failure)) =>
+        replace(worker)
+        (1, Array.emptyByteArray, s"tokenizer worker ${worker.index} failed: $failure\n")
+      case Some(Right(response)) =>
+        if (worker.process.isAlive) idle.put(worker) else replace(worker)
+        if (response._1 == WorkerTimeoutExitCode) Timeout else response
+    }
+  }
+
+  override def close(): Unit =
+    if (closed.compareAndSet(false, true)) {
+      val all = workers.asScala.toVector
+      all.foreach(w => try w.stdin.close() catch { case _: IOException => () })
+      val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(CloseSeconds)
+      all.foreach { w =>
+        if (!w.process.waitFor(math.max(0L, deadline - System.nanoTime()), TimeUnit.NANOSECONDS)) kill(w.process)
+      }
+    }
+
+  private def addWorker(index: Int): Unit = {
+    val process = new ProcessBuilder(workerPath).redirectError(ProcessBuilder.Redirect.INHERIT).start()
+    val worker = new Worker(index, process)
+    within(ReadySeconds)(readLine(worker.stdout)) match {
+      case Some(Right("READY")) =>
+        workers.add(worker)
+        idle.put(worker)
+      case notReady =>
+        kill(process)
+        throw new IllegalStateException(s"tokenizer worker [$workerPath] did not start: $notReady")
+    }
+  }
+
+  private def replace(worker: Worker): Unit = {
+    workers.remove(worker)
+    kill(worker.process)
+    if (!closed.get()) addWorker(worker.index)
+  }
+
+  private def request(worker: Worker, bytes: Array[Byte], filename: String, timeoutSeconds: Int) = {
+    val name = filename.getBytes(UTF_8)
+    worker.stdin.write(s"REQ ${bytes.length} ${name.length} $timeoutSeconds\n".getBytes(UTF_8))
+    worker.stdin.write(name)
+    worker.stdin.write(bytes)
+    worker.stdin.flush()
+    readLine(worker.stdout).split(' ') match {
+      case Array("RES", exit, outLength, errLength) =>
+        val out = readBytes(worker.stdout, outLength.toInt)
+        (exit.toInt, out, new String(readBytes(worker.stdout, errLength.toInt), UTF_8))
+      case header => throw new IOException(s"malformed response header [${header.mkString(" ")}]")
+    }
+  }
+}
+
+object TokenizerWorkerPool {
+  /** What a worker answers when it killed the tokenizer for running out of time. */
+  val WorkerTimeoutExitCode = 124
+  private val Timeout = (BlobExec.TimeoutExitCode, Array.emptyByteArray, "")
+  /** The worker kills its own tokenizer at the budget. This only catches a wedged worker. */
+  private val BackstopSlackSeconds = 5L
+  private val ReadySeconds = 30L
+  private val CloseSeconds = 2L
+  private val MaxHeaderBytes = 8192
+
+  /** Kill the children first: once the worker is gone they are no longer its descendants. */
+  private def kill(process: Process): Unit = {
+    process.descendants().forEach(child => { child.destroyForcibly(); () })
+    process.destroyForcibly()
+    ()
+  }
+
+  /** Runs `read` on a daemon thread, so that a worker which never answers cannot block the caller. */
+  private def within[A](seconds: Long)(read: => A): Option[Either[Throwable, A]] = {
+    val result = new AtomicReference[Either[Throwable, A]]()
+    val done = new CountDownLatch(1)
+    val reader = new Thread(() => {
+      try result.set(Right(read))
+      catch { case failure: Throwable => result.set(Left(failure)) }
+      finally done.countDown()
+    }, "tokenizer-worker-read")
+    reader.setDaemon(true)
+    reader.start()
+    if (done.await(seconds, TimeUnit.SECONDS)) Some(result.get()) else None
+  }
+
+  private def readLine(in: DataInputStream): String = {
+    val line = new java.io.ByteArrayOutputStream
+    var next = in.read()
+    while (next != '\n') {
+      if (next == -1) throw new EOFException("tokenizer worker closed its stdout")
+      if (line.size() >= MaxHeaderBytes) throw new IOException(s"header longer than $MaxHeaderBytes bytes")
+      line.write(next)
+      next = in.read()
+    }
+    line.toString(UTF_8)
+  }
+
+  private def readBytes(in: DataInputStream, length: Int): Array[Byte] = {
+    val bytes = new Array[Byte](length)
+    in.readFully(bytes)
+    bytes
+  }
+}

blobExec/src/main/scala/cregit/blobexec/Walker.scala

@@ -60,6 +60,7 @@ final class Walker(
     destinationMayContainObjects: Boolean = true,
     blobTimeoutSeconds: Int = BlobExec.DefaultTimeoutSeconds,
     stallTimeoutSeconds: Int = Walker.DefaultStallTimeoutSeconds,
+    workerPool: Option[TokenizerWorkerPool] = None,
     // The shipped list by default, so a caller cannot forget it and hand a known
     // non-terminating blob to srcml. A parameter only so a test can supply its
     // own fixture; nothing at run time chooses a different list.
@@ -803,7 +804,8 @@ final class Walker(
         inserter     = workerInserter,
         timeoutSeconds = blobTimeoutSeconds,
         onTimeout      = () => { blobsTimedOut.increment(); failure.set("timeout") },
-        onParserCrash  = () => { blobsParserCrashed.increment(); failure.set("parser-crash") }
+        onParserCrash  = () => { blobsParserCrashed.increment(); failure.set("parser-crash") },
+        workerPool     = workerPool
       )
       val res = outcome match {
         case BlobExec.Outcome.Skip if failure.get() != null =>