Fix concurrency for multiplex

examples/multiplex/.bazelrc

@@ -2,3 +2,4 @@ build --worker_verbose
 build --worker_max_instances=1
 build --worker_quit_after_build
 build --profile=/tmp/profile.gz
+build --experimental_worker_max_multiplex_instances=KotlinCompile=5

examples/multiplex/WORKSPACE

@@ -2,9 +2,13 @@ workspace(name = "multiplex")
 
 local_repository(
     name = "io_bazel_rules_kotlin",
-    path = "../../bazel-bin/rules_kotlin_release.tgz",
+    path = "../..",
 )
 
+load("@io_bazel_rules_kotlin//kotlin:dependencies.bzl", "kt_download_local_dev_dependencies")
+
+kt_download_local_dev_dependencies()
+
 load("@io_bazel_rules_kotlin//kotlin:kotlin.bzl", "kotlin_repositories", "kt_register_toolchains")
 
 kotlin_repositories()

kotlin/internal/jvm/compile.bzl

@@ -204,7 +204,7 @@ def _kotlinc_options_provider_to_flags(opts, language_version):
     if opts.x_no_optimized_callable_references:
         flags.append("-Xno-optimized-callable-references")
     if opts.x_multi_platform:
-            flags.append("-Xmulti-platform")
+        flags.append("-Xmulti-platform")
     if opts.java_parameters:
         flags.append("-java-parameters")
     return flags
@@ -335,10 +335,7 @@ def _run_merge_jdeps_action(ctx, rule_kind, toolchains, jdeps, outputs):
         input_manifests = input_manifests,
         outputs = [f for f in outputs.values()],
         executable = toolchains.kt.jdeps_merger.files_to_run.executable,
-        execution_requirements = {
-            "supports-workers": "1",
-            "supports-multiplex-workers": "1",
-        },
+        execution_requirements = toolchains.kt.execution_requirements,
         arguments = [args],
         progress_message = progress_message,
     )
@@ -484,10 +481,7 @@ def _run_kt_builder_action(
         input_manifests = input_manifests,
         outputs = [f for f in outputs.values()],
         executable = toolchains.kt.kotlinbuilder.files_to_run.executable,
-        execution_requirements = {
-            "supports-workers": "1",
-            "supports-multiplex-workers": "1",
-        },
+        execution_requirements = toolchains.kt.execution_requirements,
         arguments = [args],
         progress_message = progress_message,
         env = {

kotlin/internal/jvm/impl.bzl

@@ -137,7 +137,7 @@ def kt_jvm_import_impl(ctx):
                 files = depset(direct = [artifact.class_jar]),
                 runfiles = ctx.runfiles(
                     # Append class jar with the optional sources jar
-                    files = [artifact.class_jar] + [artifact.source_jar] if artifact.source_jar else []
+                    files = [artifact.class_jar] + [artifact.source_jar] if artifact.source_jar else [],
                 ),
             ),
             JavaInfo(

kotlin/internal/toolchains.bzl

@@ -79,6 +79,10 @@ def _kotlin_toolchain_impl(ctx):
         kotlin_home = ctx.attr.kotlin_home,
         jvm_stdlibs = java_common.merge(compile_time_providers + runtime_providers),
         js_stdlibs = ctx.attr.js_stdlibs,
+        execution_requirements = {
+            "supports-workers": "1",
+            "supports-multiplex-workers": "1" if ctx.attr.experimental_multiplex_workers else "0",
+        },
         experimental_use_abi_jars = ctx.attr.experimental_use_abi_jars,
         experimental_strict_kotlin_deps = ctx.attr.experimental_strict_kotlin_deps,
         experimental_report_unused_deps = ctx.attr.experimental_report_unused_deps,
@@ -187,6 +191,10 @@ _kt_toolchain = rule(
             ],
             providers = [_KtJsInfo],
         ),
+        "experimental_multiplex_workers": attr.bool(
+            doc = """Run workers in multiplex mode.""",
+            default = False,
+        ),
         "experimental_use_abi_jars": attr.bool(
             doc = """Compile using abi jars. Can be disabled for an individual target using the tag
             `kt_abi_plugin_incompatible`""",
@@ -258,6 +266,7 @@ def define_kt_toolchain(
         experimental_strict_kotlin_deps = None,
         experimental_report_unused_deps = None,
         experimental_reduce_classpath_mode = None,
+        experimental_multiplex_workers = None,
         javac_options = None,
         kotlinc_options = None):
     """Define the Kotlin toolchain."""
@@ -282,6 +291,7 @@ def define_kt_toolchain(
             absolute_target("//kotlin/internal:noexperimental_use_abi_jars"): False,
             "//conditions:default": experimental_use_abi_jars,
         }),
+        experimental_multiplex_workers = experimental_multiplex_workers,
         experimental_strict_kotlin_deps = experimental_strict_kotlin_deps,
         experimental_report_unused_deps = experimental_report_unused_deps,
         experimental_reduce_classpath_mode = experimental_reduce_classpath_mode,
@@ -329,6 +339,11 @@ def kt_configure_toolchains():
         values = {"define": "kt_timings=1"},
     )
 
+    native.config_setting(
+        name = "experimental_multiplex_workers",
+        values = {"define": "kt_multiplex=1"},
+    )
+
     native.config_setting(
         name = "builder_debug_trace",
         values = {"define": "kt_trace=1"},

scripts/release.sh

@@ -56,7 +56,7 @@ shasum -a 256 bazel-bin/rules_kotlin_release.tgz > bazel-bin/rules_kotlin_releas
 
 # iterate through the examples and build them
 for ex in examples/*/; do
-  if [[ -f "$ex/WORKSPACE" ]]; then
+  if [[ -f "$ex/WORKSPACE" ]] && ! [[ -f "$ex/ignore.me" ]]; then
     (
       cd "$ex"
       bazel build ${BUILD_ARGS} --override_repository=io_bazel_rules_kotlin=$ARCHIVE_DIR //...:all

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

@@ -19,13 +19,15 @@ package io.bazel.worker
 
 import com.google.devtools.build.lib.worker.WorkerProtocol
 import com.google.devtools.build.lib.worker.WorkerProtocol.WorkRequest
+import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExecutorCoroutineDispatcher
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.asCoroutineDispatcher
+import kotlinx.coroutines.async
 import kotlinx.coroutines.flow.asFlow
+import kotlinx.coroutines.flow.buffer
 import kotlinx.coroutines.flow.collect
-import kotlinx.coroutines.flow.flowOn
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.runBlocking
 import java.nio.charset.StandardCharsets.UTF_8
@@ -58,15 +60,16 @@ class PersistentWorker(
    */
   private class BlockableDispatcher(
     private val unblockedContext: CoroutineContext,
-    private val blockingContext: ExecutorCoroutineDispatcher
-  ) {
+    private val blockingContext: ExecutorCoroutineDispatcher,
+    scope: CoroutineScope
+  ) : CoroutineScope by scope {
     companion object {
       fun <T> runIn(
         owningContext: CoroutineContext,
         exec: suspend BlockableDispatcher.() -> T
       ) =
         Executors.newCachedThreadPool().asCoroutineDispatcher().use { dispatcher ->
-          runBlocking(owningContext) { BlockableDispatcher(owningContext, dispatcher).exec() }
+          runBlocking(owningContext) { BlockableDispatcher(owningContext, dispatcher, this).exec() }
         }
     }
 
@@ -84,26 +87,31 @@ class PersistentWorker(
         blockable {
           generateSequence { WorkRequest.parseDelimitedFrom(io.input) }
         }.asFlow()
-          .flowOn(coroutineContext)
           .map { request ->
             info { "received req: ${request.requestId}" }
-            doTask("request ${request.requestId}") { ctx ->
-              request.argumentsList.run {
-                execute(ctx, toList())
+            async {
+              doTask("request ${request.requestId}") { ctx ->
+                request.argumentsList.run {
+                  execute(ctx, toList())
+                }
+              }.let { result ->
+                info { "task result ${result.status}" }
+                WorkerProtocol.WorkResponse.newBuilder().apply {
+                  output =
+                    listOf(
+                      result.log.out.toString(),
+                      io.captured.toByteArray().toString(UTF_8)
+                    ).filter { it.isNotBlank() }.joinToString("\n")
+                  exitCode = result.status.exit
+                  requestId = request.requestId
+                }.build()
               }
-            }.let { result ->
-              [email protected] { "task result ${result.status}" }
-              WorkerProtocol.WorkResponse.newBuilder().apply {
-                output =
-                  listOf(
-                    result.log.out.toString(),
-                    io.captured.toByteArray().toString(UTF_8)
-                  ).filter { it.isNotBlank() }.joinToString("\n")
-                exitCode = result.status.exit
-                requestId = request.requestId
-              }.build()
             }
           }
+          .buffer()
+          .map { deferred ->
+            deferred.await()
+          }
           .collect { response ->
             blockable {
               info {