WIP rewrite to support parallel delegates, Delay

Author: zach-klippenstein

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

@@ -2,17 +2,16 @@ 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.concurrent.Volatile
 import kotlin.coroutines.Continuation
 import kotlin.coroutines.ContinuationInterceptor
 import kotlin.coroutines.CoroutineContext
 import kotlin.coroutines.resume
 
-private typealias RunQueue = ArrayList<Runnable>
-
 /**
  * A [CoroutineDispatcher] that delegates to another dispatcher but allows stealing any work
  * scheduled on this dispatcher and performing it synchronously by calling [advanceUntilIdle].
@@ -21,7 +20,7 @@ private typealias RunQueue = ArrayList<Runnable>
  *
  * E.g.
  * ```
- * val dispatcher = WorkStealingDispatcher.wrappingDispatcherFrom(scope.coroutineContext)
+ * val dispatcher = WorkStealingDispatcher.wrapDispatcherFrom(scope.coroutineContext)
  * scope.launch(dispatcher) {
  *   while (true) {
  *     lots()
@@ -38,16 +37,22 @@ private typealias RunQueue = ArrayList<Runnable>
  * delegate scheduling behavior to. This can either be a confined or unconfined dispatcher, and its
  * behavior will be preserved transparently.
  */
-internal class WorkStealingDispatcher(
+internal open class WorkStealingDispatcher protected constructor(
   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.
+     * Creates a [WorkStealingDispatcher] that supports [Delay] if [delegateInterceptor] does.
      */
-    private const val INITIAL_QUEUE_CAPACITY = 3
+    @OptIn(InternalCoroutinesApi::class)
+    operator fun invoke(delegateInterceptor: ContinuationInterceptor): WorkStealingDispatcher {
+      return if (delegateInterceptor is Delay) {
+        DelayableWorkStealingDispatcher(delegateInterceptor, delegateInterceptor)
+      } else {
+        WorkStealingDispatcher(delegateInterceptor)
+      }
+    }
 
     /**
      * Returns a [WorkStealingDispatcher] that delegates to the [CoroutineDispatcher] from
@@ -57,70 +62,43 @@ internal class 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)
+      return invoke(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: RunQueue? = 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
+  private val queue = LinkedHashSet<DelegateDispatchedContinuation>()
   // 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
+  final override fun isDispatchNeeded(context: CoroutineContext): Boolean = true
 
-  override fun dispatch(
+  final override fun dispatch(
     context: CoroutineContext,
     block: Runnable
   ) {
-    var continuation: DispatchContinuation? = null
+    val continuation = DelegateDispatchedContinuation(context, block)
     lock.withLock {
-      val queue = this.queue ?: RunQueue(INITIAL_QUEUE_CAPACITY).also { this.queue = it }
-      queue += block
-
-      // If no dispatch is currently scheduled, then flag that we're handling it, and schedule one
-      // outside the critical section.
-      if (!dispatchScheduled) {
-        dispatchScheduled = true
-        continuation = dispatchContinuation ?: DispatchContinuation()
-          .also { dispatchContinuation = it }
-      }
+      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()
+    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 {
+  final 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.
@@ -148,118 +126,29 @@ internal class WorkStealingDispatcher(
   // 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: RunQueue? = 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.
-      queueToDrain?.runAll()
-    } while (queueToDrain != null)
+      val task = nextTask()
+      task?.releaseAndRun()
+    } while (task != 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: RunQueue?): RunQueue? {
-    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.
-        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.
-        null
-      }
-
-      else -> {
-        // There are queued tasks, so return the current queue and replace it with the recycled one.
-        queue.also {
-          this.queue = queueToRecycle?.apply { clear() }
-        }
+  private fun nextTask(): DelegateDispatchedContinuation? {
+    lock.withLock {
+      val iterator = queue.iterator()
+      if (iterator.hasNext()) {
+        val task = iterator.next()
+        iterator.remove()
+        return task
+      } else {
+        return null
       }
     }
   }
 
-  private fun RunQueue.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
-
+  private 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
@@ -289,18 +178,29 @@ internal class WorkStealingDispatcher(
      * 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
       }
-
-      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
-      })
+      runnable.run()
     }
   }
 }
+
+@OptIn(InternalCoroutinesApi::class)
+private class DelayableWorkStealingDispatcher(
+  delegate: ContinuationInterceptor,
+  delay: Delay
+) : WorkStealingDispatcher(delegate), Delay by delay