The ViewModel’s leaked Flow collectors problem

I love Kotlin Flow, specially when used to chain the data from your data layer (or use cases if you got them) to your Ui. When done right, complex screens that update themselves when data changes become surprisingly easy to understand and maintain, but as anything in our beloved Kotlin, this power comes with its own set of not so obvious risks that you need to understand. In this article we’ll go through one of them, one that I see from time to time in the Pull Request from the team.

This time let’s start with a bold, maybe polemic statement:

Don’t use Flow.collect() in your ViewModels.

Understanding the issue with collect in the ViewModel

All right, that statement needs a lot of clarification. There are scenarios were collect() won’t imply a risk, but my personal take when reviewing pull request is to inspect every collect() I see in ViewModels and, most of the time, there is a problem behind it. Lets check some code.

ViewModels observe data from repositories (or use cases) and map it to the UiLayer. A common way to do it is this one:

class MyVeryBasicViewModel(private val repository: Repository): ViewModel {
  
  private val _uiState = MutableStateFlow<UiState>(UiState("initial state"))
  
  // Expose UiState to fragment
  val uiState = _uiState

  init{
    viewModelScope.launch {
      repository.getDataFlow().collect { something -> 
        _uiState.emit(something.mapToUiState())
      }
    }
  }
}

data class UiState(val text: String)

This ViewModel defines a mutable StateFlow that will be used to post the data coming from the repo.

At its init method, the ViewModel collects the data from the repository and emits it again to the MutableStateFlow.

Do you spot any issue? Because there is none, or at least, no biggie. There are minor caveats though:

• Collection starts as soon as the ViewModel is created, but the Flow might be never consumed by the Ui. For most cases you don’t want to trigger request until there is someone collecting the StateFlow.
• Having a MutableStateFlow declared in the ViewModel means that anyone from anywhere in the ViewModel will be able to emit something. If the ViewModel grows in complexity it might get hard to follow the flow of execution and debug issues.

These two points a mere warnings but check what happen if we add a little more complexity. Let’s say that the Ui has a refresh button. It might be a pull to refresh or just a retry button displayed if a request failed.

class MyStandardViewModel(private val repository: Repository): ViewModel {
  
  private val _uiState = MutableStateFlow<UiState>(UiState("initial state"))
  val uiState = _uiState

  init {
    refresh()
  }

  /**
  * Request data again. Can be called from the outside.
  */
  public fun refresh(){
    viewModelScope.launch {
      repository.getDataFlow().collect { something -> 
        _uiState.emit(something.mapToUiState())
      }
    }
  }
}

data class UiState(val text: String)

Collecting the Flow from the repo and emitting in the _uiState has been moved a separated method, which is public so it can be called from Ui to refresh the data.

It might not look like it, but this is very wrong. Every time refresh() is called, a new collector is created and will stay alive until the ViewModel is destroyed. So, imagine a pull to refresh. Every time the user pulls, a new collector will be created. If there are 10 pulls to refresh, there will be 10 collectors emitting to _uiState. This is the problem of the leaked collectors.

Wait a minute? Every time refresh is called a collector is leaked? Well, not exactly. It depends on the kind of Flow we are collecting:

• With a Flow that emits a finite number of values and then finishes, there won’t be any problem. The Flow will end at some point and the collector will be garbage collected. Still, a Flow that does a lot of emissions might also take its time and leak the collector for the time being.
• When the Flow is a hot Flow, like one that reacts to changes in Room or SharedPreferences, then this monster is very real. A hot Flow never ends, so the collector is never released.

Even if it depends on the type of Flow, we should take into account that, from the ViewModel perspective, we don’t know which kind of Flow (hot or cold) is exposed from the underneath layers, and even if you know (because the code is yours and you checked) there is no guarantee that one day won’t change, so a properly coded ViewModel should only count with the information exposed to it and be resilient.

A big group of leaked collectors overpopulating a ViewModel (Photo by Martin Wettstein on Unsplash)

How to fix it

We already spoiled the conclusion, do not use collect in your ViewModel, but… how? Talking about the two scenarios above, let’s see what can be done, which is solely based in Flow and it’s operators.

Very basic scenario

class MyVeryBasicViewModel(private val repository: Repository): ViewModel {
  
  // Expose UiState to fragment
  val uiState = repository.getDataFlow()
    .stateIn(viewModelScope, SharingStarted.WhileSubscribed(), UiState())
}

data class UiState(val text: String? = null)

Instead of collecting to later emit in another Flow, stateIn() operator can be used to convert the Flow from the repository into a StateFlow in one lovely line. It also improves in the two mentioned caveats:

• Repository is only called when the UI start collecting uiState. If you want to start collecting in advance, SharingStarted.Eagerly will do the trick.
• There is no mutable StateFlow around.

Refresh from the UI scenario

For the second scenario, the one with the feature to refresh the content from the UI, we can do the the following:

class MyStandardViewModel(private val repository: Repository): ViewModel {
  
  // Emit here for refreshing the Ui
  private val trigger = MutableSharedFlow<Unit>(replay = 1)
  
  // UiState is reclculated with every trigger emission
  val uiState = trigger.flatMapLatest { _->
    repository.getDataFlow()
      .map{ it.mapToUiState() }
  }.stateIn(viewModelScope, SharingStarted.WhileSubscribed(), UiState())

  init {
    refresh()
  }

  /**
  * Request data again. Can be called from the outside.
  */
  public fun refresh(){
    viewModelScope.launch {
      trigger.emit(Unit)
    }
  }
}

data class UiState(val text: String)

In this case, the mechanism to re-trigger the call to our repository is a MutableSharedFlow that we keep private. Thanks to flatMapLatest(), when something is emitted in the trigger the content of the lambda is executed, in this case returning a new Flow from our repository. We are transforming this Flow into an StateFlow with stateIn() like in the previous example.

Our refresh method is only responsible from emitting in the trigger.
Note that we are using SharedFlow for the trigger instead of StateFlow. This is because StateFlow does skip the emission of the same value than before (like with Flow operator .distingUntilChanged()), but with SharedFlow that is not a problem and will emit every time there is something new.

Let’s evaluate the pros of this solution:

• Like the solution to first scenario, repository is only called when the UI starts collecting.
• We still have a mutable Flow declared in the ViewModel, but it is an ultra scoped one.
• There is no collect() in the ViewModel and the leaked collectors problem has gone for good.

Conclusion

Now you are aware of the leaked collectors problem and know a way to solve it, one which is solely based in Flow and it’s wonderful range of operators. Even if the scenario gets more tricky than this one, there will be a combination of Flow operators that will help you to avoid collecting and emitting again.

Thanks for reading, I hope that you found it useful. Have a nice Flow!

Sample code can be found in the following repo:
https://github.com/juanmeanwhile/LeakedFlows

The views, thoughts,and opinions expressed in the text belong solely to the author, and do not represent the opinion, strategy or goals of the author’s employer, organization, committee or any other group or individual.