0x01. 协程的简单实现

130 阅读5分钟

kotlin协程体就是一个由suspend关键字修饰的lambda表达式,kotlin为该类型提供了最基本的扩展方法,来创建协程实例。

val continuation = suspend {
    println("协程开始执行")
}.createCoroutine(object : Continuation<Unit> {
    override val context: CoroutineContext
        get() = EmptyCoroutineContext

    override fun resumeWith(result: Result<Unit>) {
        println("协程执行结束")
    }
})

continuation.resume(Unit)

这样就是最简单的协程形式了。我们看到和一般的由CorotineScope启动的协程不太一样,其实launch的内部最后也差不多是这种形式,只是在这些的基础上进一步地封装使得协程更好用。

代码大致分3个部分:传入一个Continuation匿名类作为协程执行结束的回调,我们一般称为completion;通过createCoroutine方法返回一个Continuation;调用resume方法启动,这里的resume是resumeWith方法的简易版本,resumeWith方法我们马上就能见到。

Continuation接口很简单,一个context变量,一个resumeWith方法。context,上下文,作为一种经典的设计模式,一般用来提供当前代码执行所需的一些参数设置或行为模式,这里先按下不表。resumeWith方法是协程恢复的入口,也就是上面resume方法真正调用的方法。但completion实现地比较简单,只是用来表示完成时的一个回调,在createCoroutine中我们会看到更加复杂的实现。

public interface Continuation<in T> {
    /**
     * The context of the coroutine that corresponds to this continuation.
     */
    public val context: CoroutineContext

    /**
     * Resumes the execution of the corresponding coroutine passing a successful or failed [result] as the
     * return value of the last suspension point.
     */
    public fun resumeWith(result: Result<T>)
}

再来看下createCoroutine方法:

public fun <T> (suspend () -> T).createCoroutine(
    completion: Continuation<T>
): Continuation<Unit> =
    SafeContinuation(createCoroutineUnintercepted(completion).intercepted(), COROUTINE_SUSPENDED)

返回了一个SafeContinuation实例,构造方法如下:

internal expect class SafeContinuation<in T> : Continuation<T> {
    internal constructor(delegate: Continuation<T>, initialResult: Any?)

    // 其余省略
}

见文知意,入参一个是被代理的Continuation对象,一个是初始值。由此看出createCoroutineUnintercepted(completion).intercepted()这一串最后返回的还是一个Continuation。先不管,继续往下看。接下来我们就要启动协程了,即调用resumeWith方法。

public actual override fun resumeWith(result: Result<T>) {
    while (true) { // lock-free loop
        val cur = this.result // atomic read
        when {
            cur === UNDECIDED -> if (RESULT.compareAndSet(this, UNDECIDED, result.value)) return
            cur === COROUTINE_SUSPENDED -> if (RESULT.compareAndSet(this, COROUTINE_SUSPENDED, RESUMED)) {
                delegate.resumeWith(result)
                return
            }
            // 这里意味着第二次调用resumeWith就会抛出异常
            else -> throw IllegalStateException("Already resumed")
        }
    }
}

一个CAS自旋实现。curr = result = initialResult,即构造时传入的COROUTINE_SUSPENDED,意味着它直接启动了代理Coroutine。这下不得不回头看看代理Coroutine的产生了。

createCoroutineUnintercepted(completion).intercepted(),一共两步,create和intercept。看看create:

public actual fun <T> (suspend () -> T).createCoroutineUnintercepted(
    completion: Continuation<T>
): Continuation<Unit> {
    // 用于调试,忽略
    val probeCompletion = probeCoroutineCreated(completion)
    return if (this is BaseContinuationImpl)
        create(probeCompletion)
    else
        createCoroutineFromSuspendFunction(probeCompletion) {
            (this as Function1<Continuation<T>, Any?>).invoke(it)
        }
}

出现了一个类型判断,判断当前的lambda是否是BaseContinuationImpl类型。这是什么?

这是编译器的魔法,suspend lambda类型在编译后再反编译可以得到一个内部类:

static final class MainActivity$onCreate$continuation$1 extends SuspendLambda implements Function1<Continuation<? super Unit>, Object> {
    int label;

    MainActivity$onCreate$continuation$1(Continuation<? super MainActivity$onCreate$continuation$1> completion) {
        super(1, completion);
    }

    public final Continuation<Unit> create(Continuation<?> completion) {
        return (Continuation<Unit>)new MainActivity$onCreate$continuation$1((Continuation)completion);
    }

    public final Object invoke(Continuation<? super Unit> completion) {
        return ((MainActivity$onCreate$continuation$1)create(completion)).invokeSuspend(Unit.INSTANCE);
    }

    public final Object invokeSuspend(Object param1Object) {
        IntrinsicsKt.getCOROUTINE_SUSPENDED();
        switch (this.label) {
        default:
            throw new IllegalStateException("call to 'resume' before 'invoke' with coroutine");
        case 0:
            break;
        } 
        ResultKt.throwOnFailure(param1Object);
        System.out.println("协程开始执行");
        return Unit.INSTANCE;
    }
}

生成的协程内部类的继承关系为SuspendLambda -> ContinuationImpl -> BaseContinuationImpl -> Continuation,由此可见它确实是一个BaseContinuationImpl类型。

那么什么时候它不是一个BaseContinuationImpl呢,这个稍后再讲。

继续上面的类型判断,既然suspend lambda是BaseContinuationImpl类型,那么就应该走第一条分支,调用create方法。

然而create方法只能由子类实现,正好我们在上面反编译的类中看到了子类的create方法

可以看到,create方法中直接创建了当前内部类的实例并返回,也就是说createCoroutineUnintercepted返回的就是协程体的Continuation实例。

然后是intercepted:

// InstrinsicsJvm.kt
public actual fun <T> Continuation<T>.intercepted(): Continuation<T> =
    (this as? ContinuationImpl)?.intercepted() ?: this

// ContinuationImpl.kt
public fun intercepted(): Continuation<Any?> =
        intercepted
            ?: (context[ContinuationInterceptor]?.interceptContinuation(this) ?: this)
                .also { intercepted = it }

从context里取出一个对象调用拦截方法,没有就返回自身。从SuspendLambda的构造看出,这里的context等于completion的context,但是completion在实现的时候设置的context是一个EmptyCoroutineContext,相当于一个空集合,自然什么都取不出来。那么直接返回自身。这里的interceptContinuation在后文还会遇到。

至此代理Continuation就搞清楚了,就是协程体的Continuation实例。回过头再看SafeContinuation的resumeWith方法,调用delegate的resumeWith也就是调用了协程体的resumeWith方法。然而我们再查看反编译的代码,并没有看到resuemWith的实现。其实实现在BaseContinuationImpl中:

public final override fun resumeWith(result: Result<Any?>) {
        // This loop unrolls recursion in current.resumeWith(param) to make saner and shorter stack traces on resume
        var current = this
        var param = result
        while (true) {
            // Invoke "resume" debug probe on every resumed continuation, so that a debugging library infrastructure
            // can precisely track what part of suspended callstack was already resumed
            probeCoroutineResumed(current)
            with(current) {
                val completion = completion!! // fail fast when trying to resume continuation without completion
                val outcome: Result<Any?> =
                    try {
                        val outcome = invokeSuspend(param)
                        if (outcome === COROUTINE_SUSPENDED) return
                        Result.success(outcome)
                    } catch (exception: Throwable) {
                        Result.failure(exception)
                    }
                releaseIntercepted() // this state machine instance is terminating
                if (completion is BaseContinuationImpl) {
                    // unrolling recursion via loop
                    current = completion
                    param = outcome
                } else {
                    // top-level completion reached -- invoke and return
                    completion.resumeWith(outcome)
                    return
                }
            }
        }
    }

看到一个循环,但并不是无限循环的,退出条件有两个:1. invokeSuspend方法返回值为COROUTINE_SUSPENDED。2.completion类型不是BaseContinuationImpl。

查看反编译代码,内部只是打印了一句话,所以只是返回Unit,用Result类包裹起来,调用releaseIntercepted,completion是我们自己实现的,所以不是BaseContinuationImpl类型,走else分支,调用completion的resumeWith方法,传入Result,结束调用。

这就是协程的最简单实现。编译器先将suspend lambda转换(这个过程叫CPS转换)为内部类,继承SuspendLambda抽象类,祖先类即是Continuation;然后通过createCoroutine方法再封装一层保证resumeWith不被二次调用,同时传入completion到协程体中,最后resumeWith依次启动SafeContinuation -> 协程体Continuation -> completion,流程结束。通过分析,我们能体会到kotlin协程实现的周密,这是kotlin语法和编译器一起成就的。

kotlin的协程是具有切换线程的能力的,具体怎么做呢?我们注意到在调用intercepted方法时,由于completion携带的是一个EmptyCoroutineContext,返回为空,所以直接返回自身了,然而奥秘就在于此。我们把completion中的context重新赋值为Dispatchers.IO,再在执行语句的同时打印线程,就会发现执行线程改变了。

// context = EmptyCoroutineContext
I/System.out: 协程开始执行
I/System.out: currThread = [main]
I/System.out: 协程执行结束
I/System.out: currThread = [main]

// context = Dispatchers.IO
I/System.out: 协程开始执行
I/System.out: currThread = [DefaultDispatcher-worker-2]
I/System.out: 协程执行结束
I/System.out: currThread = [DefaultDispatcher-worker-2]

CoroutineContext是什么结构?ContinuationInterceptor取出的是什么?有多少种键?作用都是什么?如何自定义?

且听下回分解。