Block of kotlin code behaves differently when it passed as lambda to another function

Question

import kotlinx.coroutines.*

fun main() = runBlocking {
    println("Main program starts: ${Thread.currentThread().name}")
    withContext(Dispatchers.IO.limitedParallelism(1)) {
        println("Switched to limited parallelism context: ${Thread.currentThread().name}")
        repeat(2) { i ->
            launch {
                println("Coroutine $i starts on thread ${Thread.currentThread().name}")
                delay(1000L)
            }
        }
    }
    println("Main program ends: ${Thread.currentThread().name}")
}

The above Kotlin code prints the below result

Main program starts: main
Switched to limited parallelism context: DefaultDispatcher-worker-1
Coroutine 0 starts on thread DefaultDispatcher-worker-1
Coroutine 1 starts on thread DefaultDispatcher-worker-1
Main program ends: main

But if I write a custom function for withContext(Dispatchers.IO.limitedParallelism(1)) then the behaviour of the code inside withContext changes.

import kotlinx.coroutines.*

fun main() = runBlocking {
    println("Main program starts: ${Thread.currentThread().name}")
    atomically {
        println("Switched to limited parallelism context: ${Thread.currentThread().name}")
        repeat(2) { i ->
            launch {
                println("Coroutine $i starts on thread ${Thread.currentThread().name}")
                delay(1000L)
            }
        }
    }
    println("Main program ends: ${Thread.currentThread().name}")
}

internal suspend inline fun <R> atomically(crossinline f: () -> R): R =
    withContext(Dispatchers.IO.limitedParallelism(1)) { f() }

The above Kotlin code prints the below result

Main program starts: main
Switched to limited parallelism context: DefaultDispatcher-worker-1
Coroutine 0 starts on thread main
Coroutine 1 starts on thread main
Main program ends: main

As you can see, the two-child coroutine inside the withContext is not started on the DefaultDispatcher instead it is running on the main. I would like to understand why this behaviour change happens.

Answer 1

Your problem comes from the fact that the function parameter f in your extracted function doesn't have a CoroutineScope receiver. Nested coroutine launches will therefore not be done on the child scope of withContext but on the coroutine scope one level up.

To fix it, just add a CoroutineScope receiver for f:

internal suspend inline fun <R> atomically(crossinline f: CoroutineScope.() -> R): R =
    withContext(Dispatchers.IO.limitedParallelism(1)) { f() }