Coroutines utils

KeyedMutex.kt

class KeyedMutex<K : Any> {
	private val locks = ConcurrentHashMap<K, Mutex>()

	suspend fun <T> withLock(key: K, block: suspend () -> T): T {
		val mutex = locks.getOrPut(key) { Mutex() }
		return mutex.withLock { block() }
	}
}

Race.kt

import kotlinx.coroutines.*
import kotlinx.coroutines.selects.select

// 1. Define a scope class to collect the blocks
class RaceScope<T> {
    internal val blocks = mutableListOf<suspend () -> T>()

    // This enables the "onAwait { ... }" syntax inside the block
    fun onAwait(block: suspend () -> T) {
        blocks.add(block)
    }
}

// 2. Define the builder function
suspend fun <T> raceOf(init: RaceScope<T>.() -> Unit): T = coroutineScope {
    val scope = RaceScope<T>().apply(init)
    
    if (scope.blocks.isEmpty()) {
        throw IllegalArgumentException("raceOf requires at least one onAwait block")
    }

    // Launch all collected blocks concurrently
    val deferreds = scope.blocks.map { block ->
        async { block() }
    }

    // Race them
    val fastestResult = select {
        deferreds.forEach { deferred ->
            deferred.onAwait { it }
        }
    }

    // Cancel the losers
    coroutineContext.cancelChildren()

    fastestResult
}

suspend fun requestData1(): String {
    delay(100_000)
    return "Data1"
}

suspend fun requestData2(): String {
    delay(1000)
    return "Data2"
}

suspend fun askMultipleForData(): String {
    return raceOf {
        onAwait {
            requestData1()
        }
        onAwait {
            requestData2()
        }
    }
}

suspend fun main(): Unit = coroutineScope {
    println(askMultipleForData()) // Prints "Data2" after 1 second
}