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

Question

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
}
}

Answer 1

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)
}
}