Atomic Updates on MutableStateFlow
The 1.5.1 release of Kotlin Coroutines brought with it interesting extension functions to help with “atomic” updates of values for a mutable state flow.
Some Context
I’m going to describe this problem in the context of Android, since that’s the platform I most commonly work on. But this issue is 100% Kotlin so it could easily apply to other platforms.
StateFlow is commonly used to hold and emit the UI state in the MVVM pattern often used in Android. For example, one might have a ViewModel that exposes a StateFlow of a data class to describe the view state. The view state could be described as a data class.
class MyViewModel : ViewModel() { data class ViewState(
val showLoading: Boolean = false,
val title: String = "Default Title",
val doneButtonEnabled: Boolean = true
) private val _viewState = MutableStateFlow<ViewState>(ViewState())
val viewState = _viewState.asStateFlow()}
An activity or fragment may consume said flow and use the values emitted by it to change the UI state. Note that I’m not going to cover how to safely observe flows within the android lifecycle, that isn’t the focus of this article. Manuel Vivo’s article covers that well. (Shameless plug for my own, older and slightly out of date article covering a similar topic too.)
class MyFragment : Fragment(R.layout.fragment) {
... // Note I'm ignoring details about where the flows are observed
// There are some good reasons to be careful how your flows
// are observed here, but since that isn't the focus of the
// article I'm omitting them.
viewModel.viewState
.onEach {
if (it.showLoading) {
// update the UI to show loading
}
// and so on with the other parts of the view state
}
.launchIn... ...}
Using a data class to describe UI state is a very convenient mechanism when coupled with StateFlow. You can observe the flow repeatedly and always get the most current UI state. The observed value is a trivial data class that can be decomposed into its component properties very simply. (I would note that using a sealed class as your UI state can exhibit the same behaviour described in this article, it’s just more complex to demonstrate.)
Furthermore, updating the state is easy with data classes. Data classes have a convenient copy function that allows you to update one or more properties of the data class while preserving the remaining values. So you can very trivially update the UI state:
// Update just one property, leave the rest as is without
// caring what the values are.
val newUIState = _viewState.value.copy(doneButtonEnabled = true)
_viewState.value = newUIState // emit the new UI stateOr better yet just make it all a one liner:
_viewState.value = _viewState.value.copy(doneButtonEnabled = true)That looks very simple and very straightforward. Whatever other properties are set in the UI state only the doneButtonEnabled property is updated; the remaining properties are preserved.
The Problem
There is, of course, a problem. While the code is all in one line there is a race condition to be aware of. Between the time copy function completes and the state flow’s new value is emitted another thread could have updated the view state and changed one of the properties that the current copy isn’t modifying.
Consider the following very contrived example with two “concurrent” operations:
// Assume the state flow value has the following existing properties:
// title: "Default title"
// doneButtonEnabled: false
// Label: Launch A
viewModelScope.launch(Dispatchers.IO) {
_viewState.value = _viewState.value.copy(doneButtonEnabled = true)
}
// Label: Launch B
viewModelScope.launch(Dispatchers.Default) {
_viewState.value = _viewState.value.copy(title = "New Title")
}If you started observing the state flow after both launch lambdas had completed what would you expect the value to contain?
Many would expect the final emitted data class to contain a title of New Title and a doneButtonEnabled value of true. That might be true for many executions of that code block but it won’t be for all of them given the concurrent nature of those two lambdas.
The reality is the above example is really just the following once we un-roll some of the inlined variables:
// Assume the state flow value has the following existing properties:
// title: "Default title"
// doneButtonEnabled: false
// Label: Launch A
viewModelScope.launch(Dispatchers.IO) {
val currentValueA = _viewState.value
val newValueA = currentValueA.copy(doneButtonEnabled = true)
_viewState.value = newValueA
}
// Label: Launch B
viewModelScope.launch(Dispatchers.Default) {
val currentValueB = _viewState.value
val newValueB = currentValueB.copy(title = "New Title")
_viewState.value = newValueB
}Lets walk through some of the scenarios:
Scenario 1: Serial execution
This scenario is pretty straight-forward. The timing works out so that Launch A runs before Launch B. The Launch A lambda updates the state flow value to:
title: Default Title
doneButtonEnabled: trueThat value is then observed when the Launch B lambda runs and it in turn updates the state flow value to
title: New Title
doneButtonEnabled: trueNo real issues here. The values are consistent, but really only by luck of timing.
Scenario 2: Concurrent execution
This is the scenario where things get interesting. In this scenario both blocks run concurrently. Both see the initial state flow value as
// Initial state flow value
title: Default Title
doneButtonEnabled: falseThe copy function in Launch A copies the initial data class, overwrites the doneButtonEnabled property and then the block emits the following state flow value:
// Data class produced by the Launch A lambda
title: Default Title
doneButtonEnabled: trueAround the same time, the copy function Launch B copies the initial data class that it read before Launch A emitted a value, overwrites the title property and then the block emits the following state flow value:
// Data class produced by the Launch B lambda
title: New Title
doneButtonEnabled: falseNotice the issue? Only one of the two block is going to “win” — whichever emits last. Both blocks read the current UI state and both blocks attempted to update just a portion of the UI state preserving the reset of the data. But that doesn’t happen.
Can this happen without thread hopping? The answer is yes, if the result of the copy is held for any length of time before being emitted on the flow. However, it’s most commonly going to happen when there is multiple concurrent accesses / reads of the state flow to build the copied data class followed by multiple concurrent emissions of the data class on the state flow.
Some Solutions
So how do we fix this?
Mutex
The first possible solution that comes to mind is to surround all updates to the state flow with some sort of synchronization such as a mutex so that it’s impossible for two concurrent state flow value reads and updates to the occur. Implementing this solution the example now becomes:
val mutex = Mutex()// Assume the state flow value has the following existing properties:
// title: "Default title"
// doneButtonEnabled: false
// Label: Launch A
viewModelScope.launch(Dispatchers.IO) {
mutex.withLock {
_viewState.value = _viewState.value.copy(doneButtonEnabled = true)
}
}
// Label: Launch B
viewModelScope.launch(Dispatchers.Default) {
mutex.withLock {
_viewState.value = _viewState.value.copy(title = "New title")
}
}
This isn’t a bad solution. The developer has to remember to maintain the mutex or other synchronization method and the rules that surround it. (For example, a mutex is non re-entrant so watch out for that!)
Compare And Set
The other, and in my opinion, more elegant solution is to use the new extension functions available to MutableStateFlow as of Kotlin Coroutines version 1.5.1 all of which make use of the compareAndSet function.
MutableStateFlow‘s compareAndSet function that hasn’t really been noticed by many developers. On the surface it’s not obvious how it is useful when setting a value. Well the new extension functions help with that.
The new functions are getAndUpdate, update, and updateAndGet
All three extension functions have a function parameter that will be called to generate the next value to be emitted on the state flow. Implementing this solution our example becomes:
// Assume the state flow value has the following existing properties:
// title: "Default title"
// doneButtonEnabled: false
// Label: Launch A
viewModelScope.launch(Dispatchers.IO) {
_viewState.update { it.copy(doneButtonEnabled = true) }
}
// Label: Launch B
viewModelScope.launch(Dispatchers.Default) {
_viewState.update { it.copy(title = "New title") }
}Now that is much cleaner. But, how is it safe for concurrent access? Looking at the source code to the update function we have the following:
public inline fun <T> MutableStateFlow<T>.update(function: (T) -> T) {
while (true) {
val prevValue = value
val nextValue = function(prevValue)
if (compareAndSet(prevValue, nextValue)) {
return
}
}
}We can see that a higher-order function is passed in as a parameter and applied to the previous state flow value to create the new value before an attempt is made to both compare and set it. compareAndSet then checks if the previous value has changed, say by another thread, before setting a new value. If this check fails then the update function just busy loops to try again until the value can actually be set. If another thread changes the previous value while the current thread is executing then the higher-order function used to create the new StateFlow value is repeated.
In this way the copy function of the data class ends up with fresh data and can then safely set the StateFlow value while still preserving the unmodified properties.
Updating a StateFlow’s value can pose some hidden pitfalls, especially when there are concurrent operations going on. Using one of the MutableStateFlow update extension functions can mitigate those issues without the use of verbose and potentially thread blocking mutexs.
However, a mutex may have its place, depending on your use case so don’t entirely ignore it as an option.
