cancellation-test.kt

///usr/bin/env jbang "$0" "$@" ; exit $?

//JAVA 17+
//KOTLIN 1.8.20
//DEPS org.jetbrains.kotlinx:kotlinx-coroutines-core:1.7.0-RC
//DEPS io.github.nomisrev:arrow-fx-coroutines-utils:0.1.1-alpha.31

import java.util.concurrent.atomic.AtomicBoolean
import kotlinx.coroutines.*

suspend fun <T> withLock(lock: AtomicBoolean, block: suspend CoroutineScope.() -> T): T {
    // Acquisition is awkward due to:
    // https://github.com/Kotlin/kotlinx.coroutines/issues/3504
    var acquired = false
    try {
        runInterruptible(Dispatchers.IO) {
            while (!lock.compareAndSet(false, true)) {
                Thread.onSpinWait()
                if (Thread.interrupted())
                    throw InterruptedException()
            }
            acquired = true
        }
    } catch (e: Throwable) {
        if (acquired) lock.set(false)
        throw e
    }
    try {
        return coroutineScope { block(this) }
    } finally {
        // The use of `Dispatchers.IO` here is problematic:
        withContext(NonCancellable + Dispatchers.IO) {
            delay(1)
            lock.set(false)
        }
    }
}

/** Concurrency test **/
fun main() = runBlocking {
    val lock = AtomicBoolean(false)
    repeat(10000) { index ->
        val job = launch {
            println("Starting job $index...")
            withLock(lock) {
                delay(1)
                println("Job $index done.")
            }
        }
        if (index % 2 == 0)
            launch {
                println("Cancelling job $index")
                job.cancel()
            }
    }
}