kotlin-uncancelable-sample.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 io.github.nomisrev.uncancellable
import java.util.concurrent.atomic.AtomicBoolean
import kotlinx.coroutines.*

/**
 * Uses `uncancelable` defined at:
 * [link](https://github.com/nomisRev/arrow-fx-coroutines-utils/blob/main/src/commonMain/kotlin/io/github/nomisrev/UncancellableRegion.kt)
 */
suspend fun <T> withLock(lock: AtomicBoolean, block: suspend CoroutineScope.() -> T): T =
    uncancellable {
        // AKA `poll` from Cats-Effect
        cancellable {
            // 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@uncancellable cancellable {
                block(this)
            }
        } finally {
            // No need for suspended execution here, but might as well
            // prove this finalizer is uncancelable
            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")
                try {
                    job.cancelAndJoin()
                } catch (_: CancellationException) {}
            }
    }
}