DO NOT MERGE: Refactor Persistent Workers To Use Kotlin Worker Avoid OOM

MODULE.bazel

@@ -7,6 +7,11 @@ module(
 bazel_dep(name = "rules_java", version = "8.6.3")
 bazel_dep(name = "rules_jvm_external", version = "5.3")
 bazel_dep(name = "rules_kotlin", version = "2.1.0")
+single_version_override(
+    module_name = "rules_kotlin",
+    patch_strip = 1,
+    patches = ["//third_party:rules_kotlin_worker_visibilty.patch"],
+)
 bazel_dep(name = "bazel_skylib", version = "1.7.1")
 bazel_dep(name = "platforms", version = "0.0.8")
 

rules/impl.bzl

@@ -15,6 +15,9 @@ load(
     _ANDROID_LINT_TOOLCHAIN_TYPE = "ANDROID_LINT_TOOLCHAIN_TYPE",
     _utils = "utils",
 )
+# Estimation is done by VisualVM runs + assumption of GC kicking in
+def _resource_set_callback(os_name, num_inputs):
+    return {"cpu": 5, "memory": 10000}
 
 def _run_android_lint(
         ctx,
@@ -135,12 +138,13 @@ def _run_android_lint(
         execution_requirements = {
             "supports-workers": "1",
             "supports-multiplex-workers": "1",
-            "requires-worker-protocol": "json",
+            "requires-worker-protocol": "proto",
         },
         env = {
             # https://googlesamples.github.io/android-custom-lint-rules/usage/variables.md.html
             "ANDROID_LINT_SKIP_BYTECODE_VERIFIER": ("true" if android_lint_skip_bytecode_verifier else "false"),
         },
+        resource_set = _resource_set_callback,
     )
 
 def _get_module_name(ctx):

src/cli/AndroidLintAction.kt

@@ -1,18 +1,52 @@
 package com.rules.android.lint.cli
 
 import com.rules.android.lint.worker.Worker
+import io.bazel.worker.PersistentWorker
+import io.bazel.worker.Status
+import io.bazel.worker.Work
+import io.bazel.worker.WorkerContext
 import java.io.PrintStream
 import java.nio.file.Files
+import javax.inject.Inject
 import kotlin.system.exitProcess
 
 object AndroidLintAction {
   @JvmStatic
   fun main(args: Array<String>) {
+    if ("--persistent_worker" in args) {
+      val worker = PersistentWorker()
+      worker.start(AndroidLintPersistentWorker()).run(::exitProcess)
+      return
+    }
     val worker = Worker.fromArgs(args, AndroidLintExecutor())
     val exitCode = worker.processRequests()
     exitProcess(exitCode)
   }
 
+  private class AndroidLintPersistentWorker
+    @Inject
+    constructor() : Work {
+      override fun invoke(
+        ctx: WorkerContext.TaskContext,
+        args: Iterable<String>,
+      ): Status {
+        val workingDirectory = Files.createTempDirectory("rules")
+        try {
+          val parsedArgs = AndroidLintActionArgs.parseArgs(args.toList())
+          val result = AndroidLintRunner().runAndroidLint(parsedArgs, workingDirectory)
+          return if (result != 0) Status.ERROR else Status.SUCCESS
+        } catch (exception: Exception) {
+          return Status.ERROR
+        } finally {
+          try {
+            workingDirectory.toFile().deleteRecursively()
+          } catch (e: Exception) {
+            e.printStackTrace()
+          }
+        }
+      }
+    }
+
   private class AndroidLintExecutor : Worker.WorkRequestCallback {
     override fun processWorkRequest(
       args: List<String>,

src/cli/AndroidLintPersistentWorker.kt

@@ -0,0 +1,77 @@
+package com.rules.android.lint.cli
+
+import com.rules.android.lint.worker.Worker
+import io.bazel.worker.PersistentWorker
+import io.bazel.worker.Status
+import io.bazel.worker.Work
+import io.bazel.worker.WorkerContext
+import java.io.PrintStream
+import java.nio.file.Files
+import javax.inject.Inject
+import kotlin.system.exitProcess
+
+object AndroidLintPersistentWorker {
+  @JvmStatic
+  fun main(args: Array<String>) {
+    if ("--persistent_worker" in args) {
+      val worker = PersistentWorker()
+      worker.start(AndroidLint()).run(::exitProcess)
+      return
+    }
+    val worker = Worker.fromArgs(args, AndroidLintExecutor())
+    val exitCode = worker.processRequests()
+    exitProcess(exitCode)
+  }
+
+  private class AndroidLint
+    @Inject
+    constructor() : Work {
+      override fun invoke(
+        ctx: WorkerContext.TaskContext,
+        args: Iterable<String>,
+      ): Status {
+        val workingDirectory = Files.createTempDirectory("rules")
+        try {
+          val runner = AndroidLintRunner()
+          val parsedArgs = AndroidLintActionArgs.parseArgs(args.toList())
+          val result = runner.runAndroidLint(parsedArgs, workingDirectory)
+          if (result == 0) {
+            return Status.SUCCESS
+          }
+          return Status.ERROR
+        } catch (exception: Exception) {
+          return Status.ERROR
+        } finally {
+          try {
+            workingDirectory.toFile().deleteRecursively()
+          } catch (e: Exception) {
+            e.printStackTrace()
+          }
+        }
+      }
+    }
+
+  private class AndroidLintExecutor : Worker.WorkRequestCallback {
+    override fun processWorkRequest(
+      args: List<String>,
+      printStream: PrintStream,
+    ): Int {
+      val workingDirectory = Files.createTempDirectory("rules")
+
+      try {
+        val runner = AndroidLintRunner()
+        val parsedArgs = AndroidLintActionArgs.parseArgs(args)
+        return runner.runAndroidLint(parsedArgs, workingDirectory)
+      } catch (exception: Exception) {
+        exception.printStackTrace()
+        return 1
+      } finally {
+        try {
+          workingDirectory.toFile().deleteRecursively()
+        } catch (e: Exception) {
+          e.printStackTrace()
+        }
+      }
+    }
+  }
+}

src/worker/BUILD

@@ -12,6 +12,7 @@ kt_jvm_library(
         "@rules_android_lint_deps//:com_squareup_okio_okio_jvm",
         "@rules_android_lint_deps//:io_reactivex_rxjava3_rxjava",
         "@rules_android_lint_deps//:org_reactivestreams_reactive_streams",
+        "@rules_kotlin//src/main/kotlin:worker",
     ],
 )
 

src/worker/Worker.kt

@@ -27,17 +27,10 @@ interface Worker {
   companion object {
     /**
      * Creates the appropriate worker instance using the provided worker arguments.
-     *
-     * If `--persistent_worker` exists in the arguments, an instance of PersistentWorker will
-     * be returned. Otherwise an instance of InvocationWorker will be returned.
      */
     fun fromArgs(
       args: Array<String>,
       workerMessageProcessor: WorkRequestCallback,
-    ): Worker =
-      when {
-        "--persistent_worker" in args -> PersistentWorker(workerMessageProcessor)
-        else -> InvocationWorker(args, workerMessageProcessor)
-      }
+    ): Worker = InvocationWorker(args, workerMessageProcessor)
   }
 }

third_party/rules_kotlin_worker_visibilty.patch

@@ -0,0 +1,12 @@
+diff --git a/src/main/kotlin/BUILD.bazel b/src/main/kotlin/BUILD.bazel
+index 5f04f0e..6668b05 100755
+--- a/src/main/kotlin/BUILD.bazel
++++ b/src/main/kotlin/BUILD.bazel
+@@ -17,6 +17,7 @@ load("@rules_java//java:defs.bzl", "java_binary", "java_import")
+ java_import(
+     name = "worker",
+     jars = ["kotlin_worker.jar"],
++    visibility = ["//visibility:public"],
+ )
+
+ java_import(