Kotlin flow sequential asynchronous processing

Viewed 2517

I have a flow (MutableSharedFlow, if it's relevant) and I have potentially expensive operation that I would like to execute asynchronously, while still maintaining the order. I achieved what I wanted using CompletableFuture:

private val threadPoolSize = 5
private val threadPool = Executors.newFixedThreadPool(threadPoolSize)

fun process(flow: Flow<String>) = flow
    .map { CompletableFuture.supplyAsync({ expensiveHandle(it) }, threadPool) }
    .buffer(threadPoolSize)
    .map { it.get() } // block current thread
    .flowOn(threadPool.asCoroutineDispatcher())

Thanks to combination of offloading to thread pool, fixed size buffer and thread blocking CompletableFuture#get, this code works to my expectations - up to threadPoolSize events are processed in parallel, and emitted to the flow in the order they were received.

When I replace CompletableFuture#get with extension function CompletableFuture#await from kotlinx.coroutines.future and use flow or async instead of CompletableFuture#supplyAsync, the messages are no longer processed in parallel:

fun process(flow: Flow<String>) = flow
    .map { 
        runBlocking {
            future { expensiveHandle(it) } // same behaviour with async {...}
        }
    }
    .buffer(threadPoolSize)
    .map { it.await() }
    .flowOn(threadPool.asCoroutineDispatcher())

Can I do equivalent code using coroutines/suspending functions?

3 Answers

async as well as future are extension functions of CoroutineScope. So, you need some CoroutineScope to call them.

runBlocking gives some CoroutineScope, but it's a blocking call, so its usage in suspend functions is prohibited.

You may go with GlobalScope.async, but it's also not recommended and execution would be dispatched by Dispatchers.Default, not by threadPool.asCoroutineDispatcher() as in original example with CompletableFuture.

coroutineScope and withContext functions will provide CoroutineScope, which inherits its coroutineContext from the outer scope, so flow processing will be suspended with immediately executed expensiveHandle(it) coroutine.

You need to create CoroutineScope with factory function, so that coroutines contexts won't mix:

fun process(flow: Flow<String>, threadPool: ThreadPoolExecutor): Flow<String> {
    val dispatcher = threadPool.asCoroutineDispatcher()
    return flow
        .map { CoroutineScope(dispatcher).async { expensiveHandle(it) } }
        .buffer(threadPool.poolSize)
        .map { it.await() }
        .flowOn(dispatcher)
}

Instead of mapping the flow passed as argument, try returning a new flow built with the callbackFlow builder and collect the flow inside, so you can launch several coroutines to call expensiveHandle(it) and send their respective results asap.

fun process(flow: Flow<String>) = callbackFlow {
        flow.collect {
            launch {
                send(expensiveHandle(it))
            }
        }
    }.flowOn(threadPool.asCoroutineDispatcher())

So the problem wasn't the future itself, but the surrounding runBlocking. When using custom CoroutineScope with the thread pool as underlying dispatcher, the code is working as expected (mind the change of get to await, and also I used async instead of future as it's in the core coroutine library):

private val threadPoolSize = 5
private val threadPool = Executors.newFixedThreadPool(threadPoolSize)
private val dispatcher = threadPool.asCoroutineDispatcher()
private val scope = CoroutineScope(dispatcher)

fun process(flow: Flow<String>) = flow
    .map { scope.async(expensiveHandle(it)) }
    .buffer(threadPoolSize)
    .map { it.await() }
    .flowOn(dispatcher)
Related