Introduce reusable cancellable continuations for hot loops with channels

  * Average performance improvement is around 25%
  * API is internal and targeted to specific usage
  * DispatchedTask and DispatchedContinuation are extracted to separate files for better readability and maintainability
  * Flow benchmark (temporary) rewritten to avoid take

Author: qwwdfsad

benchmarks/src/jmh/kotlin/benchmarks/CancellableContinuationBenchmark.kt

@@ -13,22 +13,6 @@ import kotlinx.coroutines.flow.*
 import org.openjdk.jmh.annotations.*
 import java.util.concurrent.*
 
-/*
- * Results:
- *
- * // Throw FlowAborted overhead
- * Numbers.primes            avgt    7  3039.185 ± 25.598  us/op
- * Numbers.primesRx          avgt    7  2677.937 ± 17.720  us/op
- *
- * // On par
- * Numbers.transformations   avgt    7    16.207 ±  0.133  us/op
- * Numbers.transformationsRx avgt    7    19.626 ±  0.135  us/op
- *
- * // Channels overhead
- * Numbers.zip               avgt    7   434.160 ±  7.014  us/op
- * Numbers.zipRx             avgt    7    87.898 ±  5.007  us/op
- *
- */
 @Warmup(iterations = 7, time = 1, timeUnit = TimeUnit.SECONDS)
 @Measurement(iterations = 7, time = 1, timeUnit = TimeUnit.SECONDS)
 @Fork(value = 1)
@@ -39,11 +23,11 @@ open class NumbersBenchmark {
 
     companion object {
         private const val primes = 100
-        private const val natural = 1000
+        private const val natural = 1000L
     }
 
-    private fun numbers() = flow {
-        for (i in 2L..Long.MAX_VALUE) emit(i)
+    private fun numbers(limit: Long = Long.MAX_VALUE) = flow {
+        for (i in 2L..limit) emit(i)
     }
 
     private fun primesFlow(): Flow<Long> = flow {
@@ -80,7 +64,7 @@ open class NumbersBenchmark {
 
     @Benchmark
     fun zip() = runBlocking {
-        val numbers = numbers().take(natural)
+        val numbers = numbers(natural)
         val first = numbers
             .filter { it % 2L != 0L }
             .map { it * it }
@@ -105,8 +89,7 @@ open class NumbersBenchmark {
 
     @Benchmark
     fun transformations(): Int = runBlocking {
-        numbers()
-            .take(natural)
+        numbers(natural)
             .filter { it % 2L != 0L }
             .map { it * it }
             .filter { (it + 1) % 3 == 0L }.count()
@@ -120,4 +103,4 @@ open class NumbersBenchmark {
             .filter { (it + 1) % 3 == 0L }.count()
             .blockingGet()
     }
-}
+}

benchmarks/src/jmh/kotlin/benchmarks/flow/NumbersBenchmark.kt

@@ -0,0 +1,96 @@
+/*
+ * Copyright 2016-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package benchmarks.tailcall
+
+import kotlinx.coroutines.*
+import kotlin.coroutines.*
+import kotlin.coroutines.intrinsics.*
+
+@Suppress("INVISIBLE_MEMBER", "INVISIBLE_REFERENCE")
+public abstract class SimpleChannel {
+    companion object {
+        const val NULL_SURROGATE: Int = -1
+    }
+
+    @JvmField
+    protected var producer: Continuation<Unit>? = null
+    @JvmField
+    protected var enqueuedValue: Int = NULL_SURROGATE
+    @JvmField
+    protected var consumer: Continuation<Int>? = null
+
+    suspend fun send(element: Int) {
+        require(element != NULL_SURROGATE)
+        if (offer(element)) {
+            return
+        }
+
+        return suspendSend(element)
+    }
+
+    private fun offer(element: Int): Boolean {
+        if (consumer == null) {
+            return false
+        }
+
+        consumer!!.resume(element)
+        consumer = null
+        return true
+    }
+
+    suspend fun receive(): Int {
+        // Cached value
+        if (enqueuedValue != NULL_SURROGATE) {
+            val result = enqueuedValue
+            enqueuedValue = NULL_SURROGATE
+            producer!!.resume(Unit)
+            return result
+        }
+
+        return suspendReceive()
+    }
+
+    abstract suspend fun suspendReceive(): Int
+    abstract suspend fun suspendSend(element: Int)
+}
+
+class NonCancellableChannel : SimpleChannel() {
+    override suspend fun suspendReceive(): Int = suspendCoroutineUninterceptedOrReturn {
+        consumer = it.intercepted()
+        COROUTINE_SUSPENDED
+    }
+
+    override suspend fun suspendSend(element: Int) = suspendCoroutineUninterceptedOrReturn<Unit> {
+        enqueuedValue = element
+        producer = it.intercepted()
+        COROUTINE_SUSPENDED
+    }
+}
+
+class CancellableChannel : SimpleChannel() {
+    override suspend fun suspendReceive(): Int = suspendAtomicCancellableCoroutine {
+        consumer = it.intercepted()
+        COROUTINE_SUSPENDED
+    }
+
+    override suspend fun suspendSend(element: Int) = suspendAtomicCancellableCoroutine<Unit> {
+        enqueuedValue = element
+        producer = it.intercepted()
+        COROUTINE_SUSPENDED
+    }
+}
+
+class CancellableReusableChannel : SimpleChannel() {
+    override suspend fun suspendReceive(): Int = suspendAtomicCancellableCoroutineReusable {
+        consumer = it.intercepted()
+        COROUTINE_SUSPENDED
+    }
+
+    override suspend fun suspendSend(element: Int) = suspendAtomicCancellableCoroutineReusable<Unit> {
+        enqueuedValue = element
+        producer = it.intercepted()
+        COROUTINE_SUSPENDED
+    }
+}

benchmarks/src/jmh/kotlin/benchmarks/tailcall/SimpleChannel.kt

@@ -0,0 +1,61 @@
+/*
+ * Copyright 2016-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package benchmarks.tailcall
+
+import kotlinx.coroutines.*
+import org.openjdk.jmh.annotations.*
+import java.util.concurrent.*
+
+@Warmup(iterations = 5, time = 1, timeUnit = TimeUnit.SECONDS)
+@Measurement(iterations = 5, time = 1, timeUnit = TimeUnit.SECONDS)
+@Fork(value = 1)
+@BenchmarkMode(Mode.AverageTime)
+@OutputTimeUnit(TimeUnit.MICROSECONDS)
+@State(Scope.Benchmark)
+open class SimpleChannelBenchmark {
+
+    private val iterations = 10_000
+
+    @Volatile
+    private var sink: Int = 0
+
+    @Benchmark
+    fun cancellable() = runBlocking {
+        val ch = CancellableChannel()
+        launch {
+            repeat(iterations) { ch.send(it) }
+        }
+
+        launch {
+            repeat(iterations) { sink = ch.receive() }
+        }
+    }
+
+    @Benchmark
+    fun cancellableReusable() = runBlocking {
+        val ch = CancellableReusableChannel()
+        launch {
+            repeat(iterations) { ch.send(it) }
+        }
+
+        launch {
+            repeat(iterations) { sink = ch.receive() }
+        }
+    }
+
+    @Benchmark
+    fun nonCancellable() = runBlocking {
+        val ch = NonCancellableChannel()
+        launch {
+            repeat(iterations) { ch.send(it) }
+        }
+
+        launch {
+            repeat(iterations) {
+                sink = ch.receive()
+            }
+        }
+    }
+}

benchmarks/src/jmh/kotlin/benchmarks/tailcall/SimpleChannelBenchmark.kt

@@ -223,6 +223,42 @@ public suspend inline fun <T> suspendAtomicCancellableCoroutine(
         cancellable.getResult()
     }
 
+/**
+ *  Suspends coroutine similar to [suspendAtomicCancellableCoroutine], but an instance of [CancellableContinuationImpl] is reused if possible.
+ */
+internal suspend inline fun <T> suspendAtomicCancellableCoroutineReusable(
+    crossinline block: (CancellableContinuation<T>) -> Unit
+): T = suspendCoroutineUninterceptedOrReturn { uCont ->
+        val cancellable = getOrCreateCancellableContinuation(uCont.intercepted())
+        block(cancellable)
+        cancellable.getResult()
+    }
+
+internal fun <T> getOrCreateCancellableContinuation(delegate: Continuation<T>): CancellableContinuationImpl<T> {
+    // If used outside of our dispatcher
+    if (delegate !is DispatchedContinuation<T>) {
+        return CancellableContinuationImpl(delegate, resumeMode = MODE_ATOMIC_DEFAULT)
+    }
+
+    /*
+     * Attempt to claim reusable instance.
+     *
+     * suspendAtomicCancellableCoroutineReusable { // <- claimed
+     *     // Any asynchronous cancellation is "postponed" while this block
+     *     // is being executed
+     * } // postponed cancellation is checked here.
+     *
+     * Claim can fail for the following reasons:
+     * 1) Someone tried to make idempotent resume.
+     *    Idempotent resume is internal (used only by us) and is used only in `select`,
+     *    thus leaking CC instance for indefinite time.
+     * 2) Continuation was cancelled. Then we should prevent any further reuse and bail out.
+     */
+    val claimed = delegate.claimReusableCancellableContinuation() ?: return CancellableContinuationImpl(delegate, MODE_ATOMIC_DEFAULT)
+    claimed.resetState()
+    return claimed
+}
+
 /**
  * @suppress **Deprecated**
  */