Kotlin 协程源码阅读笔记 —— Mutex

suspend fun helloWorld(): String {
        synchronized(this) {
            val hello = hello()
            val world = world()
            return hello + world
        }
    }

val lock = Mutex()
    suspend fun helloWorld(): String {
        lock.lock()
        val hello = hello()
        val world = world()
        lock.unlock()
        return hello + world
    }

@Suppress("FunctionName")
public fun Mutex(locked: Boolean = false): Mutex =
    MutexImpl(locked)

override suspend fun lock(owner: Any?) {
    // 尝试获取锁
    if (tryLock(owner)) return
    // 获取锁失败，挂起当前协程，直到锁可用后再恢复
    lockSuspend(owner)
}

override fun tryLock(owner: Any?): Boolean = when (tryLockImpl(owner)) {
    TRY_LOCK_SUCCESS -> true
    TRY_LOCK_FAILED -> false
    TRY_LOCK_ALREADY_LOCKED_BY_OWNER -> error("This mutex is already locked by the specified owner: $owner")
    else -> error("unexpected")
}

private fun tryLockImpl(owner: Any?): Int {
    while (true) {
        // 获取锁
        if (tryAcquire()) {
            // 获取锁成功
            assert { this.owner.value === NO_OWNER }
            // 设置新的 owner
            this.owner.value = owner
            return TRY_LOCK_SUCCESS
        } else {
            // 获取锁失败
            // The semaphore permit acquisition has failed.
            // However, we need to check that this mutex is not
            // locked by our owner.
            if (owner == null) return TRY_LOCK_FAILED
            // 检查 owner 状态
            when (holdsLockImpl(owner)) {
                // This mutex is already locked by our owner.
                HOLDS_LOCK_YES -> return TRY_LOCK_ALREADY_LOCKED_BY_OWNER
                // This mutex is locked by another owner, `trylock(..)` must return `false`.
                HOLDS_LOCK_ANOTHER_OWNER -> return TRY_LOCK_FAILED
                // This mutex is no longer locked, restart the operation.
                HOLDS_LOCK_UNLOCKED -> continue
            }
        }
    }
}

private fun holdsLockImpl(owner: Any?): Int {
    while (true) {
        // Is this mutex locked?
        if (!isLocked) return HOLDS_LOCK_UNLOCKED
        val curOwner = this.owner.value
        // Wait in a spin-loop until the owner is set
        if (curOwner === NO_OWNER) continue // 

override fun tryAcquire(): Boolean {
    while (true) {
        // Get the current number of available permits.
        val p = _availablePermits.value
        // Is the number of available permits greater
        // than the maximal one because of an incorrect
        // `release()` call without a preceding `acquire()`?
        // Change it to `permits` and start from the beginning.
        if (p > permits) {
            coerceAvailablePermitsAtMaximum()
            continue
        }
        // Try to decrement the number of available
        // permits if it is greater than zero.
        if (p 

private suspend fun lockSuspend(owner: Any?) = suspendCancellableCoroutineReusable { cont ->
    val contWithOwner = CancellableContinuationWithOwner(cont, owner)
    acquire(contWithOwner)
}

protected fun acquire(waiter: CancellableContinuation) = acquire(
    waiter = waiter,
    suspend = { cont -> addAcquireToQueue(cont as Waiter) },
    onAcquired = { cont -> cont.resume(Unit, onCancellationRelease) }
)

private inline fun  acquire(waiter: W, suspend: (waiter: W) -> Boolean, onAcquired: (waiter: W) -> Unit) {
    while (true) {
        // Decrement the number of available permits at first.
        // 先获取 _availablePermits 值，然后再减 1
        val p = decPermits()
        // Is the permit acquired?
        if (p > 0) {
            // 表示获取锁成功
            onAcquired(waiter)
            return
        }
        // Permit has not been acquired, try to suspend.
        // 执行挂起操作。
        if (suspend(waiter)) return
    }
}

private fun addAcquireToQueue(waiter: Waiter): Boolean {
    val curTail = this.tail.value
    // 获取 id
    val enqIdx = enqIdx.getAndIncrement()
    // 获取创建 Segment 的方法
    val createNewSegment = ::createSegment
    // 从链表尾部开始查找 Segment，如果没有查找到就创建一个新的 Semgent
    val segment = this.tail.findSegmentAndMoveForward(id = enqIdx / SEGMENT_SIZE, startFrom = curTail,
        createNewSegment = createNewSegment).segment // cannot be closed
        
    // 计算 Continuation 存储在 Segemnt 中的 index    
    val i = (enqIdx % SEGMENT_SIZE).toInt()
    // the regular (fast) path -- if the cell is empty, try to install continuation
    // 通过 CAS 的方式将 Continuation 添加到 Segment 中去
    if (segment.cas(i, null, waiter)) { // installed continuation successfully
        // 添加成功，注册 Continuation 被取消的监听，取消后会通知 Segment
        waiter.invokeOnCancellation(segment, i)
        return true
    }
    // ...
    // CAS 操作失败，在 acquire() 方法中会进行重试。
    return false // broken cell, need to retry on a different cell
}

private fun createSegment(id: Long, prev: SemaphoreSegment?) = SemaphoreSegment(id, prev, 0)

@Suppress("NOTHING_TO_INLINE")
internal inline fun > AtomicRef.findSegmentAndMoveForward(
    id: Long,
    startFrom: S,
    noinline createNewSegment: (id: Long, prev: S) -> S
): SegmentOrClosed {
    while (true) {
        val s = startFrom.findSegmentInternal(id, createNewSegment)
        // 检查查询到的 Segment 状态
        if (s.isClosed || moveForward(s.segment)) return s
    }
}

internal fun > S.findSegmentInternal(
    id: Long,
    createNewSegment: (id: Long, prev: S) -> S
): SegmentOrClosed {
    var cur: S = this
    // 当前的 id 如果小于目标 id 或者当前已经 remove了 执行查找
    while (cur.id > AtomicRef.moveForward(to: S): Boolean = loop { cur ->
    if (cur.id >= to.id) return true
    if (!to.tryIncPointers()) return false
    if (compareAndSet(cur, to)) { // the segment is moved
        if (cur.decPointers()) cur.remove()
        return true
    }
    if (to.decPointers()) to.remove() // undo tryIncPointers
}

override fun unlock(owner: Any?) {
    while (true) {
        // Is this mutex locked?
        // 检查锁状态
        check(isLocked) { "This mutex is not locked" }
        // Read the owner, waiting until it is set in a spin-loop if required.
        // 检查 owner 状态
        val curOwner = this.owner.value
        if (curOwner === NO_OWNER) continue // 

override fun release() {
    while (true) {
        // 获取 _availablePermits 后，再把它的值加 1
        val p = _availablePermits.getAndIncrement()
        if (p >= permits) {
            // Revert the number of available permits
            // back to the correct one and fail with error.
            coerceAvailablePermitsAtMaximum()
            error("The number of released permits cannot be greater than $permits")
        }
        // 如果 p 大于等于 0，表示没有 Continuation 在等待锁，直接返回
        if (p >= 0) return
        // 尝试从等待队列中恢复一个 Continuation
        if (tryResumeNextFromQueue()) return
    }
}

private fun tryResumeNextFromQueue(): Boolean {
    val curHead = this.head.value
    // 从 deqIdx 中获取基础 id
    val deqIdx = deqIdx.getAndIncrement()
    
    // 计算 Segment 的 id
    val id = deqIdx / SEGMENT_SIZE
    val createNewSegment = ::createSegment
    
    // 和插入队列一样，先去查找一个 Segment，这里不同的是从 head 开始查找
    val segment = this.head.findSegmentAndMoveForward(id, startFrom = curHead,
        createNewSegment = createNewSegment).segment // cannot be closed
    segment.cleanPrev()
    // 判断查找到的 Segemnt 的 id 和目标的 id 是否对应，如果不对应就继续查找
    if (segment.id > id) return false
    // 获取对应的 Continuation 在 Segemnt 中对应的位置
    val i = (deqIdx % SEGMENT_SIZE).toInt()
    // 获取到对应的 Continuation
    val cellState = segment.getAndSet(i, PERMIT) // set PERMIT and retrieve the prev cell state
    when {
        cellState === null -> {
            // Acquire has not touched this cell yet, wait until it comes for a bounded time
            // The cell state can only transition from PERMIT to TAKEN by addAcquireToQueue
            repeat(MAX_SPIN_CYCLES) {
                if (segment.get(i) === TAKEN) return true
            }
            // Try to break the slot in order not to wait
            return !segment.cas(i, PERMIT, BROKEN)
        }
        cellState === CANCELLED -> return false // the acquirer has already been cancelled
        
        // 恢复对应的 Continuation
        else -> return cellState.tryResumeAcquire()
    }
}

private fun Any.tryResumeAcquire(): Boolean = when(this) {
    is CancellableContinuation -> {
        this as CancellableContinuation
        // 尝试 Resume
        val token = tryResume(Unit, null, onCancellationRelease)
        if (token != null) {
            // 可以 Resume，通过 completeResume() 方法确认执行下发 Resume
            completeResume(token)
            true
        } else false
    }
    is SelectInstance -> {
        trySelect(this@SemaphoreImpl, Unit)
    }
    else -> error("unexpected: $this")
}

override fun tryResume(value: Unit, idempotent: Any?, onCancellation: ((cause: Throwable) -> Unit)?): Any? {
            // 校验 owner 状态是 NO_OWNER
            assert { this@MutexImpl.owner.value === NO_OWNER }
            // 执行被代理的 Continaution 的 tryResume 方法
            val token = cont.tryResume(value, idempotent) {
            
                // 这个 Lambda 会在 Dispatcher 把任务取消时才执行，也就是表示 resume 失败了。
                assert { this@MutexImpl.owner.value.let { it === NO_OWNER ||it === owner } }
                this@MutexImpl.owner.value = owner
                // 重新解锁
                unlock(owner)
            }
            
            // token 不为空表示 tryResume 成功
            if (token != null) {
                assert { this@MutexImpl.owner.value === NO_OWNER }
                // 将 owner 修改为当前 Continaution 的 owner.
                this@MutexImpl.owner.value = owner
            }
            return token
        }