Rudimentary Flow.shareIn with WhenSubscribed behavior until the official operator is launched

FlowShareIn.kt

fun <T> Flow<T>.shareIn(scope: CoroutineScope): Flow<T> {
    val upstream = this
    var broadcastChannel: BroadcastChannel<T>? = null
    val mutex = Mutex()
    var subscriberCount = 0

    return flow {
        val activeChannel = mutex.withLock {
            subscriberCount += 1
            broadcastChannel ?: upstream.broadcastIn(scope).also {
                broadcastChannel = it
            }
        }

        try {
            activeChannel.consumeEach { emit(it) }
        } finally {
            mutex.withLock {
                subscriberCount -= 1

                if (subscriberCount == 0) {
                    broadcastChannel?.cancel()
                    broadcastChannel = null
                }
            }
        }
    }
}