Repository navigation
perf: keep tokenizer workers alive for step 2 #95
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Closed
Closed
Changes from all commits
Commits
Show all changes
25 commits
Select commit
Hold shift + click to select a range
72604e4
perf(blobExec): Reuse tokenizer workers
EllianCarlos 77e332f
perf(tokenizeByBlobId): Add persistent tokenizer worker
EllianCarlos e75edff
fix(tokenWorker): Silence perl exec warning on the failure fallback
EllianCarlos 99eb302
ci: Run the tokenizer worker protocol tests
EllianCarlos d60935a
perf(tokenWorker): Optional per-request trace via BFG_WORKER_TRACE
EllianCarlos 192395f
perf(tokenWorker): Run the srcML chain in-process
EllianCarlos 5e6623c
fix: restore the exec bit on tokenize.pl and tokenBySha.pl
EllianCarlos 4d8cd87
fix(tokenWorker): route .C and .H to C++, as tokenBySha.pl does
EllianCarlos c439ea4
fix(tests): name the timeout process without exec -a
EllianCarlos 5db0fc1
fix(blobExec): clear worker diagnostics per request
EllianCarlos 360f4d9
fix(tokenWorker): give helper stderr to its request
EllianCarlos 9d67aa0
fix(blobExec): refuse --tokenizer-worker outside pipelines
EllianCarlos a4b0d8d
refactor(blobExec): take --tokenizer-worker=<path>
EllianCarlos 1b4d214
refactor(blobExec): pass the budget to each worker request
EllianCarlos 4404bd5
refactor(blobExec): reuse ChildRunner's lifetime and kill
EllianCarlos 8c9c563
refactor(blobExec): drop the worker env map
EllianCarlos e64608c
refactor(blobExec): drop the worker pool shutdown hook
EllianCarlos 16506b5
refactor(tokenWorker): drop the BFG_WORKER_TRACE trace
EllianCarlos b3a7d4c
test: drop test-only hooks from the worker code
EllianCarlos 79a4786
docs(tokenWorker): bring WORKER_PROTOCOL.md up to date
EllianCarlos 08723d1
refactor(tests): split the token worker driver into checks
EllianCarlos 653743b
style(tokenWorker): drop comments that restate the code
EllianCarlos d34faca
refactor(tokenWorker): split the long worker subs
EllianCarlos dd713bc
refactor(blobExec): share the pool's timed read
EllianCarlos b8d031d
refactor(tokenize): take the DECL lines out of merge
EllianCarlos File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
258 changes: 258 additions & 0 deletions
258
blobExec/src/main/scala/cregit/blobexec/TokenizerWorkerPool.scala
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,258 @@ | ||
| package cregit.blobexec | ||
|
|
||
| import java.io.{BufferedInputStream, BufferedOutputStream, BufferedReader, EOFException, IOException, InputStream, InputStreamReader} | ||
| 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._ | ||
|
|
||
| final class TokenizerWorkerPool( | ||
| workerCommand: Seq[String], | ||
| size: Int | ||
| ) extends AutoCloseable { | ||
|
|
||
| require(workerCommand.nonEmpty, "workerCommand must not be empty") | ||
| require(size > 0, "worker pool size must be positive") | ||
|
|
||
| private val WorkerTimeoutExitCode = 124 | ||
| private val ReadySeconds = 30L | ||
|
|
||
| private final class Response(val exitCode: Int, val stdout: Array[Byte], val stderr: String) | ||
|
|
||
| private final class Worker( | ||
| val index: Int, | ||
| val process: Process, | ||
| val stdin: BufferedOutputStream, | ||
| val stdout: BufferedInputStream | ||
| ) { | ||
| private val diagnostics = new StringBuilder | ||
|
|
||
| def appendDiagnostic(line: String): Unit = diagnostics.synchronized { | ||
| diagnostics.append(line).append('\n') | ||
| } | ||
|
|
||
| def takeDiagnostics(): String = diagnostics.synchronized { | ||
| try diagnostics.toString finally diagnostics.clear() | ||
| } | ||
| } | ||
|
|
||
| private val idle = new ArrayBlockingQueue[Worker](size) | ||
| private val workers = ConcurrentHashMap.newKeySet[Worker]() | ||
| private val closed = new AtomicBoolean(false) | ||
|
|
||
| try { | ||
| (0 until size).foreach { index => | ||
| val worker = startWorker(index) | ||
| workers.add(worker) | ||
| idle.put(worker) | ||
| } | ||
| } catch { | ||
| case failure: Throwable => | ||
| close() | ||
| throw failure | ||
| } | ||
|
|
||
| def invoke( | ||
| bytes: Array[Byte], | ||
| origSha: String, | ||
| filename: String, | ||
| fullPath: String, | ||
| timeoutSeconds: Int | ||
| ): ChildRunner.Outcome = { | ||
| require(timeoutSeconds > 0, "worker timeout must be positive") | ||
| if (closed.get()) throw new IllegalStateException("tokenizer worker pool is closed") | ||
|
|
||
| val worker = idle.take() | ||
| worker.takeDiagnostics() | ||
| try writeRequest(worker, bytes, origSha, filename, fullPath, timeoutSeconds) | ||
| catch { case failure: IOException => return failedWorker(worker, failure) } | ||
|
|
||
| val backstopSeconds = ChildRunner.maxLifetimeSeconds(timeoutSeconds) | ||
| readWithin(s"token-worker-response-${worker.index}", backstopSeconds)(readResponse(worker.stdout)) match { | ||
| case None => | ||
| replaceWorker(worker) | ||
| ChildRunner.Outcome.Killed(s"no response within ${backstopSeconds}s; worker replaced") | ||
| case Some(Left(failure)) => | ||
| failedWorker(worker, failure) | ||
| case Some(Right(result)) => | ||
| returnWorker(worker) | ||
| if (result.exitCode == WorkerTimeoutExitCode) | ||
| ChildRunner.Outcome.Killed(s"no exit within ${timeoutSeconds}s") | ||
| else ChildRunner.Outcome.Exited(result.exitCode, result.stdout, result.stderr) | ||
| } | ||
| } | ||
|
|
||
| override def close(): Unit = { | ||
| if (closed.compareAndSet(false, true)) { | ||
| val snapshot = workers.asScala.toVector | ||
| snapshot.foreach(worker => try worker.stdin.close() catch { case _: IOException => () }) | ||
| val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(2) | ||
| snapshot.foreach { worker => | ||
| val remaining = deadline - System.nanoTime() | ||
| if (remaining > 0) worker.process.waitFor(remaining, TimeUnit.NANOSECONDS) | ||
| } | ||
| snapshot.filter(_.process.isAlive).foreach(_.process.destroyForcibly()) | ||
| snapshot.foreach(_.process.waitFor(2, TimeUnit.SECONDS)) | ||
| workers.clear() | ||
| idle.clear() | ||
| } | ||
| } | ||
|
|
||
| private def startWorker(index: Int): Worker = { | ||
| val builder = new ProcessBuilder(workerCommand: _*) | ||
| val process = builder.start() | ||
| val worker = new Worker( | ||
| index, | ||
| process, | ||
| new BufferedOutputStream(process.getOutputStream), | ||
| new BufferedInputStream(process.getInputStream) | ||
| ) | ||
| startStderrForwarder(worker) | ||
|
|
||
| readWithin(s"token-worker-ready-$index", ReadySeconds)(readHeader(worker.stdout)) match { | ||
| case Some(Right("READY")) => worker | ||
| case notReady => | ||
| process.destroyForcibly() | ||
| throw startFailure(index, notReady) | ||
| } | ||
| } | ||
|
|
||
| private def startFailure(index: Int, ready: Option[Either[Throwable, String]]): IllegalStateException = | ||
| ready match { | ||
| case None => | ||
| new IllegalStateException(s"tokenizer worker $index did not print READY within $ReadySeconds seconds") | ||
| case Some(Right(line)) => | ||
| new IllegalStateException(s"tokenizer worker $index printed [$line] instead of READY") | ||
| case Some(Left(failure)) => | ||
| new IllegalStateException(s"tokenizer worker $index failed before READY", failure) | ||
| } | ||
|
|
||
| private def readWithin[A](threadName: String, 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() | ||
| }, | ||
| threadName | ||
| ) | ||
| reader.setDaemon(true) | ||
| reader.start() | ||
| if (done.await(seconds, TimeUnit.SECONDS)) Some(result.get()) else None | ||
| } | ||
|
|
||
| private def startStderrForwarder(worker: Worker): Unit = { | ||
| def forward(line: String): Unit = { | ||
| worker.appendDiagnostic(line) | ||
| System.err.println(s"tokenWorker[${worker.index}]: $line") | ||
| } | ||
| val thread = new Thread( | ||
| () => { | ||
| val reader = new BufferedReader(new InputStreamReader(worker.process.getErrorStream, UTF_8)) | ||
| try Iterator.continually(reader.readLine()).takeWhile(_ != null).foreach(forward) | ||
| catch { | ||
| case failure: IOException if worker.process.isAlive => | ||
| forward(s"stderr forwarding failed: ${failure.getMessage}") | ||
| case _: IOException => () | ||
| } finally try reader.close() catch { case _: IOException => () } | ||
| }, | ||
| s"token-worker-stderr-${worker.index}" | ||
| ) | ||
| thread.setDaemon(true) | ||
| thread.start() | ||
| } | ||
|
|
||
| private def writeRequest( | ||
| worker: Worker, | ||
| bytes: Array[Byte], | ||
| origSha: String, | ||
| filename: String, | ||
| fullPath: String, | ||
| timeoutSeconds: Int | ||
| ): Unit = { | ||
| val filenameBytes = filename.getBytes(UTF_8) | ||
| val pathBytes = fullPath.getBytes(UTF_8) | ||
| val header = s"REQ $origSha ${bytes.length} ${filenameBytes.length} ${pathBytes.length} $timeoutSeconds\n" | ||
| .getBytes(UTF_8) | ||
| worker.stdin.write(header) | ||
| worker.stdin.write(filenameBytes) | ||
| worker.stdin.write(pathBytes) | ||
| worker.stdin.write(bytes) | ||
| worker.stdin.flush() | ||
| } | ||
|
|
||
| private def readResponse(input: InputStream): Response = { | ||
| val fields = readHeader(input).split(" ", -1) | ||
| if (fields.length != 4 || fields(0) != "RES") | ||
| throw new IOException(s"malformed tokenizer worker response header [${fields.mkString(" ")}]") | ||
|
|
||
| val exitCode = parseInt(fields(1), "exit") | ||
| val stdoutLength = parseLength(fields(2), "outLen") | ||
| val stderrLength = parseLength(fields(3), "errLen") | ||
| val stdout = readExactly(input, stdoutLength) | ||
| val stderr = new String(readExactly(input, stderrLength), UTF_8) | ||
| new Response(exitCode, stdout, stderr) | ||
| } | ||
|
|
||
| private def readHeader(input: InputStream): String = { | ||
| val bytes = new java.io.ByteArrayOutputStream | ||
| var next = input.read() | ||
| while (next != -1 && next != '\n') { | ||
| if (bytes.size() >= 8192) throw new IOException("tokenizer worker response header exceeds 8192 bytes") | ||
| bytes.write(next) | ||
| next = input.read() | ||
| } | ||
| if (next == -1) throw new EOFException("tokenizer worker closed stdout") | ||
| new String(bytes.toByteArray, UTF_8) | ||
| } | ||
|
|
||
| private def readExactly(input: InputStream, length: Int): Array[Byte] = { | ||
| val bytes = new Array[Byte](length) | ||
| var offset = 0 | ||
| while (offset < length) { | ||
| val read = input.read(bytes, offset, length - offset) | ||
| if (read == -1) throw new EOFException(s"tokenizer worker closed stdout after $offset of $length bytes") | ||
| offset += read | ||
| } | ||
| bytes | ||
| } | ||
|
|
||
| private def parseInt(value: String, name: String): Int = | ||
| value.toIntOption.getOrElse(throw new IOException(s"invalid $name [$value] in tokenizer response")) | ||
|
|
||
| private def parseLength(value: String, name: String): Int = { | ||
| val length = parseInt(value, name) | ||
| if (length < 0) throw new IOException(s"negative $name [$value] in tokenizer response") | ||
| length | ||
| } | ||
|
|
||
| private def failedWorker(worker: Worker, failure: Throwable): ChildRunner.Outcome = { | ||
| val workerExit = if (worker.process.waitFor(100, TimeUnit.MILLISECONDS)) Some(worker.process.exitValue()) else None | ||
| val diagnostics = worker.takeDiagnostics() | ||
| replaceWorker(worker) | ||
| if (workerExit.contains(WorkerTimeoutExitCode)) | ||
| ChildRunner.Outcome.Killed(s"worker exited $WorkerTimeoutExitCode (timeout)") | ||
| else { | ||
| val stderr = if (diagnostics.nonEmpty) diagnostics else Option(failure.getMessage).getOrElse(failure.toString) | ||
| ChildRunner.Outcome.Exited(workerExit.filter(_ != 0).getOrElse(1), Array.emptyByteArray, stderr) | ||
| } | ||
| } | ||
|
|
||
| private def returnWorker(worker: Worker): Unit = { | ||
| if (worker.process.isAlive && !closed.get()) idle.put(worker) | ||
| else replaceWorker(worker) | ||
| } | ||
|
|
||
| private def replaceWorker(worker: Worker): Unit = { | ||
| workers.remove(worker) | ||
| try worker.stdin.close() catch { case _: IOException => () } | ||
| if (worker.process.isAlive) ChildRunner.killTree(worker.process) | ||
| if (!closed.get()) { | ||
| val replacement = startWorker(worker.index) | ||
| workers.add(replacement) | ||
| idle.put(replacement) | ||
| } | ||
| } | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
[Spec, medium] The pool starts in modes that never use it.
Only
Walker.executeBlobTaskcallsinvoke. OnlywalkCommitsPipelinedandwalkCommitsTreescallexecuteBlobTask. The serial walker and--shard=K/Ngo throughresolveMisses, which callsBlobExec.runwithout the pool. Butmainstarts the pool in every mode when the flag is set (line 309). Sorun_pipeline_process.sh --mode serial --tokenizer-workerstarts one worker per CPU, each with its own ctags and srcml2token, and then still spawnstokenBySha.plfor each blob.The PR body says these modes ignore the flag. They do not ignore it: they start idle processes. Reject the flag outside the pipeline modes (
--shardalready rejects--pipeline, so this check covers it). Cost: +4 lines.