Introduce WorkStealingDispatcher.

Author: zach-klippenstein

workflow-runtime/build.gradle.kts

@@ -17,11 +17,14 @@ kotlin {
   if (targets == "kmp" || targets == "js") {
     js(IR) { browser() }
   }
+
+  compilerOptions.freeCompilerArgs.add("-Xexpect-actual-classes")
 }
 
 dependencies {
   commonMainApi(project(":workflow-core"))
   commonMainApi(libs.kotlinx.coroutines.core)
+  commonMainImplementation("androidx.collection:collection:1.5.0")
 
   commonTestImplementation(libs.kotlinx.coroutines.test.common)
   commonTestImplementation(libs.kotlin.test.core)

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,318 @@
+package com.squareup.workflow1.internal
+
+import androidx.collection.MutableObjectList
+import com.squareup.workflow1.internal.WorkStealingDispatcher.Companion.wrapDispatcherFrom
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.Runnable
+import kotlin.concurrent.Volatile
+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.wrappingDispatcherFrom(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 class WorkStealingDispatcher(
+  private val delegateInterceptor: ContinuationInterceptor
+) : CoroutineDispatcher() {
+
+  companion object {
+    /**
+     * The initial storage capacity for the task queue. We use a small queue capacity since in most
+     * cases the queue should be processed very soon after enqueuing.
+     */
+    private const val INITIAL_QUEUE_CAPACITY = 3
+
+    /**
+     * 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 WorkStealingDispatcher(delegateInterceptor = baseDispatcher)
+    }
+  }
+
+  /** Used to synchronize access to the mutable properties of this class. */
+  private val lock = Lock()
+
+  // region Access to these properties must always be synchronized with lock.
+  /**
+   * The queue of unconsumed work items. When there is no contention on the dispatcher, only one
+   * queue will ever be allocated. Only when [dispatch] is called while the queue is being processed
+   * (either by [advanceUntilIdle] or a [DispatchContinuation]) then a new queue will be allocated,
+   * but when the processing is done the old one will be placed back here to be re-used.
+   */
+  @Volatile
+  private var queue: MutableObjectList<Runnable>? = null
+
+  @Volatile
+  private var dispatchScheduled = false
+
+  /**
+   * Cached [DispatchContinuation] used to delegate to the [delegateInterceptor]'s dispatching
+   * behavior from [dispatch]. This is initialized the first call to [dispatch] that needs dispatch,
+   * and then never changed.
+   */
+  @Volatile
+  private var dispatchContinuation: DispatchContinuation? = null
+  // endregion
+
+  /**
+   * Always returns true since we always need to track what work is waiting so we can advance it.
+   */
+  override fun isDispatchNeeded(context: CoroutineContext): Boolean = true
+
+  override fun dispatch(
+    context: CoroutineContext,
+    block: Runnable
+  ) {
+    var continuation: DispatchContinuation? = null
+    lock.withLock {
+      val queue = this.queue
+        ?: MutableObjectList<Runnable>(initialCapacity = INITIAL_QUEUE_CAPACITY)
+          .also { this.queue = it }
+      queue += block
+      println("OMG dispatch: queue size = ${queue.size}")
+
+      // If no dispatch is currently scheduled, then flag that we're handling it, and schedule one
+      // outside the critical section.
+      if (!dispatchScheduled) {
+        println("OMG scheduling dispatch")
+        dispatchScheduled = true
+        continuation = dispatchContinuation ?: DispatchContinuation()
+          .also { dispatchContinuation = it }
+      }
+    }
+
+    // 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()
+  }
+
+  /**
+   * Throws [UnsupportedOperationException]. We can't allow the default implementation to run in
+   * case the delegate dispatcher would throw.
+   */
+  @ExperimentalCoroutinesApi
+  override fun limitedParallelism(parallelism: Int): CoroutineDispatcher {
+    // We could probably support this by forwarding the call to the delegate then wrapping that
+    // dispatcher with a WSD that advances when this one does, but we don't need this for our use
+    // cases and getting all the behavior correct might be hard, so don't bother for now.
+    throw UnsupportedOperationException(
+      "limitedParallelism is not supported for WorkStealingDispatcher"
+    )
+  }
+
+  /**
+   * "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() {
+    var wasDispatchScheduled = false
+    advanceUntilIdle(
+      onStartLocked = {
+        // If no dispatch was scheduled, then set the flag so that any new dispatch calls that
+        // happen while we're draining the queue won't schedule one unnecessarily since we'll just
+        // handle them.
+        // Note that we could "cancel" the dispatch if this is true here, since we're stealing all
+        // its work, but we can't cancel that task so it will just have to noop.
+        wasDispatchScheduled = dispatchScheduled.also {
+          if (!it) dispatchScheduled = true
+        }
+      },
+      onFinishedLocked = {
+        // If we set this flag above, then clear it now so future dispatch calls schedule normally.
+        dispatchScheduled = wasDispatchScheduled
+      }
+    )
+  }
+
+  /**
+   * Executes queued work items until there are none left, then returns.
+   *
+   * @param onStartLocked Called while [lock] is held exactly 1 time before any tasks are executed.
+   * @param onFinishedLocked Called while [lock] is held exactly 1 time after all tasks are finished
+   * executing.
+   */
+  private inline fun advanceUntilIdle(
+    onStartLocked: () -> Unit = {},
+    onFinishedLocked: () -> Unit
+  ) {
+    var queueToDrain: MutableObjectList<Runnable>? = null
+    do {
+      lock.withLock {
+        // Will only be null on first loop, since if it's null after this critical section we'll
+        // exit the loop.
+        if (queueToDrain == null) {
+          onStartLocked()
+        }
+
+        // We're about to overwrite queueToDrain, so put the old one back so future calls to
+        // dispatch might not need to allocate a new queue.
+        queueToDrain = consumeQueueLocked(queueToRecycle = queueToDrain).also {
+          if (it == null) {
+            onFinishedLocked()
+          }
+        }
+      }
+
+      // Drain the queue outside the critical section to ensure we don't deadlock if any of the
+      // runnables try to dispatch.
+      println("OMG draining queue of size ${queueToDrain?.size}")
+      queueToDrain?.runAll()
+    } while (queueToDrain != null)
+  }
+
+  /**
+   * If there are work items queued up, returns the queue, otherwise returns null. MUST ONLY BE
+   * CALLED while [lock] is held.
+   *
+   * If [queueToRecycle] is non-null then we try to place it back in the [queue] property for the
+   * next call to [dispatch] (after clearing it) so it won't have to allocate a new one. After this
+   * method returns [queueToRecycle] is unsafe to use for the calling code since it might be
+   * modified by another thread.
+   */
+  private fun consumeQueueLocked(
+    queueToRecycle: MutableObjectList<Runnable>?
+  ): MutableObjectList<Runnable>? {
+    if (queueToRecycle != null && queueToRecycle === queue) {
+      throw IllegalArgumentException("Cannot recycle queue with itself")
+    }
+
+    // Note: clear() iterates through the list to null everything out, so avoid calling it unless
+    // necessary.
+    val queue = this.queue
+    return when {
+      queue == null -> {
+        // The next dispatch would allocate a new queue, so recycle one if possible.
+        println("OMG recycling queue to null")
+        this.queue = queueToRecycle?.apply { clear() }
+        null
+      }
+
+      queue.isEmpty() -> {
+        // There's nothing to process in an empty queue, so don't return it at all. And since the
+        // next dispatch call already has a queue to use, so just let the recycled one be GC'd and
+        // don't bother clearing it.
+        println("OMG queue is empty, not recycling")
+        null
+      }
+
+      else -> {
+        // There are queued tasks, so return the current queue and replace it with the recycled one.
+        println("OMG recycling queue to non-null")
+        queue.also {
+          this.queue = queueToRecycle?.apply { clear() }
+        }
+      }
+    }
+  }
+
+  private fun MutableObjectList<Runnable>.runAll() {
+    forEach {
+      it.run()
+    }
+  }
+
+  /**
+   * A reusable continuation that is used to access the coroutine runtime's resumption behavior for
+   * both confined and unconfined dispatchers. See [resumeOnDelegateDispatcher] for more information
+   * on how this works.
+   *
+   * [WorkStealingDispatcher] guarantees that only one instance of this class will be created per
+   * dispatcher, and that it will never be resumed more than once concurrently, so it's safe to
+   * reuse.
+   */
+  private inner class DispatchContinuation : Continuation<Unit> {
+    override val context: CoroutineContext get() = delegateInterceptor
+
+    /**
+     * 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>) {
+      println("OMG DispatchContinuation resuming")
+      intercepted?.let {
+        if (it !== this) {
+          delegateInterceptor.releaseInterceptedContinuation(it)
+        }
+        intercepted = null
+      }
+
+      println("OMG DispatchContinuation draining queue")
+      advanceUntilIdle(onFinishedLocked = {
+        // Set this in the lock when we're about to return so that any dispatch calls waiting
+        // on the lock will know to schedule a fresh dispatch.
+        dispatchScheduled = false
+      })
+      println("OMG DispatchContinuation done draining queue")
+    }
+  }
+}