Can I build a Kotlin SharedFlow where the consumer dictates replay length?

Viewed 496

Question
When instantiatiang a Kotlin MutableSharedFlow<T> class it allows you to specify replay length of n >= 0. All consumers will get n number of events replayed. Is it a good way to extend or wrap MutableSharedFlow so that the consumer dictactes how many (if any) events he/she wants replayed?

Example desired consumer code

flow.collectWithReplay(count = 1) { event -> ... }

Count would of course have to be equal or less than the upper boundary decided by the flow instance.

Rationale
Some times you want to act differently upon events that are old and new. An example is when the event contains one-time information that is irellevant after consumed once (e.g. data for an error dialog). You may still want to know that the last state was an error, but since it is old you don't show a dialog again. You'd then call flow.replayCache.lastOrNull() to get the old and then subscribe to new using .collectWitReplay(0).

Other times you don't want that distinction and then it would be a hassle to do the two calls separately. .collectWithReplay(1) then yields less and prettier code.

Solution attempted
I have made a solution using my own 1-element replay cache, which solves a special case for n=1. It would be trivial to extend to any n - that's not the point, but I dislike a couple of things about it:

a) It doesn't utilize the built in replay mechanism of SharedFlow
b) It's not thread-safe. collectWithReplay might lose an event emitted in between its line 1 and 2
c) Not sure if I lose any performance by losing inline on the collect method signature

open class FlowEventBus<T>() {

    private val _flow = MutableSharedFlow<T>(replay = 0)

    var latest: T? = null
        private set

    suspend fun emit(event: T) {
        latest = event
       _flow.emit(event)  // suspends until all subscribers receive the event
    }

    /** Consumers who only wants events occuring from now on subscribe here */
    suspend fun collect(action: suspend (value: T) -> Unit) = _flow.collect(action)

    /** Consumers who wants the last event emitted as well as future events subscribe here */
    suspend fun collectWithReplay(action: suspend (value: T) -> Unit) {
        latest?.let { action(it) } // Replay any cached event
        _flow.collect(action)      // Listen for new events
    }
}
1 Answers

Answer to main question
Here the foundation for a solution based on the suggestion from @tenfour04

val mainFlow = MutableSharedFlow<String>(10)

If consumers want a different replay value, the do this:

val flowForTwo = mainFlow.shareIn(threadPoolScope, SharingStarted.Eagerly, 2)
flowForTwo.collect { }

You'll be creating a new SharedFlow each time you do this though, so performance may suffer.

See working test

Variation: Event bus with zero or one replay

Here is a solution where the flow is wrapped in an event bus and the consumer may decide between replay length of 0 or 1. This solution comes with some race condition quirks when you emit and collect very close in time. Run and understand this failing unit before using in production. I don't know how to fix it, or if it's worth fixing. You might be better off just using a variation of my original idea.

/**
 * FlowEventBus where consumer can decide between single replay or no replay when collecting.
 * Warning: It has some concurrency issues that is apparent when you run the tests
 */
class FlowEventBus<T> {
    private val threadPoolScope = CoroutineScope(Dispatchers.Default + SupervisorJob())

    private val eventsWithSingleReplay = MutableSharedFlow<T>(replay = 1) // private mutable shared flow
    private val eventsWithoutReplay = eventsWithSingleReplay.shareIn(threadPoolScope, SharingStarted.Eagerly, replay = 0)

    val latest: T?
        get() = eventsWithSingleReplay.replayCache.lastOrNull()

    /** Emit a new event */
    suspend fun emit(event: T) = eventsWithSingleReplay.emit(event)

    /** Consumers who only wants events occuring from now on subscribe here */
    suspend fun collect(action: suspend (value: T) -> Unit) = eventsWithoutReplay.collect(action)

    /** Consumers who wants the last event emitted as well as future events subscribe here */
    suspend fun collectWithReplay(action: suspend (value: T) -> Unit) {
        eventsWithSingleReplay.collect(action)
    }
} 
Related