Is there a way to only run the molecule when its being collected?

[joakimtall]

We have molecule in a view model that is kept alive even when navigating to other tabs, for smooth tab navigation. The compose UI is not shown when going to another tab and will stop collecting, but the molecule in the view model continues to recompose. The molecule is collecting an upstream flow that emits periodically.

The compose UI collectes the molecule state flow with collectAsState, and the molecule is launched on viewModelScope with frame clock

[JakeWharton]

Today you would have to do this with a custom frame clock which wrapped a real one and did not send frames while no one was collecting.

I'm not sure we want to do this generally. One problem is has is that when you start collecting again the initial value is going to be wildly out-of-date which might not be what some want, and probably not something that could be the default. Maybe it could be a flag that you set when creating the StateFlow, but even then someone could read the value without collecting and never see updates.

[joakimtall]

Our workaround for this is to return a StateFlow created using the stateIn function with the WhileSubscribed option. This should have the very same problems that you are describing, right?

Lazy question: Is it straightforward to construct such a custom frame clock?

[JakeWharton]

We have a version of it in the repo: https://github.com/cashapp/molecule/blob/trunk/molecule-runtime/src/commonMain/kotlin/app/cash/molecule/GatedFrameClock.kt

The only difference is that you would want to wrap it around a regular clock since this one ticks as fast as possible (it drives the immediate recomposing clock option of the library).

[joakimtall]

Is it possible to stop collection of an upstream flow this way? I guess that probably is not possible, and would need the whole molecule to be torn down and lose its state -- hope I'm wrong though?

[JakeWharton]

We don't have anything like that, no. That's somewhat out of our control since Compose is the one collecting flows.

[mslalith]

@joakimtall I wrote an extension based on subscriptionCount. This works for my use case. I might've probably missed some corner cases

fun <T : Any> CoroutineScope.launchActiveMolecule(
    clock: RecompositionClock = RecompositionClock.ContextClock,
    body: @Composable () -> T
): StateFlow<T> {
    var mutableStateFlow: MutableStateFlow<T>? = null
    var localBody by mutableStateOf<(@Composable () -> T)?>(value = null)

    launchMolecule(
        clock = clock,
        emitter = { value ->
            val outputFlow = mutableStateFlow
            if (outputFlow != null) {
                outputFlow.value = value
            } else {
                mutableStateFlow = MutableStateFlow(value).flowWhileActive { localBody = if (it) body else null }
            }
        },
        body = {
            localBody?.invoke() ?: mutableStateFlow?.value ?: body()
        }
    )

    return mutableStateFlow!!
}

context (CoroutineScope)
@OptIn(ExperimentalCoroutinesApi::class)
fun <T> MutableStateFlow<T>.flowWhileActive(
    started: SharingStarted = SharingStarted.WhileSubscribed(stopTimeoutMillis = 5000),
    onSubscriptionChanged: (Boolean) -> Unit = {}
): MutableStateFlow<T> {
    val mutableStateFlow = MutableStateFlow(value = value)

    launch {
        started.command(mutableStateFlow.subscriptionCount)
            .mapLatest {
                when (it) {
                    SharingCommand.START -> true
                    SharingCommand.STOP, SharingCommand.STOP_AND_RESET_REPLAY_CACHE -> false
                }
            }
            .distinctUntilChanged()
            .collectLatest { onSubscriptionChanged(it) }
    }

    launch {
        started.command(mutableStateFlow.subscriptionCount)
            .distinctUntilChanged()
            .flatMapLatest {
                when (it) {
                    SharingCommand.START -> this@flowWhileActive
                    SharingCommand.STOP, SharingCommand.STOP_AND_RESET_REPLAY_CACHE -> emptyFlow()
                }
            }
            .collect(mutableStateFlow)
    }

    return mutableStateFlow
}

[joakimtall]

This looks great! Will try it

[joakimtall]

My own adaption of this. Seems to work great. I could not figure out your second block of SharingStarted.command collection. When would that be needed?

Also I think localBody?.invoke() ?: mutableStateFlow?.value ?: body() will break if a molecule actually wants to return null, so tried to fix that

/**
 * This function is a copy of the original launchMolecule StateFlow variant that has been augmented to disable
 * when there are no subscribers according to the SharingStarted strategy. It will simply replace the body of
 * the molecule with a body returning the latest flow value when there are no subscribers.
 *
 * This replacement will cause any upstream flows to stop being collected and is a perfect fit to use with
 * .collectAsStateWithLifecycle from compose UI to stop collection on both lifecycle pause and when the
 * compose UI is removed from composition while the view model is still alive
 *
 * Credit to https://github.com/cashapp/molecule/discussions/225#discussioncomment-7741919
 */
@SuppressLint("StateFlowValueCalledInComposition") // we read flow.value without collect in the body
fun <T> CoroutineScope.launchMolecule(
    started: SharingStarted,
    mode: RecompositionMode,
    body: @Composable () -> T,
): StateFlow<T> {
    var flow: MutableStateFlow<T>? = null
    var bodyOrNothing by mutableStateOf<(@Composable () -> T)?>(null)

    launchMolecule(
        mode = mode,
        emitter = { value ->
            val outputFlow = flow
            if (outputFlow != null) {
                outputFlow.value = value
            } else {
                flow = SubscriberAwareMutableStateFlow(value, started) { hasSubscribers ->
                    bodyOrNothing = if (hasSubscribers) body else null
                }
            }
        },
        /* This is called on each molecule recomposition. If bodyOrNothing is set we invoke it. If not we
         * return the latest flow value. If we dont have a flow to get latest value we are at the first
         * composition and run the body. Subsequent recompositions will always have a flow for the latest
         * value so the last body call will never be run again */
        body = {
            when (val captured = bodyOrNothing) {
                null -> when (flow) { // if bodyOrNothing is null we look if we have a flow yet
                    null -> body() // if not its first run and we run the body for first emit
                    else -> flow!!.value
                }
                else -> captured.invoke()
            }
        },
    )

    return flow!!
}

/**
 * Returns a MutableStateFlow tied to a callback signaling if there are any subscribers according to the
 * SharingStarted strategy
 */
@Suppress("FunctionName")
private fun <T> CoroutineScope.SubscriberAwareMutableStateFlow(
    value: T,
    started: SharingStarted,
    onSubscriptionChanged: (Boolean) -> Unit,
): MutableStateFlow<T> {
    val flow = MutableStateFlow(value = value)

    launch {
        started.command(flow.subscriptionCount)
            .mapLatest {
                when (it) {
                    SharingCommand.START -> true
                    SharingCommand.STOP, SharingCommand.STOP_AND_RESET_REPLAY_CACHE -> false
                }
            }
            .distinctUntilChanged()
            .collectLatest { onSubscriptionChanged(it) }
    }

    return flow
}

[joakimtall]

@JakeWharton What do you think about including something like this in the molecule library? It is a very common scenario for normal compose UI / view model apps

[jingibus]

I'm not a fan of explicitly supporting this if it produces the StateFlow API.

The SharingStarted.WhileSubscribed behavior produces a StateFlow that will yield garbage from the .value API when there are no subscribers: there are no subscribers, so that value is by definition not up-to-date. If we get rid of the .value API, I'm okay with it because Flow doesn't have that same drawback (although I'm not sure it has much utility if we do)