fix(backend-cpu): CoroutineSchedule deadlock that hangs the jvm CI leg

.github/workflows/build.yml

@@ -40,14 +40,16 @@ jobs:
               jvmTest
               :skainet-apps:skainet-plan:test
               :skainet-backends:benchmarks:jvm-cpu-publish:test
-            # This leg runs nearly every module's jvmTest sequentially, each forking its own
-            # 8g-heap Test JVM (build.gradle.kts `maxHeapSize`). ci-gradle.properties already caps
-            # org.gradle.workers.max=2 repo-wide, but on this leg two 8g forks running at once plus
-            # the 4g Gradle client JVM (GRADLE_OPTS below) can still exceed the 16 GB runner — the
-            # same OOM signature described above ("The operation was canceled" with no BUILD
-            # FAILED), just needing the extra memory pressure SKEEP-005's concurrency tests
-            # (KernelDispatchConcurrencyTest, CoroutineSchedule, parallel SDPA) added to tip it
-            # over. Force this leg fully serial so at most one 8g fork runs at a time.
+            # This leg runs nearly every module's jvmTest, each forking its own 8g-heap Test JVM
+            # (build.gradle.kts `maxHeapSize`). ci-gradle.properties caps org.gradle.workers.max=2
+            # repo-wide; this leg goes fully serial so at most one 8g fork runs next to the 4g
+            # Gradle client JVM (GRADLE_OPTS below) on the 16 GB runner.
+            #
+            # History: the 40-minute reds after SKEEP-005 (#1262) were NOT memory — the runner had
+            # 11 GB free when the job was killed. `CoroutineSchedule.forRange` blocked a
+            # Dispatchers.Default worker on children that needed the same 4-thread pool, so
+            # `skainet-backend-cpu:jvmTest` deadlocked in CoroutineScheduleTest on the 4-vCPU
+            # runner. Fixed in the schedule itself; the serial setting stays for memory headroom.
             extraGradleArgs: -Dorg.gradle.workers.max=1
           # verifyNpmPins guards the npm-* pins in gradle/libs.versions.toml against
           # lockfile drift; it belongs on the leg that already has the JS toolchain.

docs/modules/ROOT/pages/explanation/schedules.adoc

@@ -50,16 +50,21 @@ public interface Schedule {
 
 It is deliberately small, non-suspending and dependency-free: ops are synchronous on every
 Kotlin target, and `skainet-lang-core` takes no coroutine dependency. The JVM implementation,
-`CoroutineSchedule` in `skainet-backend-cpu`, is structured concurrency in the textbook sense —
-one region is one `coroutineScope`:
-
-* the calling thread runs the first chunk itself, so no core idles blocked on the join;
-* the other chunks are `launch`ed on the dispatcher (`Dispatchers.Default` by default);
-* the scope closes only when every child has finished or been cancelled — the first failure
-  cancels the siblings and is rethrown to the caller;
-* a region reached from inside another region runs inline, so the dispatcher never waits on
-  itself; `CoroutineSchedule.dedicated(n)` owns its own pool for callers that already live on
-  `Dispatchers.Default`.
+`CoroutineSchedule` in `skainet-backend-cpu`, keeps the structured-concurrency guarantees without
+ever letting the caller block on work that still needs a pool thread. One region is one shared
+queue of chunks:
+
+* the calling thread runs the first chunk itself, then keeps claiming chunks until the queue is
+  empty, so no core idles blocked on the join;
+* `tasks - 1` helpers are dispatched to the dispatcher (`Dispatchers.Default` by default) and
+  claim chunks from the same queue; a helper the pool never started finds the queue empty and
+  exits without running a body;
+* the region returns only when every *claimed* chunk has finished — the first failure stops the
+  chunks that have not started and is rethrown to the caller once the running ones are done;
+* because the caller only waits for chunks that are already running somewhere, a region entered
+  from a thread of the dispatcher's own pool cannot deadlock, however many pool threads are busy;
+* a region reached from inside another region runs inline; `CoroutineSchedule.dedicated(n)` owns
+  its own pool for callers that want isolation from `Dispatchers.Default`.
 
 == The contract a body must keep
 

docs/modules/skeep/pages/005-schedules-structured-concurrency.adoc

@@ -93,13 +93,22 @@ cancelled; all writes happen-before the return.
 
 === The JVM implementation (`skainet-backend-cpu`)
 
-`CoroutineSchedule(dispatcher = Dispatchers.Default, parallelism = cores, sink)`: a region is
-`runBlocking { coroutineScope { launch(dispatcher) { … } … ; body(first chunk) } }` — the caller
-runs the first chunk itself, so no core idles on the join; a nested region (detected with a
-thread-local) runs inline so the dispatcher never waits on itself; `dedicated(parallelism)` owns
-its own pool for callers that already live on `Dispatchers.Default`. `parallelChunks(outputDim,
-schedule)` replaces the old island; the Panama Q4_K/Q5_K SPI kernels gained a schedule-aware
-overload; `ScheduleRegion` events report what ran when a sink is attached.
+`CoroutineSchedule(dispatcher = Dispatchers.Default, parallelism = cores, sink)`: a region is a
+shared queue of `tasks` chunks — `tasks - 1` helpers are dispatched to the dispatcher, the caller
+runs the first chunk itself and then drains the queue, and finally waits only for chunks some
+thread has already claimed. A caller can therefore always finish a region on its own, which is
+what makes it safe to enter from `Dispatchers.Default`. A nested region (detected with a
+thread-local) runs inline; `dedicated(parallelism)` owns its own pool for callers that want
+isolation from `Dispatchers.Default`. `parallelChunks(outputDim, schedule)` replaces the old
+island; the Panama Q4_K/Q5_K SPI kernels gained a schedule-aware overload; `ScheduleRegion`
+events report what ran when a sink is attached.
+
+NOTE: The first merged version ran a region as
+`runBlocking { coroutineScope { launch(dispatcher) { … } … } }`. A `coroutineScope` waits for
+every child, including the ones the pool never got to start, so once every `Dispatchers.Default`
+worker was inside a region the pool had nobody left to run the children and the region hung.
+That is exactly the 4-vCPU CI runner's situation; it never reproduced on a 14-core laptop. The
+chunk-queue region above replaced it.
 
 === First consumer: `scaledDotProductAttention`
 
@@ -169,8 +178,10 @@ bodies. Behaviour change on the JVM only: `DirectCpuExecutionContext()` now runs
 
 == Risks
 
-* **Nested `runBlocking` on `Dispatchers.Default`** — mitigated by caller participation, inline
-  nested regions and `dedicated()`; tested from inside a Default worker.
+* **Blocking a `Dispatchers.Default` worker on the pool it belongs to** — the region never
+  waits for a chunk that still needs a pool thread: the caller drains the queue itself and only
+  joins chunks already running. Tested from inside a Default worker and, deterministically, from
+  every thread of a saturated two-thread pool.
 * **A body that touches the context** — the contract is documented on `forRange`; the parity
   tests run bodies on foreign threads.
 * **JIT-dependent lane reductions** — the Panama matmul's `reduceLanes` order changes when the

skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/schedule/CoroutineSchedule.kt

@@ -3,30 +3,40 @@ package sk.ainet.exec.schedule
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.asCoroutineDispatcher
-import kotlinx.coroutines.coroutineScope
-import kotlinx.coroutines.launch
-import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.asExecutor
 import sk.ainet.context.schedule.Schedule
 import sk.ainet.lang.memory.ExperimentalMemoryApi
 import sk.ainet.lang.memory.trace.NoopTraceSink
 import sk.ainet.lang.memory.trace.TraceClock
 import sk.ainet.lang.memory.trace.TraceEvent
 import sk.ainet.lang.memory.trace.TraceSink
+import java.util.concurrent.CountDownLatch
+import java.util.concurrent.Executor
 import java.util.concurrent.ExecutorService
 import java.util.concurrent.Executors
 import java.util.concurrent.atomic.AtomicInteger
+import java.util.concurrent.atomic.AtomicReference
 
 /**
- * The JVM [Schedule] (SKEEP-005): a region is one `coroutineScope` — structured concurrency —
- * whose children run on [dispatcher] while the *calling thread runs the first chunk itself*, so
- * no core sits idle blocked on the join. The scope closes only after every child finished or was
- * cancelled, which gives the contract its guarantees: the first failure cancels the siblings and
- * is rethrown, no task outlives [forRange], and every task's writes happen-before the return.
+ * The JVM [Schedule] (SKEEP-005): a region is a queue of `tasks` chunks shared between the
+ * calling thread and `tasks - 1` helpers dispatched to [dispatcher]. The caller runs the first
+ * chunk itself, then keeps claiming chunks until the queue is empty, and only then waits — and
+ * only for chunks a thread has already claimed. That is what makes the region safe to enter
+ * from a thread of the dispatcher's own pool: a caller never blocks on work that still needs a
+ * pool thread, because it does that work itself. A helper that the pool never got round to
+ * running finds the queue empty and exits without touching the body.
+ *
+ * The contract's guarantees hold as before: the first failure stops every chunk that has not
+ * started, is rethrown once every running chunk has finished (later failures are attached as
+ * suppressed), no body outlives [forRange], and every chunk's writes happen-before the return.
  *
  * A region reached from inside another region (a body that calls into a parallel op despite the
- * contract) runs inline: the dispatcher never waits on itself, and a nested `runBlocking` on
- * `Dispatchers.Default` cannot starve the pool. Callers that already run on `Dispatchers.Default`
- * and want true isolation use [dedicated].
+ * contract) runs inline. Callers that want isolation from `Dispatchers.Default` use [dedicated].
+ *
+ * A previous version ran the region as `runBlocking { coroutineScope { launch(dispatcher) … } }`.
+ * A `coroutineScope` waits for every child, including the ones the pool never started, so a
+ * region entered from `Dispatchers.Default` deadlocked as soon as every worker was inside one —
+ * routinely on a 4-vCPU CI runner, never on a 14-core laptop.
  */
 @OptIn(ExperimentalMemoryApi::class)
 public open class CoroutineSchedule @JvmOverloads constructor(
@@ -42,6 +52,8 @@ public open class CoroutineSchedule @JvmOverloads constructor(
 
     final override val name: String = "$label($parallelism)"
 
+    private val executor: Executor = dispatcher.asExecutor()
+
     override fun forRange(n: Int, grain: Int, body: (start: Int, end: Int) -> Unit) {
         val tasks = Schedule.tasksFor(n, grain, parallelism)
         if (tasks == 0) return
@@ -51,36 +63,86 @@ public open class CoroutineSchedule @JvmOverloads constructor(
         }
         val chunk = Schedule.chunkFor(n, tasks)
         val started = if (sink.isEnabled) TraceClock.nowNanos() else 0L
+        val region = Region(n, chunk, tasks, body)
         inRegion.set(true)
         try {
-            runBlocking {
-                coroutineScope {
-                    var start = chunk
-                    while (start < n) {
-                        val s = start
-                        val e = minOf(start + chunk, n)
-                        launch(dispatcher) { inRegion(s, e, body) }
-                        start = e
-                    }
-                    body(0, minOf(chunk, n))
-                }
-            }
+            repeat(tasks - 1) { executor.execute(region) }
+            region.runChunk(0)
+            region.drain()
         } finally {
             inRegion.set(false)
         }
+        region.awaitClaimed()
+        region.rethrow()
         if (sink.isEnabled) {
             val now = TraceClock.nowNanos()
             sink.emit(TraceEvent.ScheduleRegion(op = "forRange", schedule = name, elements = n, tasks = tasks, durationNanos = now - started, timeNanos = now))
         }
     }
 
-    private fun inRegion(start: Int, end: Int, body: (Int, Int) -> Unit) {
-        val previous = inRegion.get()
-        inRegion.set(true)
-        try {
-            body(start, end)
-        } finally {
-            inRegion.set(previous)
+    /** One `forRange` call: the chunk queue, the completion latch and the first failure. */
+    private class Region(
+        private val n: Int,
+        private val chunk: Int,
+        private val tasks: Int,
+        private val body: (Int, Int) -> Unit,
+    ) : Runnable {
+        /** Next chunk to claim; the caller takes chunk 0 before the helpers start. */
+        private val next = AtomicInteger(1)
+
+        /** Counts down once per chunk, whether it ran or was skipped after a failure. */
+        private val remaining = CountDownLatch(tasks)
+
+        private val failure = AtomicReference<Throwable?>(null)
+
+        /** Helper entry point: run on a pool thread, claim chunks until none are left. */
+        override fun run() {
+            val previous = inRegion.get()
+            inRegion.set(true)
+            try {
+                drain()
+            } finally {
+                inRegion.set(previous)
+            }
+        }
+
+        fun drain() {
+            while (true) {
+                val i = next.getAndIncrement()
+                if (i >= tasks) return
+                runChunk(i)
+            }
+        }
+
+        fun runChunk(i: Int) {
+            try {
+                if (failure.get() == null) body(i * chunk, minOf((i + 1) * chunk, n))
+            } catch (t: Throwable) {
+                if (!failure.compareAndSet(null, t)) failure.get()!!.addSuppressed(t)
+            } finally {
+                remaining.countDown()
+            }
+        }
+
+        /**
+         * Waits for the chunks other threads have claimed. By the time the caller gets here it
+         * has drained the queue, so every outstanding chunk is on a thread that is running it.
+         */
+        fun awaitClaimed() {
+            var interrupted = false
+            while (true) {
+                try {
+                    remaining.await()
+                    break
+                } catch (_: InterruptedException) {
+                    interrupted = true
+                }
+            }
+            if (interrupted) Thread.currentThread().interrupt()
+        }
+
+        fun rethrow() {
+            failure.get()?.let { throw it }
         }
     }
 
@@ -97,7 +159,7 @@ public open class CoroutineSchedule @JvmOverloads constructor(
 
         /**
          * A schedule with its own daemon pool of `parallelism - 1` workers (the caller is the last
-         * worker), for code that already runs on `Dispatchers.Default`. Close it when done.
+         * worker), for code that wants isolation from `Dispatchers.Default`. Close it when done.
          */
         @JvmStatic
         @JvmOverloads

skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/schedule/CoroutineScheduleTest.kt

@@ -1,13 +1,17 @@
 package sk.ainet.exec.schedule
 
 import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.asCoroutineDispatcher
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.runBlocking
 import sk.ainet.context.schedule.Schedule
 import sk.ainet.lang.memory.ExperimentalMemoryApi
 import sk.ainet.lang.memory.trace.RecordingTraceSink
 import sk.ainet.lang.memory.trace.TraceEvent
 import java.util.BitSet
+import java.util.concurrent.CyclicBarrier
+import java.util.concurrent.Executors
+import java.util.concurrent.TimeUnit
 import java.util.concurrent.atomic.AtomicInteger
 import kotlin.test.Test
 import kotlin.test.assertEquals
@@ -60,16 +64,20 @@ class CoroutineScheduleTest {
     }
 
     @Test
-    fun callerThreadRunsTheFirstChunkAndWorkersRunTheRest() {
+    fun callerThreadRunsTheFirstChunkAndWorkersHelpWithTheRest() {
         val schedule = CoroutineSchedule(parallelism = 4)
         val caller = Thread.currentThread()
         val onCaller = AtomicInteger()
         val elsewhere = AtomicInteger()
-        schedule.forRange(4000, grain = 1) { _, _ ->
+        var first: Thread? = null
+        schedule.forRange(4000, grain = 1) { s, _ ->
+            if (s == 0) first = Thread.currentThread()
             if (Thread.currentThread() === caller) onCaller.incrementAndGet() else elsewhere.incrementAndGet()
+            Thread.sleep(30)   // long enough for the pool to claim the other chunks
         }
-        assertEquals(1, onCaller.get(), "the caller runs exactly one chunk itself")
-        assertEquals(3, elsewhere.get(), "the other chunks run on the dispatcher")
+        assertSame(caller, first, "the caller runs the first chunk itself")
+        assertTrue(elsewhere.get() >= 1, "at least one chunk ran on the dispatcher")
+        assertEquals(4, onCaller.get() + elsewhere.get(), "every chunk ran exactly once")
     }
 
     @Test
@@ -119,6 +127,33 @@ class CoroutineScheduleTest {
         assertEquals(4 * 4096, total.get())
     }
 
+    /**
+     * The CI deadlock: every thread of the schedule's own pool enters a region at once, so no
+     * pool thread is free to run the helpers. The caller must finish the region by itself. With
+     * the old `runBlocking { coroutineScope { … } }` region this hung forever on a pool of any
+     * size, because the scope waited for children the pool could never start.
+     */
+    @Test
+    fun aRegionEnteredFromEveryThreadOfItsOwnPoolStillCompletes() {
+        val poolSize = 2
+        val pool = Executors.newFixedThreadPool(poolSize)
+        try {
+            val schedule = CoroutineSchedule(dispatcher = pool.asCoroutineDispatcher(), parallelism = 4)
+            val allInside = CyclicBarrier(poolSize)
+            val total = AtomicInteger()
+            val futures = List(poolSize) {
+                pool.submit {
+                    allInside.await(10, TimeUnit.SECONDS)
+                    schedule.forRange(4096, grain = 1) { s, e -> total.addAndGet(e - s) }
+                }
+            }
+            futures.forEach { it.get(30, TimeUnit.SECONDS) }
+            assertEquals(poolSize * 4096, total.get())
+        } finally {
+            pool.shutdownNow()
+        }
+    }
+
     @Test
     fun dedicatedScheduleOwnsItsPoolAndCloses() {
         CoroutineSchedule.dedicated(parallelism = 3).use { schedule ->