Move to the official Bazel worker API

MODULE.bazel

@@ -11,6 +11,8 @@ bazel_dep(name = "rules_java", version = "7.2.0")
 bazel_dep(name = "rules_python", version = "0.23.1")
 bazel_dep(name = "rules_cc", version = "0.0.9")
 bazel_dep(name = "rules_android", version = "0.1.1")
+bazel_dep(name = "bazel_worker_java", version = "0.0.1")
+bazel_dep(name = "bazel_worker_api", version = "0.0.1")
 
 rules_kotlin_extensions = use_extension("//src/main/starlark/core/repositories:bzlmod_setup.bzl", "rules_kotlin_extensions")
 use_repo(

src/main/kotlin/io/bazel/kotlin/builder/BUILD

@@ -30,7 +30,6 @@ java_library(
         "//src/main/kotlin/io/bazel/worker",
         "//src/main/protobuf:deps_java_proto",
         "//src/main/protobuf:kotlin_model_java_proto",
-        "//src/main/protobuf:worker_protocol_java_proto",
         "//third_party:dagger",
         "@kotlin_rules_maven//:javax_annotation_javax_annotation_api",
         "@kotlin_rules_maven//:javax_inject_javax_inject",

src/main/kotlin/io/bazel/kotlin/builder/tasks/BUILD.bazel

@@ -31,7 +31,6 @@ kt_bootstrap_library(
         "//src/main/kotlin/io/bazel/worker",
         "//src/main/protobuf:deps_java_proto",
         "//src/main/protobuf:kotlin_model_java_proto",
-        "//src/main/protobuf:worker_protocol_java_proto",
         "@kotlin_rules_maven//:com_google_protobuf_protobuf_java",
         "@kotlin_rules_maven//:com_google_protobuf_protobuf_java_util",
         "@kotlin_rules_maven//:javax_inject_javax_inject",

src/main/kotlin/io/bazel/kotlin/builder/utils/BUILD.bazel

@@ -10,7 +10,6 @@ kt_bootstrap_library(
     deps = [
         "//src/main/protobuf:deps_java_proto",
         "//src/main/protobuf:kotlin_model_java_proto",
-        "//src/main/protobuf:worker_protocol_java_proto",
         "@bazel_tools//tools/java/runfiles",
         "@kotlin_rules_maven//:com_google_protobuf_protobuf_java",
         "@kotlin_rules_maven//:com_google_protobuf_protobuf_java_util",

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

@@ -1,25 +1,13 @@
-load("@rules_java//java:defs.bzl", "java_binary", "java_import")
-
 # General purpose Bazel worker implemented Kotlin. Best suited for jvm based tools.
 load("//src/main/kotlin:bootstrap.bzl", "kt_bootstrap_library")
 
-java_binary(
-    name = "worker_proto_bundle_bin",
-    runtime_deps = ["//src/main/protobuf:worker_protocol_java_proto"],
-)
-
-java_import(
-    name = "worker_proto",
-    jars = [":worker_proto_bundle_bin_deploy.jar"],
-)
-
 kt_bootstrap_library(
     name = "worker",
     srcs = glob(["*.kt"]),
     visibility = [
         "//:__subpackages__",
     ],
     deps = [
-        ":worker_proto",
+        "@bazel_worker_java//src/main/java/com/google/devtools/build/lib/worker:work_request_handlers",
     ],
 )

src/main/kotlin/io/bazel/worker/PersistentWorker.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/protobuf/BUILD

@@ -8,11 +8,6 @@ alias(
     actual = "@bazel_tools//src/main/protobuf:deps_java_proto",
 )
 
-alias(
-    name = "worker_protocol_java_proto",
-    actual = "@bazel_tools//src/main/protobuf:worker_protocol_java_proto",
-)
-
 proto_library(
     name = "kotlin_model_proto",
     srcs = [":kotlin_model.proto"],
@@ -29,6 +24,5 @@ java_library(
     exports = [
         ":deps_java_proto",
         ":kotlin_model_java_proto",
-        ":worker_protocol_java_proto",
     ],
 )

src/main/starlark/core/repositories/download.bzl

@@ -59,6 +59,22 @@ def kt_download_local_dev_dependencies():
         url = "https://github.com/bazelbuild/rules_jvm_external/releases/download/%s/rules_jvm_external-%s.tar.gz" % (versions.RULES_JVM_EXTERNAL_TAG, versions.RULES_JVM_EXTERNAL_TAG),
     )
 
+    maybe(
+        http_archive,
+        name = "bazel_worker_java",
+        sha256 = versions.BAZEL_WORKER_JAVA.sha256,
+        strip_prefix = "bazel-worker-api-{version}/java/".format(version = versions.BAZEL_WORKER_JAVA.version),
+        urls = [url.format(version = versions.BAZEL_WORKER_JAVA.version) for url in versions.BAZEL_WORKER_JAVA.url_templates],
+    )
+
+    maybe(
+        http_archive,
+        name = "bazel_worker_api",
+        sha256 = versions.BAZEL_WORKER_JAVA.sha256,
+        strip_prefix = "bazel-worker-api-{version}/proto/".format(version = versions.BAZEL_WORKER_JAVA.version),
+        urls = [url.format(version = versions.BAZEL_WORKER_JAVA.version) for url in versions.BAZEL_WORKER_JAVA.url_templates],
+    )
+
     versions.use_repository(
         name = "rules_license",
         rule = http_archive,

src/main/starlark/core/repositories/versions.bzl

@@ -153,5 +153,12 @@ versions = struct(
         ],
         sha256 = "d3234179bcff1886d53d67c11eca47f7f3cf7b63c349d16965f6db51b7f3dd9a",
     ),
+    BAZEL_WORKER_JAVA = version(
+        version = "0.0.1",
+        url_templates = [
+            "https://github.com/bazelbuild/bazel-worker-api/releases/download/v{version}/bazel-worker-api-v{version}.tar.gz",
+        ],
+        sha256 = "b341e3fba0a3dd0ab7bfdc7e256fad711a1f9e9255563a74c305676046b5a184",
+    ),
     use_repository = _use_repository,
 )