kotlin coroutines withContext() 切换线程源码分析

1.代码分析

GlobalScope.launch(Dispatchers.Default) {
    print("执行")

    val x = withContext(EmptyCoroutineContext) {
        print("当前线程：${Thread.currentThread().name}")
        delay(10)
        "withContext"
    }

    print(x)
}

1-1 kotlin show kotlin Bytecode

public final class FlowLaunchKt {
   public static final void flowLaunch() {
      BuildersKt.launch$default((CoroutineScope)GlobalScope.INSTANCE, (CoroutineContext)Dispatchers.getDefault(), (CoroutineStart)null, (Function2)(new Function2((Continuation)null) {
         int label;

         @Nullable
         public final Object invokeSuspend(@NotNull Object $result) {
            
            //是否挂起标识
            Object var3 = IntrinsicsKt.getCOROUTINE_SUSPENDED();
            
            //当前的结果 
            Object var10000;
            
            //状态机label
            switch (this.label) {
               case 0:
                  ResultKt.throwOnFailure($result);
                  Coroutine2Kt.printlnLog("执行");
                  CoroutineContext var4 = (CoroutineContext)Dispatchers.getIO();
                  
                  //withContext{} 协程体
                  Function2 var10001 = (Function2)(new Function2((Continuation)null) {
                     int label;

                     @Nullable
                     public final Object invokeSuspend(@NotNull Object $result) {
                        Object var2 = IntrinsicsKt.getCOROUTINE_SUSPENDED();
                        switch (this.label) {
                           case 0:
                              ResultKt.throwOnFailure($result);
                              StringBuilder var10000 = (new StringBuilder()).append("当前线程：");
                              Thread var10001 = Thread.currentThread();
                              Intrinsics.checkNotNullExpressionValue(var10001, "Thread.currentThread()");
                              Coroutine2Kt.printlnLog(var10000.append(var10001.getName()).toString());
                              this.label = 1;
                              
                              // delay(,this)函数是否挂起，是return,传入var10001的
                              // continuation 在合适的实际resume()->invokeSuspend(label=1)
                              if (DelayKt.delay(10L, this) == var2) {
                                 return var2;
                              }
                              break;
                           case 1:
                           
                               // 继续执行
                              ResultKt.throwOnFailure($result);
                              break;
                           default:
                              throw new IllegalStateException("call to 'resume' before 'invoke' with coroutine");
                        }
                        
                        // withContext(){} 返回值 
                        return "withContext";
                     }
                     
                     //创建Function2 并传入launch()函数的contiuation 对象
                     @NotNull
                     public final Continuation create(@Nullable Object value, @NotNull Continuation completion) {
                        Intrinsics.checkNotNullParameter(completion, "completion");
                        Function2 var3 = new <anonymous constructor>(completion);
                        return var3;
                     }
                      
                     //var1 ==当前withContext的 CoroutineScope() ,
                     //var2 == launch() 的CoroutineScope()
                     public final Object invoke(Object var1, Object var2) {
                     //
                        return ((<undefinedtype>)this.create(var1, (Continuation)var2)).invokeSuspend(Unit.INSTANCE);
                     }
                  });
                  this.label = 1;
                  var10000 = BuildersKt.withContext(var4, var10001, this);
                  
                  //var10000 == "withContext"
                  if (var10000 == var3) {
                     return var3;
                  }
                  break;
               case 1:
                  ResultKt.throwOnFailure($result);
                  var10000 = $result;
                  break;
               default:
                  throw new IllegalStateException("call to 'resume' before 'invoke' with coroutine");
            }
            
            //打印 "withContext"
            String x = (String)var10000;
            Coroutine2Kt.printlnLog(x);
            return Unit.INSTANCE;
         }

         @NotNull
         public final Continuation create(@Nullable Object value, @NotNull Continuation completion) {
            Intrinsics.checkNotNullParameter(completion, "completion");
            Function2 var3 = new <anonymous constructor>(completion);
            return var3;
         }

         public final Object invoke(Object var1, Object var2) {
            return ((<undefinedtype>)this.create(var1, (Continuation)var2)).invokeSuspend(Unit.INSTANCE);
         }
      }), 2, (Object)null);
   }

2.withContext()

public suspend fun <T> withContext(
 context: CoroutineContext,
 block: suspend CoroutineScope.() -> T
): T {
 contract {
     callsInPlace(block, InvocationKind.EXACTLY_ONCE)
 }
 
 // suspendCoroutineUninterceptedOrReturn 消除回调的挂起函数，
 // 函数返回值意义：CoroutineSingletons { COROUTINE_SUSPENDED, UNDECIDED, RESUMED }
 // 当 Any==COROUTINE_SUSPENDED 为真挂起函数
 // uCont：外部挂起函数的Continuation对象 通过resume()返回结果&恢复
 // (切换状态机label,继续执行invokeSuspend())
 
 return suspendCoroutineUninterceptedOrReturn sc@ { uCont ->
     val oldContext = uCont.context
     val newContext = oldContext.newCoroutineContext(context)
     newContext.ensureActive()
     //如 context= EmptyCoroutineContext 
     // 新旧上下文相等，直接在新 ScopeCoroutine 执行当前block()
     if (newContext === oldContext) {
         val coroutine = ScopeCoroutine(newContext, uCont)
         return@sc coroutine.startUndispatchedOrReturn(coroutine, block)
     }
     
     //dispatcher相同，其它有变化(CoroutineName..)
     if (newContext[ContinuationInterceptor] == oldContext[ContinuationInterceptor]) {
         
         //在新的上下文 UndispatchedCoroutine 中执行block()
         val coroutine = UndispatchedCoroutine(newContext, uCont)
         withCoroutineContext(newContext, null) {
             return@sc coroutine.startUndispatchedOrReturn(coroutine, block)
         }
     }
     
     //dispatcher 变化 Dispatchers.Main->Dispatchers.IO
     // 在 DispatchedCoroutine开启
     val coroutine = DispatchedCoroutine(newContext, uCont)
     block.startCoroutineCancellable(coroutine, coroutine)
     //suspendCoroutineUninterceptedOrReturn 返回值
-       //  coroutine.getResult() == COROUTINE_SUSPENDED 协程挂起
     coroutine.getResult()
 }
}

2-1 suspendCoroutineUninterceptedOrReturn 函数

//Implementation of suspendCoroutineUninterceptedOrReturn is intrinsic
//编译器内建函数，由kotlin 编译器实现

public suspend inline fun <T> suspendCoroutineUninterceptedOrReturn(crossinline block: (Continuation<T>) -> Any?): T {
    contract { callsInPlace(block, InvocationKind.EXACTLY_ONCE) }
    throw NotImplementedError("Implementation of suspendCoroutineUninterceptedOrReturn is intrinsic")
}

1.类型参数 block: (Continuation<T>) -> Any?
    接收一个Lambda,并且 Any 返回值标志该函数是否真正挂起
    Any == CoroutineSingletons.COROUTINE_SUSPENDED 挂起标识
    
2.crossinline ：不允许出现ruturn语句 但可以可以局部return(return sc@ "hello")

2-1-1 使用suspendCoroutineUninterceptedOrReturn 挂起函数

suspendCancellableCoroutine & suspendCoroutine 底层都是 suspendCoroutineUninterceptedOrReturn实现

suspend fun testSuspendCoroutine() = (suspendCoroutineUninterceptedOrReturn {uCont-> 
    thread { 
        uCont.resume("")
    }
    //挂起标识 
    return@suspendCoroutineUninterceptedOrReturn COROUTINE_SUSPENDED
}

@Nullable
public static final Object testSuspendCoroutine(@NotNull Continuation $completion) {
   int var2 = false;
   ThreadsKt.thread$default(false, false, (ClassLoader)null, (String)null, 0, (Function0)(new FlowLaunchKt$testSuspendCoroutine$2$1($completion)), 31, (Object)null);
   
   Object var10000 = IntrinsicsKt.getCOROUTINE_SUSPENDED();
   if (var10000 == IntrinsicsKt.getCOROUTINE_SUSPENDED()) {
      DebugProbesKt.probeCoroutineSuspended($completion);
   }

   return var10000;
}

2-1-2 crossinline

inline fun block(action: () -> Unit) {
    println("before")
    action()
    println("after")
}

inline fun blockNoReturn(crossinline action: () -> Unit) {
    println("before")
    action()
    println("after")
}

fun mainxxx() {
    //'return' is not allowed here
    // can return s@
    //输出 after
    blockNoReturn s@{
        return@s
    }

    //不会输出after
    block {
        return
    }
}

3. withContext 切换线程后如何恢复到原来的线程？

graph TD
DispatchedContinuation.resumeCancellableWith --> dispatcher.dispatch -->
 
DispatchedTask.run --> launch协程体continuation.resume --> BaseContinuationImpl.resume --> invokeSuspend!=COROUTINE_SUSPENDED --> completion.DispatchedCoroutine.resume --> afterResume --> uContinuation.interceted.resumeCancellableWith

public suspend fun <T> withContext(
    context: CoroutineContext,
    block: suspend CoroutineScope.() -> T
): T {
   
    return suspendCoroutineUninterceptedOrReturn sc@ { uCont ->
        ...
        // 新dispatcher 
        val coroutine = DispatchedCoroutine(newContext, uCont)
        block.startCoroutineCancellable(coroutine, coroutine)
        coroutine.getResult()
    }
}

internal fun <R, T> (suspend (R) -> T).startCoroutineCancellable(
   receiver: R, completion: Continuation<T>,
   onCancellation: ((cause: Throwable) -> Unit)? = null
) =

   //创建协程体 创建DispatchedContinuation
   runSafely(completion) {
       createCoroutineUnintercepted(receiver, completion).intercepted().resumeCancellableWith(Result.success(Unit), onCancellation)
   }

ContinuationImpl.kt

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

CoroutineDispatcher.kt
 public final override fun <T> interceptContinuation(continuation: Continuation<T>): Continuation<T> =
        DispatchedContinuation(this, continuation)

BaseContimuationImpl:continuation ---  DispatchedTask.run() 协程体 continuation 
                                                 
  1. standaloneCoroutine d3fdc5d  -> block.startCoroutineCancellable(receiver, completion) -> Continuationimpl  23786 -> interceptContinuation(this) 23786 -> dispatchedTask.run-> continuation.resume()  

2. dispatchedCoroutine 159411e  -> block.startCoroutineCancellable(receiver, completion) -> Continuationimpl  23931 -> interceptContinuation(this) 24058 -> dispatchedTask.run-> continuation.resume()  

DispatchedContinuation.kt

// by continuation 类委托
internal class DispatchedContinuation<in T>(
    @JvmField val dispatcher: CoroutineDispatcher,
    @JvmField val continuation: Continuation<T>
) : DispatchedTask<T>(MODE_UNINITIALIZED), CoroutineStackFrame, Continuation<T> by continuation {
       
       
    //外部协程体 suspendLambda 的代理对象 continuation
    // DispatchedTask delegate,task执行外部，通过 
    // continuation.resume 继续状态机label invokeSuspend()执行
    override val delegate: Continuation<T>
        get() = this

    
    @Suppress("NOTHING_TO_INLINE")
    inline fun resumeCancellableWith(
        result: Result<T>,
        noinline onCancellation: ((cause: Throwable) -> Unit)?
    ) {
        val state = result.toState(onCancellation)
        if (dispatcher.isDispatchNeeded(context)) {
            _state = state
            resumeMode = MODE_CANCELLABLE
            ①
            //线程分发
            dispatcher.dispatch(context, this)
        } else {
            executeUnconfined(state, MODE_CANCELLABLE) {
                if (!resumeCancelled(state)) {
                    resumeUndispatchedWith(result)
                }
            }
        }
    }
}

DispatchedTask.kt

public final override fun run() {
    val taskContext = this.taskContext
    var fatalException: Throwable? = null
    try {
        val delegate = delegate as DispatchedContinuation<T>
        val continuation = delegate.continuation
           ...
            if (job != null && !job.isActive) {
                val cause = job.getCancellationException()
                cancelCompletedResult(state, cause)
                continuation.resumeWithStackTrace(cause)
            } else {
                if (exception != null) {
                    continuation.resumeWithException(exception)
                } else {
                     ② launch continuation 状态机 label 的恢复
                     // 执行 BaseContinuationImpl.resume()
                    // block 执行完毕，协程恢复
                    continuation.resume(getSuccessfulResult(state))
                }
            }
        }
    } 
}

BaseContinuationImpl.kt

@SinceKotlin("1.3")
internal abstract class BaseContinuationImpl(
    public val completion: Continuation<Any?>?
) : Continuation<Any?>, CoroutineStackFrame, Serializable {
    public final override fun resumeWith(result: Result<Any?>) {
       
        var current = this
        var param = result
        while (true) {
            with(current) {
                val outcome: Result<Any?> =
                    try {
                        ③//协程体label状态机的流转
                        val outcome = invokeSuspend(param)
                        //COROUTINE_SUSPENDED 协程挂起，等待resume 
                        if (outcome === COROUTINE_SUSPENDED) return
                        Result.success(outcome)
                    } catch (exception: Throwable) {
                        Result.failure(exception)
                    }
                releaseIntercepted()
                if (completion is BaseContinuationImpl) {
                    current = completion
                    param = outcome
                } else {
                    ④ withContext协程体执行完毕
                    //如果当前 completion == DispatchedCoroutine
                    ⑨ launch协程体执行完毕
                    // 如果 completion == StandaloneCoroutine
                    completion.resumeWith(outcome)
                    return
                }
            }
        }
    }

3-1 DispatchedContine.kt

//uCont外部continuation 对象
internal class DispatchedCoroutine<in T>(
    context: CoroutineContext,
    uCont: Continuation<T>
) : ScopeCoroutine<T>(context, uCont) {
    // todo: we may some-how abstract it via inline class
    private val _decision = atomic(UNDECIDED)

    private fun tryResume(): Boolean {
        _decision.loop { decision ->
            when (decision) {
                UNDECIDED -> if (this._decision.compareAndSet(UNDECIDED, RESUMED)) return true
                SUSPENDED -> return false
                else -> error("Already resumed")
            }
        }
    }

    override fun afterCompletion(state: Any?) {
        afterResume(state)
    }
    
   
    override fun afterResume(state: Any?) {
        if (tryResume()) return 
        
         //⑧ uCont： StandaloneCoroutine 恢复执行
        //恢复外部协程的关键
       
        //private var intercepted: Continuation<Any?>? = null
        //public fun intercepted(): Continuation<Any?> =
        //   intercepted
        //        ?: //(context[ContinuationInterceptor]?.interceptContinuation(this) ?: this)
                   // .also { intercepted = it }
        //重新获取intercepted == DispathcedContinuation(),执行task
        //并且state=="withContext"传递给外部协程体的continuation
        uCont.intercepted().resumeCancellableWith(recoverResult(state, uCont))
    }
    
    //结果返回值
    fun getResult(): Any? {
        
        //默认值UNDECIDED -> trysuspend() -> COROUTINE_SUSPENDED 协程挂起
        if (trySuspend()) return COROUTINE_SUSPENDED
        // 检查当前job 是否异常终止
        val state = this.state.unboxState()
        if (state is CompletedExceptionally) throw state.cause
        // 否在直接返回当前值
        return state as T
    }
}

public abstract class AbstractCoroutine<in T>(
    parentContext: CoroutineContext,
    initParentJob: Boolean,
    active: Boolean
) : JobSupport(active), Job, Continuation<T>, CoroutineScope {
    
     //⑥ withContext协程恢复 
    public final override fun resumeWith(result: Result<T>) {
        //检查当前协程
        //result.toState() == Success("withContext")
        
        //state job状态符号
        //COMPLETING_ALREADY -- when already complete or completing
        // COMPLETING_RETRY -- when need to retry due to interference
        // COMPLETING_WAITING_CHILDREN -- when made completing and is waiting for children
        // final state -- when completed, for call to afterCompletion
        val state = makeCompletingOnce(result.toState())
        // 等待子协程完成，父协程才能完成
        //launch{
        //    launch1{
        //        delay(1000)
        //        print("launch1")
        //    }
        //    launch2{
        //        delay(2000)
        //        print("launch1")
        //    }
        //    print("launch")
        // }
        // logcat: launch  launch1  launch2 
        //此时 launch 需要等待launch1,launch2的完成，才能完成
        if (state === COMPLETING_WAITING_CHILDREN) return
        afterResume(state)
    }
    
    //⑦ 协程恢复后的操作
    protected open fun afterResume(state: Any?): Unit = afterCompletion(state)


}