Making cold Flows lifecycle-aware

With the introduction of SharedFlow and StateFlow, many developers are migrating from LiveData in the UI layer, to take advantage of the goodies of the Flow API, and for a more consistent API accross all layers, but sadly, and as Christophe Beyls explains on his post, the migration is complicated when the view's lifecycle enters into equation.
The version 2.4 of lifecycle:lifecycle-runtime-ktx introduced APIs to help on this side: repeatOnLifecycle and flowWithLifecycle (to learn more about those, check the article: A safer way to collect flows from Android UIs), on this article, we'll try them, and we'll dicuss a minor issue that they introduce in some cases, and we'll see if we can come up with a more flexible solution.

The problem

To explain the problem, let's imagine we have a sample app that listens to location updates when it's active, and whenever a new location is available, it'll make an API call to retrieve some nearby locations.
So for listening to location updates, we'll write a LocationObserver class that offers a cold Flow returning them

class LocationObserver(private val context: Context) {
    fun observeLocationUpdates(): Flow<Location> {
        return callbackFlow {
            Log.d(TAG, "observing location updates")

            val client = LocationServices.getFusedLocationProviderClient(context)
            val locationRequest = LocationRequest
                .create()
                .setPriority(LocationRequest.PRIORITY_HIGH_ACCURACY)
                .setInterval(0)
                .setFastestInterval(0)

            val locationCallback = object : LocationCallback() {
                override fun onLocationResult(locationResult: LocationResult?) {
                    if (locationResult != null) {
                        Log.d(TAG, "got location ${locationResult.lastLocation}")
                        trySend(locationResult.lastLocation)
                    }
                }
            }

            client.requestLocationUpdates(
                locationRequest,
                locationCallback,
                Looper.getMainLooper()
            )

            awaitClose {
               Log.d(TAG, "stop observing location updates")
               client.removeLocationUpdates(locationCallback)
            }
        }
    }
}

then we'll use this class in our ViewModel

class MainViewModel(application: Application) : AndroidViewModel(application) {
    private val locationObserver = LocationObserver(application)

    private val hasLocationPermission = MutableStateFlow(false)

private val locationUpdates: Flow<Location> = hasLocationPermission
        .filter { it }
        .flatMapLatest { locationObserver.observeLocationUpdates() }

val viewState: Flow<ViewState> = locationUpdates
        .mapLatest { location ->
            val nearbyLocations = api.fetchNearbyLocations(location.latitude, location.longitude)
            ViewState(
                isLoading = false,
                location = location,
                nearbyLocations = nearbyLocations
            )
        }

    fun onLocationPermissionGranted() {
        hasLocationPermission.value = true
    }
}

For the sake of simplicity, we are using an AndroidViewModel to have access to the Context directly, and we won't handle different edge cases about location permissions and settings.

Now, all we have to do in our Fragment, is to listen to the react to the viewState updates, and update the UI:

viewLifecycleOwner.lifecycleScope.launchWhenStarted {
   viewModel.viewState
       .onEach { viewState ->
            binding.render(viewState)
        }
       .launchIn(this)
}

where the FragmentMainBinding#render is an extension that can update the UI.

Now if we try to run the app, when we put it to the background, we'll see that the LocationObserver is still listening to location updates, then fetching the nearby places, even though the UI is ignoring them.

Our first attempt to solve this, is to use the new API flowWithLifecycle

viewLifecycleOwner.lifecycleScope.launchWhenStarted {
   viewModel.viewState
       .flowWithLifecycle(viewLifecycleOwner.lifecycle)
       .onEach { viewState ->
            binding.render(viewState)
        }
       .launchIn(this)
}

If we run the app, now, we'll notice that it prints the following line to Logcat each time it goes to background

D/LocationObserver: stop observing location updates

So the new APIs fix the issue, but there is an issue, whenever the app goes to background then we come back, we lose the data we had before, and we hit the API another time even if the location hasn't changed, this occurs because flowWithLifecycle will cancel the upstream each time the used lifecycle goes below the passed State (which is Started for us) and restart it again when the state is restaured.

Solution using the official APIs

The official solution while keeping using flowWithLifecycle is explained in Jose Alcérreca's article, and it's to use stateIn but with a special timeout set to 5 seconds to account for configuration changes, so we need to add the following statement to our viewState's Flow to this

stateIn(
            scope = viewModelScope,
            started = SharingStarted.WhileSubscribed(stopTimeoutMillis = 5000L),
            initialValue = ViewState(isLoading = true)
        )

This works well, except, the stopping/restarting of the Flow each time the app goes to background creates another issue, let's say for example that we don't need to fetch the nearby places unless the location has changed by a minimum distance, so let's change our code to the following

val viewState: Flow<ViewState> = locationUpdates
        .distinctUntilChanged { l1, l2 -> l1.distanceTo(l2) <= 300 }
        .mapLatest { location ->
            val nearbyLocations = api.fetchNearbyLocations(location.latitude, location.longitude)
            ViewState(
                isLoading = false,
                location = location,
                nearbyLocations = nearbyLocations
            )
        }
        .stateIn(
            scope = viewModelScope,
            started = SharingStarted.WhileSubscribed(stopTimeoutMillis = 5000L),
            initialValue = ViewState(isLoading = true)
        )

If we run the app now, then we put it to background for longer than 5 seconds, and re-open it, we will notice that the we re-fetch the nearby locations even if the location didn't change at all, and while this is not a big issue for most cases, it can be costly on some situations: slow network, or slow APIs, or heavy calculations...

An alternative solution: making the Flows lifecycle-aware

What if we could make our locationUpdates flow lifecycle-aware, to stop it without any explicit interaction from the Fragment? This way, we will be able to stop listening to location updates, without having to restart the whole Flow, and re-run all the intermediate operators if the location didn't change, and we could even collect our viewState Flow regularly using launchWhenStarted, since we will be sure it won't run as we are not emitting any locations.

If only we can have an internal hot flow inside our ViewModel that let's us observe the View's state:

private val lifeCycleState = MutableSharedFlow<Lifecycle.State>(replay = 1, onBufferOverflow = BufferOverflow.DROP_OLDEST)

Then we would be able to have an extension that stops then restarts our upstream Flow depending on the lifecycle:

fun <T> Flow<T>.whenAtLeast(requiredState: Lifecycle.State): Flow<T> {
    return lifeCycleState.map { state -> state.isAtLeast(requiredState) }
            .distinctUntilChanged()
            .flatMapLatest {
                // flatMapLatest will take care of cancelling the upstream Flow
                if (it) this else emptyFlow()
            }
}

Actually, we can implement this using LifecycleEventObserver API

private val lifecycleObserver = object : LifecycleEventObserver {
    override fun onStateChanged(source: LifecycleOwner, event: Lifecycle.Event) {
        lifeCycleState.tryEmit(event.targetState)
        if (event.targetState == Lifecycle.State.DESTROYED) {
            source.lifecycle.removeObserver(this)
        }
    }
}

Which we can use to hook up to the Fragment's lifecycle events:

fun startObservingLifecycle(lifecycle: Lifecycle) {
    lifecycle.addObserver(lifecycleObserver)
}

Having this, we can now update our locationUpdates Flow to the following

private val locationUpdates: Flow<Location> = hasLocationPermission
    .filter { it }
    .flatMapLatest { locationObserver.observeLocationUpdates() }
    .whenAtLeast(Lifecycle.State.STARTED)

And we get to observe our viewState Flow regularly in the Fragment, without worrying about keeping the GPS on when the app goes to the background

viewLifecycleOwner.lifecycleScope.launchWhenStarted {
    viewModel.viewState
        .onEach { viewState ->
            binding.render(viewState)
        }
        .launchIn(this)
}

The extension whenAtLeast is flexible in the sense that it can be applied to any Flow in the chain, and not only during the collection, and as we saw, applying it to the upstream triggering Flow (location updates in our case), resulted in less calculations:

• The intermediate operators including the nearby places fetching doesn't run unless needed.
• We won't re-emit the result to the UI on coming back from background since we won't cancel the collection.

If you want to check the full code on Github: https://github.com/hichamboushaba/FlowLifecycle, and the full code contains a sample on how we can unit test our ViewModels with those changes in place.

Conclusion

As you can see, using kotlin Flows on the UI layer is still not always straightforward, but still, I prefer it over LiveData, to have access to all its APIs, and for a consistent API across all layers, for the issue we discussed here, I personally just put the explained logic in a BaseFragment/BaseViewModel, and then I can use it in all the screens without any boilerplate, and it worked well for me on my personal app since 2018.
Let's hope that this feature request gets implemented to offer a way to pause the collection cooperatively without cancelling it, which would fix the issue completely.

What do you think about the solutions explained on this article? And did I miss a simpler solution for this issue? Let me know on the comments.