WhileSubscribed(5000)

WhileSubscribed(5000)

Trampolining keeps the (blood) flow pumping!

👋 Hi, this is P.Y., I work as an Android Distinguished Engineer at Block. I know nothing about Compose so please let me know if I messed up on Twitter!

I've been hacking on a new LeakCanary standalone app to visualize leaks, which is going to be 100% Compose. As I started following tutorials and looking at sample apps, I noticed a strange pattern in both Now In Android and tivi:

class AuthorViewModel : ViewModel() {

    val authorUiState: StateFlow<AuthorUiState> = authorUiStateStream()
        .stateIn(
            scope = viewModelScope,
            // {-_-} 5000?!
            started = SharingStarted.WhileSubscribed(5_000),
            initialValue = AuthorUiState.Loading
        )
}

@Composable
fun AuthorScreen(authorUiState: AuthorUiState) {
  when (authorUiState) {
    AuthorUiState.Loading -> {
      // ...
    }
  }
}

I wonder what this WhileSubscribed(5_000) is all about! Let's look at the source:

Screenshot of WhileSubscribed sources

Using WhileSubscribed() without any timeout would make sense here, i.e. we want to keep the sharing coroutine running as long as there's a UI consuming it. When that UI goes away, why would we want to wait an additional 5 seconds before we stop sharing?

I find more details in a post from the Android Developers blog:

Tip for Android apps! You can use WhileSubscribed(5000) most of the time to keep the upstream flow active for 5 seconds more after the disappearance of the last collector. That avoids restarting the upstream flow in certain situations such as configuration changes. This tip is especially helpful when upstream flows are expensive to create and when these operators are used in ViewModels.

Surprise Surprise, it's config changes once again, the bane of my Android career...

On Twitter Gabor Varadi pointed out that CoroutineLiveData has the same 5000 ms default timeout.

In my experience, introducing random delays does not properly solve whatever underlying issue I'm running into. I'm also not comfortable with the idea that every ViewModel exposing state as a flow should now have this weird 5000 magic number baked in.

What happens when we remove the delay?

class AuthorViewModel : ViewModel() {

    val authorUiState: StateFlow<AuthorUiState> = authorUiStateStream()
        .stateIn(
            scope = viewModelScope,
            // Timeout is now 0!
            started = SharingStarted.WhileSubscribed(),
            initialValue = AuthorUiState.Loading
        )
}

@Composable
fun AuthorScreen(authorUiState: AuthorUiState)
// ...

When I rotate the screen, the sharing coroutine is stopped & restarted, whereas previously it stayed in started state and didn't restart the upstream flow.

Here's why: as part of a configuration change, the activity is destroyed. Then the activity is recreated, and resumed. Somewhere during that recreation the scope that AuthorScreen used to collect authorUiState completes, which brings the subscription count to 0, and stops the state sharing. And then on the first frame post resume, sharing restarts, which re-triggers the authorUiStateStream() but also immediately reuses the latest cached value (WhileSubscribed.replayExpiration defaults to never).

My next idea was to use SharingStarted.Lazily instead, which starts sharing on subscribe and never stops:

class AuthorViewModel : ViewModel() {

    val authorUiState: StateFlow<AuthorUiState> = authorUiStateStream()
        .stateIn(
            scope = viewModelScope,
            // Start on subscribe and never stop!
            started = SharingStarted.Lazily,
            initialValue = AuthorUiState.Loading
        )
}

@Composable
fun AuthorScreen(authorUiState: AuthorUiState)
// ...

The state is shared with the viewModelScope scope so the sharing will stop as soon as the view model is cleared.

This works, but there's one limitation: the upstream flow stays active for as long as the ViewModel is around, which is usually tied to an Activity or Fragment lifecycle. If a state flow is tied only to a small part of the UI that then goes away and unsubscribes, we'll be keeping that flow active for no good reason.

The same is true of SharingStarted.WhileSubscribed(5_000) of course. That 5_000 timeout is meant as "wait long enough to be sure that if we went through a config change we'd have time to resubscribe before we stop sharing". Unfortunately, this also means that whenever that state is unsubscribed we wait an additional 5 seconds before stopping the upstream flow.

Can we create a SharingStarted that behaves like SharingStarted.WhileSubscribed but will also wait for config changes to settle and for the UI to have a chance to resubscribe before stopping the upstream flow? Let's call it WhileSubscribedOrRetained:

class AuthorViewModel : ViewModel() {

    val authorUiState: StateFlow<AuthorUiState> = authorUiStateStream()
        .stateIn(
            scope = viewModelScope,
            started = WhileSubscribedOrRetained,
            initialValue = AuthorUiState.Loading
        )
}

@Composable
fun AuthorScreen(authorUiState: AuthorUiState)
// ...

Yes we can! I was chatting with Romain Guy and he jokingly suggested:

Romain Guy suggesting post(post(…))

Little did he know... that's exactly what I did, except with even more posting!

The implementation is inspired from WhileSubscribed. Thanks Bill Phillips for suggesting CompletableDeferred:

object WhileSubscribedOrRetained : SharingStarted {

  private val handler = Handler(Looper.getMainLooper())

  override fun command(subscriptionCount: StateFlow<Int>): Flow<SharingCommand> = subscriptionCount
  .transformLatest { count ->
    if (count > 0) {
      emit(SharingCommand.START)
    } else {
      val posted = CompletableDeferred<Unit>()
      // This code is perfect. Do not change a thing.
      Choreographer.getInstance().postFrameCallback {
        handler.postAtFrontOfQueue {
          handler.post {
            posted.complete(Unit)
          }
        }
      }
      posted.await()
      emit(SharingCommand.STOP)
    }
  }
  .dropWhile { it != SharingCommand.START }
  .distinctUntilChanged()

  override fun toString(): String = "WhileSubscribedOrRetained"
}

Wait, what?!

Choreographer.getInstance().postFrameCallback {
  handler.postAtFrontOfQueue {
    handler.post {
      posted.complete(Unit)
    }
  }
}

Ok so this is the fun part: the subscriptionCount updates are dispatched async on the main thread. When a config change occurs, the activity is destroyed, recreated and resumed. As part of the teardown, the subscription count decrement event is posted and runs right after Activity.onResume() but also right before the first frame renders. So we can't stop the subscription right there, as the first composition (where we resubscribe) will happen as part of the first frame.

The resubscription happens during the first frame callback, and the subscription count increment is posted and runs after the first frame.

Last but not least, composition runs during the measure part of a traversal, whereas callbacks manually posted with Choreographer.getInstance().postFrameCallback() run before that, during the animation phase.

So, this is what we do:

  • We receive the decrement to 0 during a post that runs in between Activity.onResume() and the first frame.
  • We schedule a frame callback with Choreographer.getInstance().postFrameCallback {} so that we can be called during the animation part of the first frame.
  • We know that a bit later during that frame callback the resubscription will happen and will trigger a post. We want to run code after that post runs.
  • So we enqueue a post at the front of the main thread queue with handler.postAtFrontOfQueue {}, which runs immediately after the frame callback.
  • And from that we enqueue a post at the back of the main thread queue with handler.post{}, which is guaranteed to run after the subscription count increase notification.
  • When the subscription count increase notification runs, transformLatest ensures that the work to stop (which was suspended with posted.await()) is canceled.
  • If the subscription count doesn't increase, we proceed with stopping the sharing.

Phew, so much trampolining!

The result is nice though: the flow stays active during config changes, and stops immediately when the subscribed UI goes away.

The root cause: state lifecycle

We have state presented on an active screen. We want the lifecycle of that state to be tied to when that screen is visible, and survive config changes.

Our core issue is that we're shoving state in a ViewModel that has a longer lifecycle than what we want the state to have.

The fix is to use fine grained scopes that are tied to the app navigation state ("what am I currently doing?") which does not get torn down on config changes. If you use ViewModels from Jetpack navigation, you already get that, so you can replace WhileSubscribed(5000) with Lazily!

DALL·E 2022-08-30 15.41.41 - An Android jumping on a trampoline, vaporware, low angle.png

Generated by DALL-E, prompt: "An Android jumping on a trampoline, vaporware, low angle"

Â