Introduce WorkStealingDispatcher.

Author: zach-klippenstein

build.gradle.kts

@@ -1,4 +1,5 @@
 import com.squareup.workflow1.buildsrc.shardConnectedCheckTasks
+import org.gradle.api.tasks.testing.logging.TestExceptionFormat.FULL
 import org.jetbrains.dokka.gradle.AbstractDokkaLeafTask
 import java.net.URL
 
@@ -101,6 +102,17 @@ subprojects {
     .configureEach { mustRunAfter(tasks.matching { it is Sign }) }
 }
 
+subprojects {
+  tasks.withType(Test::class.java)
+    .configureEach {
+      testLogging {
+        // This prints exception messages and stack traces to the log when tests fail. Makes it a
+        // lot easier to see what failed in CI. If this gets too noisy, just remove it.
+        exceptionFormat = FULL
+      }
+    }
+}
+
 // This task is invoked by the documentation site generator script in the main workflow project (not
 // in this repo), which also expects the generated files to be in a specific location. Both the task
 // name and destination directory are defined in this script:

workflow-runtime/build.gradle.kts

@@ -17,6 +17,10 @@ kotlin {
   if (targets == "kmp" || targets == "js") {
     js(IR) { browser() }
   }
+
+  // Needed for expect class Lock, which is not public API, so this doesn't add any binary compat
+  // risk.
+  compilerOptions.freeCompilerArgs.add("-Xexpect-actual-classes")
 }
 
 dependencies {

workflow-runtime/src/appleMain/kotlin/com/squareup/workflow1/internal/Synchronization.apple.kt

@@ -0,0 +1,14 @@
+package com.squareup.workflow1.internal
+
+import platform.Foundation.NSLock
+
+internal actual typealias Lock = NSLock
+
+internal actual inline fun <R> Lock.withLock(block: () -> R): R {
+  lock()
+  try {
+    return block()
+  } finally {
+    unlock()
+  }
+}

workflow-runtime/src/commonMain/kotlin/com/squareup/workflow1/internal/Synchronization.kt

@@ -0,0 +1,5 @@
+package com.squareup.workflow1.internal
+
+internal expect class Lock()
+
+internal expect inline fun <R> Lock.withLock(block: () -> R): R

workflow-runtime/src/commonMain/kotlin/com/squareup/workflow1/internal/WorkStealingDispatcher.kt

@@ -0,0 +1,239 @@
+package com.squareup.workflow1.internal
+
+import com.squareup.workflow1.internal.WorkStealingDispatcher.Companion.wrapDispatcherFrom
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.Delay
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.InternalCoroutinesApi
+import kotlinx.coroutines.Runnable
+import kotlin.coroutines.Continuation
+import kotlin.coroutines.ContinuationInterceptor
+import kotlin.coroutines.CoroutineContext
+import kotlin.coroutines.resume
+
+/**
+ * A [CoroutineDispatcher] that delegates to another dispatcher but allows stealing any work
+ * scheduled on this dispatcher and performing it synchronously by calling [advanceUntilIdle].
+ *
+ * The easiest way to create one is by calling [wrapDispatcherFrom].
+ *
+ * E.g.
+ * ```
+ * val dispatcher = WorkStealingDispatcher.wrapDispatcherFrom(scope.coroutineContext)
+ * scope.launch(dispatcher) {
+ *   while (true) {
+ *     lots()
+ *     of()
+ *     suspending()
+ *     calls()
+ *   }
+ * }
+ * …
+ * dispatcher.advanceUntilIdle()
+ * ```
+ *
+ * @param delegateInterceptor The [CoroutineDispatcher] or other [ContinuationInterceptor] to
+ * delegate scheduling behavior to. This can either be a confined or unconfined dispatcher, and its
+ * behavior will be preserved transparently.
+ */
+internal open class WorkStealingDispatcher protected constructor(
+  private val delegateInterceptor: ContinuationInterceptor,
+  lock: Lock?,
+  queue: LinkedHashSet<DelegateDispatchedContinuation>?
+) : CoroutineDispatcher() {
+  companion object {
+    /**
+     * Creates a [WorkStealingDispatcher] that supports [Delay] if [delegateInterceptor] does.
+     */
+    operator fun invoke(delegateInterceptor: ContinuationInterceptor): WorkStealingDispatcher =
+      createMatchingDelayability(
+        delegateInterceptor = delegateInterceptor,
+        lock = null,
+        queue = null
+      )
+
+    /**
+     * Returns a [WorkStealingDispatcher] that delegates to the [CoroutineDispatcher] from
+     * [context]. If the context does not specify a dispatcher, [Dispatchers.Default] is used.
+     */
+    fun wrapDispatcherFrom(context: CoroutineContext): WorkStealingDispatcher {
+      // If there's no dispatcher in the context then the coroutines runtime will fall back to
+      // Dispatchers.Default anyway.
+      val baseDispatcher = context[ContinuationInterceptor] ?: Dispatchers.Default
+      return invoke(delegateInterceptor = baseDispatcher)
+    }
+
+    @OptIn(InternalCoroutinesApi::class)
+    private fun createMatchingDelayability(
+      delegateInterceptor: ContinuationInterceptor,
+      lock: Lock?,
+      queue: LinkedHashSet<DelegateDispatchedContinuation>?
+    ): WorkStealingDispatcher {
+      return if (delegateInterceptor is Delay) {
+        DelayableWorkStealingDispatcher(
+          delegate = delegateInterceptor,
+          delay = delegateInterceptor,
+          lock = lock,
+          queue = queue
+        )
+      } else {
+        WorkStealingDispatcher(
+          delegateInterceptor = delegateInterceptor,
+          lock = lock,
+          queue = queue
+        )
+      }
+    }
+  }
+
+  /** Used to synchronize access to the mutable properties of this class. */
+  private val lock = lock ?: Lock()
+
+  // region Access to these properties must always be synchronized with lock.
+  private val queue = queue ?: LinkedHashSet()
+  // endregion
+
+  /**
+   * Always returns true since we always need to track what work is waiting so we can advance it.
+   */
+  final override fun isDispatchNeeded(context: CoroutineContext): Boolean = true
+
+  final override fun dispatch(
+    context: CoroutineContext,
+    block: Runnable
+  ) {
+    val continuation = DelegateDispatchedContinuation(context, block)
+    lock.withLock {
+      queue += continuation
+    }
+
+    // Trampoline the dispatch outside the critical section to avoid deadlocks.
+    // This will either synchronously run block or dispatch it, depending on what resuming a
+    // continuation on the delegate dispatcher would do.
+    continuation.resumeOnDelegateDispatcher()
+  }
+
+  /**
+   * Calls [limitedParallelism] on [delegateInterceptor] and wraps the returned dispatcher with
+   * a [WorkStealingDispatcher] that this instance will steal from.
+   *
+   * This satisfies the limited parallelism requirements because [advanceUntilIdle] always runs
+   * tasks with a parallelism of 1 (i.e. serially).
+   */
+  @ExperimentalCoroutinesApi
+  final override fun limitedParallelism(parallelism: Int): CoroutineDispatcher {
+    if (delegateInterceptor !is CoroutineDispatcher) {
+      throw UnsupportedOperationException(
+        "limitedParallelism is not supported for WorkStealingDispatcher with " +
+          "non-dispatcher delegate"
+      )
+    }
+
+    val limitedDelegate = delegateInterceptor.limitedParallelism(parallelism)
+    return createMatchingDelayability(
+      delegateInterceptor = limitedDelegate,
+      lock = lock,
+      queue = queue
+    )
+  }
+
+  /**
+   * "Steals" work that was scheduled on this dispatcher but hasn't had a chance to run yet and runs
+   * it, until there is no work left to do. If the work schedules more work, that will also be ran
+   * before the method returns.
+   *
+   * Some care must be taken with this method: if it is called from a coroutine that is dispatched
+   * by this dispatcher, then it may run inside the stack frame of another call and in that case the
+   * inner call may think the dispatcher is idle when the outer call still has more tasks to run.
+   *
+   * It is technically safe to call from multiple threads, even in parallel, although the behavior
+   * is undefined. E.g. One thread might return from this method before the other has finished
+   * running all tasks.
+   */
+  // If we need a strong guarantee for reentrant calls we could use a ThreadLocal so the inner call
+  // could steal work from the outer one.
+  // If we need a strong guarantee for calling from multiple threads we could just run this method
+  // with a separate lock so all threads would just wait on the first one to finish running, but
+  // that could deadlock if any of the dispatched coroutines call this method reentrantly.
+  fun advanceUntilIdle() {
+    do {
+      val task = nextTask()
+      task?.releaseAndRun()
+    } while (task != null)
+  }
+
+  private fun nextTask(): DelegateDispatchedContinuation? {
+    lock.withLock {
+      val iterator = queue.iterator()
+      if (iterator.hasNext()) {
+        val task = iterator.next()
+        iterator.remove()
+        return task
+      } else {
+        return null
+      }
+    }
+  }
+
+  protected inner class DelegateDispatchedContinuation(
+    override val context: CoroutineContext,
+    private val runnable: Runnable
+  ) : Continuation<Unit> {
+    /**
+     * Cache for intercepted coroutine so we can release it from [resumeWith].
+     * [WorkStealingDispatcher] guarantees only one resume call will happen until the continuation
+     * is done, so we don't need to guard this property with a lock.
+     */
+    private var intercepted: Continuation<Unit>? = null
+
+    /**
+     * Resumes this continuation on [delegateInterceptor] by intercepting it and resuming the
+     * intercepted continuation. When a dispatcher returns false from [isDispatchNeeded], then when
+     * continuations intercepted by it are resumed, they may either be ran in-place or scheduled to
+     * a special thread-local queue. The only way to access this queue is to have the dispatcher
+     * intercept a continuation and resume the intercepted continuation.
+     */
+    fun resumeOnDelegateDispatcher() {
+      val intercepted = delegateInterceptor.interceptContinuation(this).also {
+        this.intercepted = it
+      }
+
+      // If delegate is a CoroutineDispatcher, intercepted will be a special Continuation that will
+      // check the delegate's isDispatchNeeded to decide whether to call dispatch() or to enqueue it
+      // to the thread-local unconfined queue.
+      intercepted.resume(Unit)
+    }
+
+    /**
+     * DO NOT CALL DIRECTLY! Call [resumeOnDelegateDispatcher] instead.
+     */
+    override fun resumeWith(result: Result<Unit>) {
+      val shouldRun = lock.withLock {
+        queue.remove(this)
+      }
+
+      if (shouldRun) {
+        releaseAndRun()
+      }
+    }
+
+    fun releaseAndRun() {
+      intercepted?.let {
+        if (it !== this) {
+          delegateInterceptor.releaseInterceptedContinuation(it)
+        }
+        intercepted = null
+      }
+      runnable.run()
+    }
+  }
+}
+
+@OptIn(InternalCoroutinesApi::class)
+private class DelayableWorkStealingDispatcher(
+  delegate: ContinuationInterceptor,
+  delay: Delay,
+  lock: Lock?,
+  queue: LinkedHashSet<DelegateDispatchedContinuation>?
+) : WorkStealingDispatcher(delegate, lock, queue), Delay by delay