WIP rewrite to support parallel delegates, Delay

Author: zach-klippenstein

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

@@ -0,0 +1,25 @@
+@file:Suppress(
+  "CANNOT_OVERRIDE_INVISIBLE_MEMBER",
+  "INVISIBLE_MEMBER",
+  "INVISIBLE_REFERENCE",
+)
+
+package com.squareup.workflow1.internal
+
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.DefaultDelay
+import kotlinx.coroutines.Delay
+import kotlinx.coroutines.InternalCoroutinesApi
+import kotlin.coroutines.ContinuationInterceptor
+
+/**
+ * Ugly, terrible, no-good hack to support propagating [Delay] implementations to custom
+ * [CoroutineDispatcher]s without needing them to do all the internal suppressions we need to make
+ * this work.
+ *
+ * See https://github.com/Kotlin/kotlinx.coroutines/issues/4476.
+ */
+@OptIn(InternalCoroutinesApi::class)
+internal abstract class DelayPropagatingCoroutineDispatcher(
+  private val delegate: ContinuationInterceptor
+) : CoroutineDispatcher(), Delay by (delegate as? Delay ?: DefaultDelay)

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

@@ -5,7 +5,6 @@ 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
@@ -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()
@@ -40,15 +39,9 @@ private typealias RunQueue = ArrayList<Runnable>
  */
 internal class WorkStealingDispatcher(
   private val delegateInterceptor: ContinuationInterceptor
-) : CoroutineDispatcher() {
+) : DelayPropagatingCoroutineDispatcher(delegateInterceptor) {
 
   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.
@@ -65,25 +58,7 @@ internal class WorkStealingDispatcher(
   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
 
   /**
@@ -95,24 +70,15 @@ internal class WorkStealingDispatcher(
     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()
   }
 
   /**
@@ -148,118 +114,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 +166,23 @@ 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()
     }
   }
 }

workflow-runtime/src/commonTest/kotlin/com/squareup/workflow1/internal/WorkStealingDispatcherTest.kt

@@ -5,6 +5,7 @@ import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.Runnable
 import kotlinx.coroutines.currentCoroutineContext
+import kotlinx.coroutines.delay
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.test.StandardTestDispatcher
 import kotlinx.coroutines.test.runTest
@@ -352,7 +353,7 @@ class WorkStealingDispatcherTest {
   }
 
   @Test fun integration_confined_whenAdvanceUntilIdle() = runTest {
-    val testDispatcher = StandardTestDispatcher()
+    val testDispatcher = StandardTestDispatcher(testScheduler)
     val dispatcher = WorkStealingDispatcher(testDispatcher)
 
     expect(0)
@@ -366,7 +367,7 @@ class WorkStealingDispatcherTest {
   }
 
   @Test fun integration_confined_whenDispatched() = runTest {
-    val testDispatcher = StandardTestDispatcher()
+    val testDispatcher = StandardTestDispatcher(testScheduler)
     val dispatcher = WorkStealingDispatcher(testDispatcher)
 
     expect(0)
@@ -380,7 +381,7 @@ class WorkStealingDispatcherTest {
   }
 
   @Test fun integration_yield_whenAdvanceUntilIdle() = runTest {
-    val testDispatcher = StandardTestDispatcher()
+    val testDispatcher = StandardTestDispatcher(testScheduler)
     val dispatcher = WorkStealingDispatcher(testDispatcher)
 
     launch(dispatcher) {
@@ -399,7 +400,7 @@ class WorkStealingDispatcherTest {
   }
 
   @Test fun integration_yield_whenDispatched() = runTest {
-    val testDispatcher = StandardTestDispatcher()
+    val testDispatcher = StandardTestDispatcher(testScheduler)
     val dispatcher = WorkStealingDispatcher(testDispatcher)
 
     launch(dispatcher) {
@@ -417,6 +418,47 @@ class WorkStealingDispatcherTest {
     expect(4)
   }
 
+  @Test fun integration_delay_whenAdvanceUntilIdle() = runTest {
+    val testDispatcher = StandardTestDispatcher(testScheduler)
+    val dispatcher = WorkStealingDispatcher(testDispatcher)
+
+    launch(dispatcher) {
+      expect(0)
+      delay(20)
+      expect(4)
+    }
+    launch(dispatcher) {
+      expect(1)
+      delay(10)
+      expect(3)
+    }
+
+    dispatcher.advanceUntilIdle()
+    expect(2)
+
+    testScheduler.advanceUntilIdle()
+    expect(5)
+  }
+
+  @Test fun integration_delay_whenDispatched() = runTest {
+    val testDispatcher = StandardTestDispatcher(testScheduler)
+    val dispatcher = WorkStealingDispatcher(testDispatcher)
+
+    launch(dispatcher) {
+      expect(0)
+      delay(20)
+      expect(3)
+    }
+    launch(dispatcher) {
+      expect(1)
+      delay(10)
+      expect(2)
+    }
+
+    testDispatcher.scheduler.advanceUntilIdle()
+    expect(4)
+  }
+
   private fun CoroutineDispatcher.dispatch(block: () -> Unit) {
     dispatch(this, Runnable { block() })
   }