Implementation principle of suspend function

Section 1: State Machine

State machine is a mathematical calculation model.
A state machine usually has the following parts:

  • State, a state machine must have at least two states
  • Input, input trigger state machine execution
  • Operation, corresponding to different states, perform different operations on the input value
  • After the operation is performed, the current state of the state machine is switched

State machine example:

enum class State{ S0, S1, S2 }

var initialState = S0

fun runStateMachine(input: Any): Any?{
    var value: Any?
    when(initalState){
        S0 -> {
            value = calculateS0(input)
            initialState = S1
            return value
        }
        S1 -> {
            value = calculateS1(input)
            initialState = S2
            return value
        }
        S2 -> {
            initialState = S0
            return calculateS2(input)
        }
    }
}
1.1 kotlin source code
private suspend fun test1(name: String): String{
    val codeStr = test2(name)
    return "$name -> $codeStr"
}

private suspend fun test2(str: String): String{   
    val strList = stringToList(str)    
    return strList.joinToString(separator = ",", transform = { it.code.toString() })
}

private fun stringToList(str: String): List<Char>{    
    return str.toList() 
}
1.2 code generated by compiler (details are deleted)
private fun test1(name: String, ctn: Continuation<Any?>): Any{
    //The Continuation class of this method
    class Test1ContinuationImpl(cpl: Continuation<Any?>): ContinuationImpl(cpl){
        var v0: String? = null
        var result: Any? = null
        var label: Int = 0
        
        override fun invokeSuspend(result: Any?){
            this.result = result
            test1(null, this)
        }
    }
    //Initialize Continuation instance
    val continuation = ctn as? Test1ContinuationImpl ?: Test1ContinuationImpl(ctn)
    //implement
    var result: Any? = null
    when(continuation.label){
        0 -> {  //First hanging point
            continuation.v0 = name
            continuation.label = 1
            result = test2(name, continuation as Continuation)
            if(result == COROUTINE_SUSPENDED){
                return COROUTINE_SUSPENDED
            }
        }
        1 -> { //After execution, the result is returned
            name = continuation.v0
            result = continuation.result
            ctn.completion.resume( "$name -> ${result as String}")
        }
        else -> throw IllegalStateException("call to 'resume' before 'invoke' with coroutine")
    }
    return Unit   
}

private fun test2(str: String, ctn: Continuation<Any?>): Any{
    val strList = stringToList(str)  
    ctn.resume(strList.joinToString(separator = ",", transform = { it.code.toString() }))
    return Unit
}

Conclusion:

  1. The method modified by suspend will automatically add a Continuation parameter after compilation, which comes from the suspend method or coroutine calling it;

  2. The Kotlin compiler will automatically create a state machine for each suspend method. If no other suspend methods are called in the method, it will not be created;

  3. The number of state machines in the state machine is related to the number of suspend methods invoked in the method, and the number of States is equal to the number of calls + 1.

  4. COROUTINE_SUSPENDED is an enumeration value used to indicate that execution is suspended and will not return results immediately.

PS: conclusion 1 explains why the suspend method can only be called in a coroutine or suspend method, because non suspend methods cannot provide the Continuation parameter.

Section II: Continuation

Continuation is a very simple interface, as follows:

public interface Continuation<in T> {
    public val context: CoroutineContext
    public fun resumeWith(result: Result<T>)
}

When the suspend method is completed, the resumeWith method will be called to recover from the suspended state and return the Result value. The return value is wrapped in the Result object. If the execution fails, Result.Failure will be returned.

2.1 BaseContinuationImpl

It can be seen from the source code of 1.2 that during the initialization of Continuation, it initially becomes a ContinuationImpl object, and ContinuationImpl inherits from BaseContinuationImpl. Here we sort out the relationship between them:

//The actual structure of these two classes is more complex than that shown below. Here, only the inheritance relationship between them is reflected, and other parts are omitted
abstract class ContinuationImpl(completion: Continuation<Any?>?): BaseContinuationImpl(ctn)
abstract class BaseContinuationImpl(completion: Continuation<Any?>?): Continuation<Any?>

We need to focus on the implementation of the resumeWith method in BaseContinuationImpl, which is the core of the suspend method to implement suspension and recovery.

public final override fun resumeWith(result: Result<Any?>) {
    var current = this    
    var param = result    
    while (true) {     
        with(current) {            
            val completion = completion!!
            val outcome: Result<Any?> = try {
                val outcome = invokeSuspend(param)    //  Note 1             
                if (outcome === COROUTINE_SUSPENDED) return                          
                Result.success(outcome)                
            } catch (exception: Throwable) {                     
                Result.failure(exception)                
            }       
                  //Note 2
            if (completion is BaseContinuationImpl) {
                current = completion                
                param = outcome            
            } else {               
                completion.resumeWith(outcome)                
                return            
            }        
        }    
    }
}

notes:

  1. Through the source code of 1.2, we can know that a Continuation object is instantiated in the Suspend method, and it calls itself in the invokeSuspend method of Continuation, which is the input point of the state machine. So where is the invokeSuspend method called? It's in comment 1 of the above code.

  2. The function of Continuation is just like its name. It chains the suspend method calls together. In conclusion 1 of the first section, we already know that the Continuation parameter of suspend comes from the suspend method or coroutine calling it. And when the suspend method initializes its own Continuation object, it will wrap this parameter in its own Continuation object. In this way, the Continuation object of each suspend method can access the Continuation object of the upper level method. The function of the code in Note 2 is to peel off the Continuation layer by layer until the top-level Continuation is obtained, and then call its resumeWith method to return the calculation result.

  3. The Continuation objects instantiated in the suspend method inherit from BaseContinuationImpl, while the outermost Continuation object inherits from AbstractCoroutine, which we will introduce in Section 3.

Section 3: launch process

This section is recommended to read against the source code

Let's first look at the launch method of CoroutineScope, because the suspend method runs in a CoroutineScope in the end.

public fun CoroutineScope.launch(    
    context: CoroutineContext = EmptyCoroutineContext,    
    start: CoroutineStart = CoroutineStart.DEFAULT,    
    block: suspend CoroutineScope.() -> Unit): Job {
      // Note 1
      val newContext = newCoroutineContext(context)
      // Note 2
      val coroutine = if (start.isLazy)        
          LazyStandaloneCoroutine(newContext, block) else        
          StandaloneCoroutine(newContext, active = true)    
      coroutine.start(start, coroutine, block)    
      return coroutine
}
  1. The first parameter is usually passed to CoroutineDispatcher. If the CoroutineScope used does not specify a thread scheduler, the method newCoroutineContext in note 1 will add Dispatchers.Default scheduler to the CoroutineContext;

  2. The code in Note 2 uses CoroutineContext to create Coroutine. In fact, it is also a Continuation. The relationship between them is as follows:

  3. Finally, calling the start method of Coroutine will cause the following series of calls:
    (1). AbstractCoroutine.start(...);
    (2). CoroutineStart.invoke(...);
    (3). block.startCoroutineCancellable(...).

Let's focus on the last step:

internal fun <R, T> (suspend (R) -> T).startCoroutineCancellable(    
       receiver: R, completion: Continuation<T>,    
       onCancellation: ((cause: Throwable) -> Unit)? = null
) = runSafely(completion) {
       //Convert block to Continuation
       createCoroutineUnintercepted(receiver,completion)
           //Add interceptor (actually setting coroutine dispatcher)
           .intercepted()
           //implement
           .resumeCancellableWith(Result.success(Unit), onCancellation)    
}

//The source code of the first step is as follows:
private inline fun <T> createCoroutineFromSuspendFunction(completion: Continuation<T>,  crossinline block: (Continuation<T>) -> Any?) : Continuation<Unit>{
   val context = completion.context
   return if (context === EmptyCoroutineContext){
           //ellipsis
   } else {
       object : ContinuationImp(completion as Continuation<Any?>, context) {
           private var label = 0
           override fun invokeSuspend(result: Result<Any?>): Any?  = 
               when(lable){
                   0 -> {
                       lable = 1
                       result.getOrThrow()
                       block(this)
                   }
                   1 -> {
                       label = 2
                       result.getOrThrow()
                   }
                   else -> error("This coroutine had already completed")
               }
       }
   }
}

This method converts the block of code written in CoroutineScope into a BaseContinuationImpl, then calls its resumeWith method, so that it will go to its invokeSuspend method and start the execution of the whole Association.

Tags: kotlin

Posted on Fri, 29 Oct 2021 04:22:52 -0400 by NovaArgon