Move to the official Bazel worker API

Author: Bencodes

MODULE.bazel

@@ -15,6 +15,8 @@ bazel_dep(name = "rules_java", version = "8.9.0")
 bazel_dep(name = "rules_python", version = "0.23.1")
 bazel_dep(name = "rules_android", version = "0.6.1")
 bazel_dep(name = "bazel_features", version = "1.25.0")
+bazel_dep(name = "bazel_worker_api", version = "0.0.4")
+bazel_dep(name = "bazel_worker_java", version = "0.0.4")
 
 rules_java_toolchains = use_extension("@rules_java//java:extensions.bzl", "toolchains")
 use_repo(rules_java_toolchains, "remote_java_tools")

src/main/kotlin/io/bazel/worker/BUILD.bazel

@@ -21,5 +21,6 @@ kt_bootstrap_library(
     ],
     deps = [
         ":worker_proto",
+        "@bazel_worker_java//src/main/java/com/google/devtools/build/lib/worker:work_request_handlers",
     ],
 )

src/main/kotlin/io/bazel/worker/CpuTimeBasedGcScheduler.kt

@@ -17,111 +17,51 @@
 
 package io.bazel.worker
 
+import com.google.devtools.build.lib.worker.ProtoWorkerMessageProcessor
+import com.google.devtools.build.lib.worker.WorkRequestHandler.WorkRequestCallback
+import com.google.devtools.build.lib.worker.WorkRequestHandler.WorkRequestHandlerBuilder
 import com.google.devtools.build.lib.worker.WorkerProtocol.WorkRequest
-import com.google.devtools.build.lib.worker.WorkerProtocol.WorkResponse
-import src.main.kotlin.io.bazel.worker.GcScheduler
-import java.io.InputStream
+import java.io.IOException
+import java.io.PrintWriter
 import java.time.Duration
-import java.util.concurrent.ExecutorCompletionService
-import java.util.concurrent.ExecutorService
-import java.util.concurrent.Executors
-import java.util.concurrent.TimeUnit
-import java.util.concurrent.atomic.AtomicLong
 
 /**
  * PersistentWorker satisfies Bazel persistent worker protocol for executing work.
  *
  * Supports multiplex (https://docs.bazel.build/versions/master/multiplex-worker.html) provided
  * the work is thread/coroutine safe.
- *
- * @param executor thread pool for executing tasks.
- * @param captureIO to avoid writing stdout and stderr while executing.
- * @param cpuTimeBasedGcScheduler to trigger gc cleanup.
  */
-class PersistentWorker(
-  private val captureIO: () -> IO,
-  private val executor: ExecutorService,
-  private val cpuTimeBasedGcScheduler: GcScheduler,
-) : Worker {
-  constructor(
-    executor: ExecutorService,
-    captureIO: () -> IO,
-  ) : this(
-    captureIO,
-    executor,
-    GcScheduler {},
-  )
-
-  constructor() : this(
-    IO.Companion::capture,
-    Executors.newCachedThreadPool(),
-    CpuTimeBasedGcScheduler(Duration.ofSeconds(10)),
-  )
-
-  override fun start(execute: Work) =
-    WorkerContext.run {
-      captureIO().use { io ->
-        val running = AtomicLong(0)
-        val completion = ExecutorCompletionService<WorkResponse>(executor)
-        val producer =
-          executor.submit {
-            io.input.readRequestAnd { request ->
-              running.incrementAndGet()
-              completion.submit {
-                doTask(
-                  name = "request ${request.requestId}",
-                  task = request.workTo(execute),
-                ).asResponseTo(request.requestId, io)
-              }
-            }
-          }
-        val consumer =
-          executor.submit {
-            while (!producer.isDone || running.get() > 0) {
-              // poll time is how long before checking producer liveliness. Too long, worker hangs
-              // when being shutdown -- too short, and it starves the process.
-              completion.poll(1, TimeUnit.SECONDS)?.run {
-                running.decrementAndGet()
-                get().writeDelimitedTo(io.output)
-                io.output.flush()
-              }
-              cpuTimeBasedGcScheduler.maybePerformGc()
-            }
-          }
-        producer.get()
-        consumer.get()
-        io.output.close()
+class PersistentWorker : Worker {
+  override fun start(execute: Work): Int {
+    return WorkerContext.run {
+      val realStdErr = System.err
+      try {
+        val workerHandler =
+          WorkRequestHandlerBuilder(
+            WorkRequestCallback { request: WorkRequest, pw: PrintWriter ->
+              return@WorkRequestCallback doTask(
+                name = "request ${request.requestId}",
+                task = request.workTo(execute),
+              ).asResponse(pw)
+            },
+            realStdErr,
+            ProtoWorkerMessageProcessor(System.`in`, System.out),
+          ).setCpuUsageBeforeGc(Duration.ofSeconds(10)).build()
+        workerHandler.processRequests()
+      } catch (e: IOException) {
+        this.error(e, { "Unknown IO exception" })
+        e.printStackTrace(realStdErr)
+        return@run 1
       }
       return@run 0
     }
+  }
 
   private fun WorkRequest.workTo(execute: Work): (sub: WorkerContext.TaskContext) -> Status =
     { ctx -> execute(ctx, argumentsList.toList()) }
 
-  private fun InputStream.readRequestAnd(action: (WorkRequest) -> Unit) {
-    while (true) {
-      WorkRequest
-        .parseDelimitedFrom(this)
-        ?.run(action)
-        ?: return
-    }
+  private fun TaskResult.asResponse(pw: PrintWriter): Int {
+    pw.print(log.out.toString())
+    return status.exit
   }
-
-  private fun TaskResult.asResponseTo(
-    id: Int,
-    io: IO,
-  ): WorkResponse =
-    WorkResponse
-      .newBuilder()
-      .apply {
-        val cap = io.readCapturedAsUtf8String()
-        // append whatever falls through standard out.
-        output =
-          listOf(
-            log.out.toString(),
-            cap,
-          ).joinToString("\n").trim()
-        exitCode = status.exit
-        requestId = id
-      }.build()
 }

src/main/kotlin/io/bazel/worker/GcScheduler.kt

@@ -1,17 +1,5 @@
 load("//kotlin:jvm.bzl", "kt_jvm_library", "kt_jvm_test")
 
-kt_jvm_test(
-    name = "IOTest",
-    srcs = [
-        "IOTest.kt",
-    ],
-    test_class = "io.bazel.worker.IOTest",
-    deps = [
-        "//src/main/kotlin/io/bazel/worker",
-        "@kotlin_rules_maven//:com_google_truth_truth",
-    ],
-)
-
 kt_jvm_test(
     name = "WorkerContextTest",
     srcs = [
@@ -58,29 +46,11 @@ kt_jvm_test(
     ],
 )
 
-kt_jvm_test(
-    name = "JavaPersistentWorkerTest",
-    srcs = [
-        "JavaPersistentWorkerTest.kt",
-    ],
-    test_class = "io.bazel.worker.JavaPersistentWorkerTest",
-    deps = [
-        ":WorkerEnvironment",
-        "//src/main/kotlin/io/bazel/worker",
-        "//src/main/protobuf:worker_protocol_java_proto",
-        "@kotlin_rules_maven//:com_google_truth_truth",
-    ],
-)
-
 test_suite(
     name = "worker_tests",
     tests = [
-        ":IOTest",
         ":InvocationWorkerTest",
-        ":JavaPersistentWorkerTest",
         ":WorkerContextTest",
-        # TODO(restingbull): Re-enable when not flaky.
-        #":WorkerEnvironmentTest",
     ],
 )