onEach changes the dispatcher in StateFlow (kotlin coroutines)

Viewed 850

Imagine the following self-contained test case

@Test
fun `stateFlow in GlobalScope`() = runBlockingTest {

    suspend fun makeHeavyRequest(): String {
        return "heavy result"
    }

    val flow1 = flowOf(Unit)
        .map { makeHeavyRequest() }
        .onEach { logThread("1: before flowOn") }
        .flowOn(testDispatcher)
        .stateIn(GlobalScope, SharingStarted.Lazily, "init state")

    val flow2 = flowOf(Unit)
        .map { makeHeavyRequest() }
        .onEach { logThread("2: before flowOn") }
        .flowOn(testDispatcher)
        .stateIn(GlobalScope, SharingStarted.Lazily, "init state")
        .onEach { logThread("2: after stateIn") }

    val flow3 = flowOf(Unit)
        .map { makeHeavyRequest() }
        .onEach { logThread("3: before flowOn") }
        .flowOn(testDispatcher)
        .onEach { logThread("3: after flowOn") }
        .stateIn(GlobalScope, SharingStarted.Lazily, "init state")

    flow1.test {
        assertEquals("heavy result", expectItem())
        cancelAndIgnoreRemainingEvents()
    }

    flow2.test {
        assertEquals("heavy result", expectItem())
        cancelAndIgnoreRemainingEvents()
    }

    flow3.test {
        assertEquals("heavy result", expectItem())
        cancelAndIgnoreRemainingEvents()
    }

}

The effect of running it will be:

Thread (1: before flowOn): Thread[main @coroutine#2,5,main]
Thread (2: before flowOn): Thread[main @coroutine#3,5,main]
Thread (2: after stateIn): Thread[main @coroutine#6,5,main]
Thread (3: before flowOn): Thread[DefaultDispatcher-worker-1 @coroutine#8,5,main]
Thread (3: after flowOn): Thread[DefaultDispatcher-worker-1 @coroutine#4,5,main]


org.opentest4j.AssertionFailedError: 
Expected :heavy result
Actual   :init state

In flow3 placing the onEach between flowOn and stateIn completely changes the dispatcher and messes up the result. Why is that?

2 Answers

The reason this happens is because the stateIn operator has some optimizations depending on whether the upstream flow is a ChannelFlow.
.flowOn(...) returns a ChannelFlow while .onEach(...) does not.

Usually this would not really matter. Why it does matter in your case, is because you expect the flow returned by stateIn to never emit the initial value. But there is a reason this parameter is mandatory and you should expect to receive the initial value. Whether you actually do mostly depends on whether the upstream flow is able to emit a value without suspending.

Now it appears that the one of the optimizations in the stateIn operator is, that it might consume a ChannelFlow without suspending. Thats why you get your expected behavior when using

.flowOn(testDispatcher) /*returns ChannelFlow*/
.stateIn(GlobalScope, SharingStarted.Lazily, "init state")
Related