Make sure to update your StateFlow safely in Kotlin!

Recently a new version of Kotlin Coroutines library was released with a few new extensions functions to help you with StateFlow updates. It all started with this issue:

https://github.com/Kotlin/kotlinx.coroutines/issues/2720

Let’s see what it’s all about… 👇

StateFlow

StateFlow is mostly used for handling state and state updates. It can be used for example in Android’s ViewModel to expose state to your views.

Let’s see an example:

class ExampleViewModel(
    private val initialState: ExampleViewState
) : ViewModel() {
    private val _uiState: MutableStateFlow<ExampleViewState> = MutableStateFlow(initialState)
    val uiState = _uiState.asStateFlow()
}

data class ExampleViewState(
    val title: String = "",
    val description: String = ""
    // other state variables
)

How to update the state?

Well, it’s really easy, since we use a data class. We can just use copy like this:

_uiState.value = _uiState.value.copy(title = "Something")

Simple, right? But you have to be careful…

Where’s the catch?

So what’s the problem? While the code is very simple, there is something you have to be aware of - concurrency.

If between the time copy function completes and the StateFlow’s new value is emitted another thread tries to update the StateFlow — by using copy and updating one of the properties that the current copy isn’t modifying —we could end up with results we were not expecting.

The solution

How do we solve this problem? By using fresh good stuff from Kotlin Coroutines for MutableStateFlow.

I’m talking about are these three methods below:

• update
• updateAndGet
• getAndUpdate

All of these take a function parameter that returns the new state that will be emitted.

Let’s see what’s inside the update method for a moment:

public inline fun <T> MutableStateFlow<T>.update(function: (T) -> T) {
    while (true) {
        val prevValue = value
        val nextValue = function(prevValue)
        if (compareAndSet(prevValue, nextValue)) {
            return
        }    
    }
}

As we can see, a function is passed as a param and it’s applied to the current StateFlow’s value. Then compareAndSet function is used to determine if the value has changed — for example by another thread. If compareAndSet returns false then the while loop will be running until it’s possible to update the state.

So how to use it properly?

_uiState.update { it.copy(title = "Something") }

That’s it! This ensures us that the state can be changed safely 👍