CountDownLatch.kt

package kotlin_playground

import kotlinx.coroutines.coroutineScope
import kotlinx.coroutines.launch
import kotlinx.coroutines.suspendCancellableCoroutine
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock
import kotlin.coroutines.Continuation
import kotlin.coroutines.resume


suspend fun main() {
    val latch = CountDownLatch(50)
    coroutineScope {
        val initial = launch {
            printStamped("Initial await")
            latch.await()
            printStamped("Initial await released!")
        }
        for (i in 1..50) {
            launch {
                printStamped("Count down $i")
                latch.countDown()
            }
        }
        initial.join()
        printStamped("After await")
        latch.await()
        printStamped("After await released!")
    }
}

private fun printStamped(any: Any) {
    println("${System.currentTimeMillis()}: $any")
}

class CountDownLatch(
    initialCount: Int
) {
    private var count = initialCount
    private val mutex = Mutex()
    private val awaiters = mutableSetOf<Continuation<Unit>>()

    suspend fun countDown() {
        if (count == 0) {
            return
        }
        mutex.withLock {
            if (count == 0) {
                return
            }
            count--
            if (count == 0) {
                releaseAll()
            }
        }
    }

    private fun releaseAll() {
        for (awaiter in awaiters) {
            awaiter.resume(Unit)
        }
    }

    suspend fun await() {
        if (count == 0) {
            return
        }
        mutex.withLock {
            if (count == 0) {
                return
            }
        }
        suspendCancellableCoroutine<Unit> {
            awaiters.add(it)
        }
    }
}