CoroutineContext and ContinuationInterceptor (instead of dispatcher)

2017-01-13 21:01:57 +03:00
parent f611e39a69
commit 8d6a913cee
22 changed files with 507 additions and 148 deletions
@@ -0,0 +1,129 @@
+/*
+ * Copyright 2010-2017 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package kotlin.coroutines
+
+/**
+ * Base class for [CoroutineContextElement] implementations.
+ */
+@SinceKotlin("1.1")
+public abstract class AbstractCoroutineContextElement : CoroutineContextElement {
+    @Suppress("UNCHECKED_CAST")
+    public override operator fun <E : CoroutineContextElement> get(key: CoroutineContextKey<E>): E? =
+            if (this.contextKey == key) this as E else null
+
+    public override fun <R> fold(initial: R, operation: (R, CoroutineContextElement) -> R): R =
+            operation(initial, this)
+
+    public override operator fun plus(context: CoroutineContext): CoroutineContext =
+            plusImpl(context)
+
+    public override fun minusKey(key: CoroutineContextKey<*>): CoroutineContext =
+            if (this.contextKey == key) EmptyCoroutineContext else this
+}
+
+/**
+ * An empty coroutine context.
+ */
+@SinceKotlin("1.1")
+public object EmptyCoroutineContext : CoroutineContext {
+    public override fun <E : CoroutineContextElement> get(key: CoroutineContextKey<E>): E? = null
+    public override fun <R> fold(initial: R, operation: (R, CoroutineContextElement) -> R): R = initial
+    public override fun plus(context: CoroutineContext): CoroutineContext = context
+    public override fun minusKey(key: CoroutineContextKey<*>): CoroutineContext = this
+    public override fun hashCode(): Int = 0
+    public override fun toString(): String = "EmptyCoroutineContext"
+}
+
+//--------------------- private impl ---------------------
+
+// this class is not exposed, but is hidden inside implementations
+// this is a left-biased list, so that `plus` works naturally
+private class CombinedContext(val left: CoroutineContext, val element: CoroutineContextElement) : CoroutineContext {
+    override fun <E : CoroutineContextElement> get(key: CoroutineContextKey<E>): E? {
+        var cur = this
+        while (true) {
+            cur.element[key]?.let { return it }
+            val next = cur.left
+            if (next is CombinedContext) {
+                cur = next
+            } else {
+                return next[key]
+            }
+        }
+    }
+
+    public override fun <R> fold(initial: R, operation: (R, CoroutineContextElement) -> R): R =
+            operation(left.fold(initial, operation), element)
+
+    public override operator fun plus(context: CoroutineContext): CoroutineContext =
+            plusImpl(context)
+
+    public override fun minusKey(key: CoroutineContextKey<*>): CoroutineContext {
+        element[key]?.let { return left }
+        val newLeft = left.minusKey(key)
+        return when (newLeft) {
+            left -> this
+            EmptyCoroutineContext -> element
+            else -> CombinedContext(newLeft, element)
+        }
+    }
+
+    private fun size(): Int =
+            if (left is CombinedContext) left.size() + 1 else 2
+
+    private fun contains(element: CoroutineContextElement): Boolean =
+            get(element.contextKey) == element
+
+    private fun containsAll(context: CombinedContext): Boolean {
+        var cur = context
+        while (true) {
+            if (!contains(cur.element)) return false
+            val next = cur.left
+            if (next is CombinedContext) {
+                cur = next
+            } else {
+                return contains(next as CoroutineContextElement)
+            }
+        }
+    }
+
+    override fun equals(other: Any?): Boolean =
+            this === other || other is CombinedContext && other.size() == size() && other.containsAll(this)
+
+    override fun hashCode(): Int = left.hashCode() + element.hashCode()
+
+    override fun toString(): String =
+            "[" + fold("") { acc, element ->
+                if (acc.isEmpty()) element.toString() else acc + ", " + element
+            } + "]"
+}
+
+private fun CoroutineContext.plusImpl(context: CoroutineContext): CoroutineContext =
+        if (context == EmptyCoroutineContext) this else // fast path -- avoid lambda creation
+            context.fold(this) { acc, element ->
+                val removed = acc.minusKey(element.contextKey)
+                if (removed == EmptyCoroutineContext) element else {
+                    // make sure interceptor is always last in the context (and thus is fast to get when present)
+                    val interceptor = removed[ContinuationInterceptor]
+                    if (interceptor == null) CombinedContext(removed, element) else {
+                        val left = removed.minusKey(ContinuationInterceptor)
+                        if (left == EmptyCoroutineContext) CombinedContext(element, interceptor) else
+                            CombinedContext(CombinedContext(left, element), interceptor)
+                    }
+                }
+            }
+
@@ -11,30 +11,26 @@ import kotlin.coroutines.intrinsics.*
 * This function creates a new, fresh instance of suspendable computation every time it is invoked.
 * To start executing the created coroutine, invoke `resume(Unit)` on the returned [Continuation] instance.
 * The [completion] continuation is invoked when coroutine completes with result of exception.
- * An optional [dispatcher] may be specified to customise dispatch of continuations between suspension points inside the coroutine.
 */
@SinceKotlin("1.1")
@Suppress("UNCHECKED_CAST")
 public fun <R, T> (suspend R.() -> T).createCoroutine(
        receiver: R,
-        completion: Continuation<T>,
-        dispatcher: ContinuationDispatcher? = null
-): Continuation<Unit> = (this as kotlin.jvm.internal.SuspendFunction1<R, T>).create(receiver, withDispatcher(completion, dispatcher))
+        completion: Continuation<T>
+): Continuation<Unit> = (this as kotlin.jvm.internal.SuspendFunction1<R, T>).create(receiver, completion)

 /**
 * Starts coroutine with receiver type [R] and result type [T].
 * This function creates and start a new, fresh instance of suspendable computation every time it is invoked.
 * The [completion] continuation is invoked when coroutine completes with result of exception.
- * An optional [dispatcher] may be specified to customise dispatch of continuations between suspension points inside the coroutine.
 */
@SinceKotlin("1.1")
@Suppress("UNCHECKED_CAST")
 public fun <R, T> (suspend R.() -> T).startCoroutine(
        receiver: R,
-        completion: Continuation<T>,
-        dispatcher: ContinuationDispatcher? = null
+        completion: Continuation<T>
 ) {
-    (this as Function2<R, Continuation<T>, Any?>).invoke(receiver, withDispatcher(completion, dispatcher))
+    (this as Function2<R, Continuation<T>, Any?>).invoke(receiver, completion)
 }

 /**
@@ -42,28 +38,24 @@ public fun <R, T> (suspend R.() -> T).startCoroutine(
 * This function creates a new, fresh instance of suspendable computation every time it is invoked.
 * To start executing the created coroutine, invoke `resume(Unit)` on the returned [Continuation] instance.
 * The [completion] continuation is invoked when coroutine completes with result of exception.
- * An optional [dispatcher] may be specified to customise dispatch of continuations between suspension points inside the coroutine.
 */
@SinceKotlin("1.1")
@Suppress("UNCHECKED_CAST")
 public fun <T> (suspend () -> T).createCoroutine(
-        completion: Continuation<T>,
-        dispatcher: ContinuationDispatcher? = null
-): Continuation<Unit> = (this as kotlin.jvm.internal.SuspendFunction0<T>).create(withDispatcher(completion, dispatcher))
+        completion: Continuation<T>
+): Continuation<Unit> = (this as kotlin.jvm.internal.SuspendFunction0<T>).create(completion)

 /**
 * Starts coroutine without receiver and with result type [T].
 * This function creates and start a new, fresh instance of suspendable computation every time it is invoked.
 * The [completion] continuation is invoked when coroutine completes with result of exception.
- * An optional [dispatcher] may be specified to customise dispatch of continuations between suspension points inside the coroutine.
 */
@SinceKotlin("1.1")
@Suppress("UNCHECKED_CAST")
 public fun <T> (suspend  () -> T).startCoroutine(
-        completion: Continuation<T>,
-        dispatcher: ContinuationDispatcher? = null
+        completion: Continuation<T>
 ) {
-    (this as Function1<Continuation<T>, Any?>).invoke(withDispatcher(completion, dispatcher))
+    (this as Function1<Continuation<T>, Any?>).invoke(completion)
 }

 /**
@@ -84,20 +76,15 @@ public inline suspend fun <T> suspendCoroutine(crossinline block: (Continuation<

 // INTERNAL DECLARATIONS

-@Suppress("INVISIBLE_REFERENCE", "INVISIBLE_MEMBER", "CANNOT_OVERRIDE_INVISIBLE_MEMBER")
-private fun <T> withDispatcher(completion: Continuation<T>, dispatcher: ContinuationDispatcher?): Continuation<T> {
-    return if (dispatcher == null) completion else
-        object : kotlin.jvm.internal.DispatchedContinuation<T>, Continuation<T> by completion {
-            override val dispatcher: ContinuationDispatcher? = dispatcher
-        }
-}
-
 private val UNDECIDED: Any? = Any()
 private val RESUMED: Any? = Any()
 private class Fail(val exception: Throwable)

@PublishedApi
 internal class SafeContinuation<in T> @PublishedApi internal constructor(private val delegate: Continuation<T>) : Continuation<T> {
+    public override val context: CoroutineContext
+        get() = delegate.context
+
    @Volatile
    private var result: Any? = UNDECIDED

@@ -0,0 +1,179 @@
+/*
+ * Copyright 2010-2017 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package test.coroutines
+
+import kotlin.test.*
+import org.junit.Test
+
+class CoroutineContextTest {
+    data class CtxA(val i: Int) : AbstractCoroutineContextElement() {
+        companion object : CoroutineContextKey<CtxA>
+        override val contextKey get() = CtxA
+    }
+
+    data class CtxB(val i: Int) : AbstractCoroutineContextElement() {
+        companion object : CoroutineContextKey<CtxB>
+        override val contextKey get() = CtxB
+    }
+
+    data class CtxC(val i: Int) : AbstractCoroutineContextElement() {
+        companion object : CoroutineContextKey<CtxC>
+        override val contextKey get() = CtxC
+    }
+
+    object Disp1 : AbstractCoroutineContextElement(), ContinuationInterceptor {
+        override fun <T> interceptContinuation(continuation: CoroutineContinuation<T>): CoroutineContinuation<T> = continuation
+        override val contextKey: CoroutineContextKey<*> = ContinuationInterceptor
+        override fun toString(): String = "Disp1"
+    }
+
+    object Disp2 : AbstractCoroutineContextElement(), ContinuationInterceptor {
+        override fun <T> interceptContinuation(continuation: CoroutineContinuation<T>): CoroutineContinuation<T> = continuation
+        override val contextKey: CoroutineContextKey<*> = ContinuationInterceptor
+        override fun toString(): String = "Disp2"
+    }
+
+    @Test
+    fun testGetPlusFold() {
+        var ctx: CoroutineContext = EmptyCoroutineContext
+        assertContents(ctx)
+        assertEquals("EmptyCoroutineContext", ctx.toString())
+
+        ctx += CtxA(1)
+        assertContents(ctx, CtxA(1))
+        assertEquals("CtxA(i=1)", ctx.toString())
+        assertEquals(CtxA(1), ctx[CtxA])
+        assertEquals(null, ctx[CtxB])
+        assertEquals(null, ctx[CtxC])
+
+        ctx += CtxB(2)
+        assertContents(ctx, CtxA(1), CtxB(2))
+        assertEquals("[CtxA(i=1), CtxB(i=2)]", ctx.toString())
+        assertEquals(CtxA(1), ctx[CtxA])
+        assertEquals(CtxB(2), ctx[CtxB])
+        assertEquals(null, ctx[CtxC])
+
+        ctx += CtxC(3)
+        assertContents(ctx, CtxA(1), CtxB(2), CtxC(3))
+        assertEquals("[CtxA(i=1), CtxB(i=2), CtxC(i=3)]", ctx.toString())
+        assertEquals(CtxA(1), ctx[CtxA])
+        assertEquals(CtxB(2), ctx[CtxB])
+        assertEquals(CtxC(3), ctx[CtxC])
+
+        ctx += CtxB(4)
+        assertContents(ctx, CtxA(1), CtxC(3), CtxB(4))
+        assertEquals("[CtxA(i=1), CtxC(i=3), CtxB(i=4)]", ctx.toString())
+        assertEquals(CtxA(1), ctx[CtxA])
+        assertEquals(CtxB(4), ctx[CtxB])
+        assertEquals(CtxC(3), ctx[CtxC])
+
+        ctx += CtxA(5)
+        assertContents(ctx, CtxC(3), CtxB(4), CtxA(5))
+        assertEquals("[CtxC(i=3), CtxB(i=4), CtxA(i=5)]", ctx.toString())
+        assertEquals(CtxA(5), ctx[CtxA])
+        assertEquals(CtxB(4), ctx[CtxB])
+        assertEquals(CtxC(3), ctx[CtxC])
+    }
+
+    @Test
+    fun testMinusKey() {
+        var ctx: CoroutineContext = CtxA(1) + CtxB(2) + CtxC(3)
+        assertContents(ctx, CtxA(1), CtxB(2), CtxC(3))
+        assertEquals("[CtxA(i=1), CtxB(i=2), CtxC(i=3)]", ctx.toString())
+
+        ctx = ctx.minusKey(CtxA)
+        assertContents(ctx, CtxB(2), CtxC(3))
+        assertEquals("[CtxB(i=2), CtxC(i=3)]", ctx.toString())
+        assertEquals(null, ctx[CtxA])
+        assertEquals(CtxB(2), ctx[CtxB])
+        assertEquals(CtxC(3), ctx[CtxC])
+
+        ctx = ctx.minusKey(CtxC)
+        assertContents(ctx, CtxB(2))
+        assertEquals("CtxB(i=2)", ctx.toString())
+        assertEquals(null, ctx[CtxA])
+        assertEquals(CtxB(2), ctx[CtxB])
+        assertEquals(null, ctx[CtxC])
+
+        ctx = ctx.minusKey(CtxC)
+        assertContents(ctx, CtxB(2))
+        assertEquals("CtxB(i=2)", ctx.toString())
+        assertEquals(null, ctx[CtxA])
+        assertEquals(CtxB(2), ctx[CtxB])
+        assertEquals(null, ctx[CtxC])
+
+        ctx = ctx.minusKey(CtxB)
+        assertContents(ctx)
+        assertEquals("EmptyCoroutineContext", ctx.toString())
+        assertEquals(null, ctx[CtxA])
+        assertEquals(null, ctx[CtxB])
+        assertEquals(null, ctx[CtxC])
+
+        assertEquals(EmptyCoroutineContext, ctx)
+    }
+
+    @Test
+    fun testPlusCombined() {
+        val ctx1 = CtxA(1) + CtxB(2)
+        val ctx2 = CtxB(3) + CtxC(4)
+        val ctx = ctx1 + ctx2
+        assertContents(ctx, CtxA(1), CtxB(3), CtxC(4))
+        assertEquals("[CtxA(i=1), CtxB(i=3), CtxC(i=4)]", ctx.toString())
+        assertEquals(CtxA(1), ctx[CtxA])
+        assertEquals(CtxB(3), ctx[CtxB])
+        assertEquals(CtxC(4), ctx[CtxC])
+    }
+
+    @Test
+    fun testLastDispatcher() {
+        var ctx: CoroutineContext = EmptyCoroutineContext
+        assertContents(ctx)
+        ctx += CtxA(1)
+        assertContents(ctx, CtxA(1))
+        ctx += Disp1
+        assertContents(ctx, CtxA(1), Disp1)
+        ctx += CtxA(2)
+        assertContents(ctx, CtxA(2), Disp1)
+        ctx += CtxB(3)
+        assertContents(ctx, CtxA(2), CtxB(3), Disp1)
+        ctx += Disp2
+        assertContents(ctx, CtxA(2), CtxB(3), Disp2)
+        ctx += (CtxB(4) + CtxC(5))
+        assertContents(ctx, CtxA(2), CtxB(4), CtxC(5), Disp2)
+    }
+
+    @Test
+    fun testEquals() {
+        val ctx1 = CtxA(1) + CtxB(2) + CtxC(3)
+        val ctx2 = CtxB(2) + CtxC(3) + CtxA(1) // same
+        val ctx3 = CtxC(3) + CtxA(1) + CtxB(2) // same
+        val ctx4 = CtxA(1) + CtxB(2) + CtxC(4) // different
+        assertEquals(ctx1, ctx2)
+        assertEquals(ctx1, ctx3)
+        assertEquals(ctx2, ctx3)
+        assertNotEquals(ctx1, ctx4)
+        assertNotEquals(ctx2, ctx4)
+        assertNotEquals(ctx3, ctx4)
+    }
+
+    private fun  assertContents(ctx: CoroutineContext, vararg elements: CoroutineContextElement) {
+        val set = ctx.fold(setOf<CoroutineContext>()) { a, b -> a + b }
+        assertEquals(listOf(*elements), set.toList())
+        for (elem in elements)
+            assertTrue(ctx[elem.contextKey] == elem)
+    }
+}