在其他執行緒上跑 Flow

Flow 與執行緒切換#

Flow 屬於 coroutine 範疇，自然也希望能享受 coroutine 的一個重要特性：輕鬆切換執行緒。在 coroutine 中我們會用 Dispatchers 搭配 coroutine builder（launch、async），或用 withContext 來切換到合適的執行緒／執行緒池。

那 Flow 是否也能在建立時直接帶入 Dispatcher 呢？看一下 flow 的簽名：

fun <T> flow(block: suspend FlowCollector<T>.() -> Unit): Flow<T>

很可惜，flow { ... } 並沒有 CoroutineContext 參數，所以沒辦法在建立時帶入 Dispatcher。

為什麼不能直接用 withContext#

直覺上會想：那我直接在 flow 內部用 withContext 包住 emit，是否可行？

class Day20 {
    fun flow(): Flow<Int> = flow {
        withContext(Dispatchers.Default) {
            repeat(10) {
                emit(it)
            }
        }
    }
}

fun main() = runBlocking {
    val day20 = Day20()
    day20.flow().collect { println(it) }
    println("done")
}

執行會直接拋出錯誤：Flow invariant is violated。錯誤訊息也會建議改用 flowOn。

Flow 對 emit 的執行緒有嚴格規則：上游 emit 端與下游 collect 端的 context 必須一致，不能在 flow 內部用 withContext 任意切換執行緒。

改用 flowOn#

把 withContext 改成 flowOn：

class Day20 {
    fun flow(): Flow<Int> = flow {
        repeat(10) {
            emit(it)
        }
    }.flowOn(Dispatchers.Default)
}

執行結果：

flowOn 的簽名：

fun <T> Flow<T>.flowOn(context: CoroutineContext): Flow<T>

它接受 CoroutineContext，所以可以丟入 Dispatchers，讓 Flow 切換到對應的執行緒／執行緒池。

flowOn 只影響上游#

flowOn 的另一個重要特性：它只會影響「呼叫它之前」的所有操作，後面的部分不受影響。所以可以連續使用多個 flowOn 來分段切換執行緒：

class Day20 {
    fun flow(): Flow<Int> = flow {
        repeat(10) {
            emit(it)
        }
    }
}

fun main() = runBlocking {
    val day20 = Day20()
    val flow = day20.flow()
    flow.map {
        println("map: ${Thread.currentThread().name}")
        it * 3
    }
    .flowOn(Dispatchers.Default)
    .filter {
        println("filter: ${Thread.currentThread().name}")
        it % 2 == 0
    }
    .flowOn(Dispatchers.IO)
    .collect {
        println("collect: ${Thread.currentThread().name}")
        println(it)
    }
    println("done")
}

輸出（節錄）：

map: DefaultDispatcher-worker-2
...
filter: DefaultDispatcher-worker-1
...
collect: main
0
collect: main
6
...
done

從輸出可以看出：

map 跑在 DefaultDispatcher（受第一個 flowOn(Dispatchers.Default) 影響）。
filter 跑在另一個 DefaultDispatcher（這裡實際上是 IO dispatcher 共用的 worker，仍與下游不同）。
collect 回到 main，這是 runBlocking 外層的 context。

預設使用外層的 Dispatcher#

如果沒有呼叫 flowOn，Flow 預設會使用「呼叫 collect 的那個 coroutine」的 Dispatcher：

class Day20 {
    val scope = CoroutineScope(Job() + Dispatchers.Default)

    fun flow2() = scope.launch {
        flow {
            println("emit: ${Thread.currentThread().name}")
            emit(10)
        }.collect {
            println("collect: ${Thread.currentThread().name} $it")
        }
    }
}

fun main() = runBlocking {
    val day20 = Day20()
    day20.flow2()
    delay(100)
    println("done")
}

emit: DefaultDispatcher-worker-1
collect: DefaultDispatcher-worker-1 10
done

這裡的 scope 帶的 Dispatchers.Default，emit 與 collect 都跑在同一條 worker 上。

launchIn 與 flowOn#

當搭配 launchIn 使用時，Flow 會被啟動在傳入的 CoroutineScope 中，可以再用 flowOn 進一步控制上游的執行緒。

小結#

不能在 flow { ... } 中用 withContext 切換執行緒，會拋出「Flow invariant is violated」。
想切換 Flow 的執行緒就用 flowOn，它接受 CoroutineContext。
flowOn 只影響它「之前」的中間運算子；可以連續多次使用，分段控制不同階段的執行緒。
如果沒有 flowOn，Flow 預設會跟著外層 coroutine 的 Dispatcher 跑。

原文出處#

原書/iThome：https://ithelp.ithome.com.tw/articles/10272159