Add StateFlow<T>.onSubscription (#4380)

Fixes #4275

Author: xit0c

kotlinx-coroutines-core/api/kotlinx-coroutines-core.api

@@ -1098,6 +1098,7 @@ public final class kotlinx/coroutines/flow/FlowKt {
 	public static synthetic fun onErrorReturn$default (Lkotlinx/coroutines/flow/Flow;Ljava/lang/Object;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)Lkotlinx/coroutines/flow/Flow;
 	public static final fun onStart (Lkotlinx/coroutines/flow/Flow;Lkotlin/jvm/functions/Function2;)Lkotlinx/coroutines/flow/Flow;
 	public static final fun onSubscription (Lkotlinx/coroutines/flow/SharedFlow;Lkotlin/jvm/functions/Function2;)Lkotlinx/coroutines/flow/SharedFlow;
+	public static final fun onSubscription (Lkotlinx/coroutines/flow/StateFlow;Lkotlin/jvm/functions/Function2;)Lkotlinx/coroutines/flow/StateFlow;
 	public static final fun produceIn (Lkotlinx/coroutines/flow/Flow;Lkotlinx/coroutines/CoroutineScope;)Lkotlinx/coroutines/channels/ReceiveChannel;
 	public static final fun publish (Lkotlinx/coroutines/flow/Flow;)Lkotlinx/coroutines/flow/Flow;
 	public static final fun publish (Lkotlinx/coroutines/flow/Flow;I)Lkotlinx/coroutines/flow/Flow;

kotlinx-coroutines-core/api/kotlinx-coroutines-core.klib.api

@@ -904,6 +904,7 @@ final fun <#A: kotlin/Any?> (kotlinx.coroutines.flow/SharedFlow<#A>).kotlinx.cor
 final fun <#A: kotlin/Any?> (kotlinx.coroutines.flow/SharedFlow<#A>).kotlinx.coroutines.flow/onSubscription(kotlin.coroutines/SuspendFunction1<kotlinx.coroutines.flow/FlowCollector<#A>, kotlin/Unit>): kotlinx.coroutines.flow/SharedFlow<#A> // kotlinx.coroutines.flow/onSubscription|[email protected]<0:0>(kotlin.coroutines.SuspendFunction1<kotlinx.coroutines.flow.FlowCollector<0:0>,kotlin.Unit>){0§<kotlin.Any?>}[0]
 final fun <#A: kotlin/Any?> (kotlinx.coroutines.flow/StateFlow<#A>).kotlinx.coroutines.flow/conflate(): kotlinx.coroutines.flow/Flow<#A> // kotlinx.coroutines.flow/conflate|[email protected]<0:0>(){0§<kotlin.Any?>}[0]
 final fun <#A: kotlin/Any?> (kotlinx.coroutines.flow/StateFlow<#A>).kotlinx.coroutines.flow/distinctUntilChanged(): kotlinx.coroutines.flow/Flow<#A> // kotlinx.coroutines.flow/distinctUntilChanged|[email protected]<0:0>(){0§<kotlin.Any?>}[0]
+final fun <#A: kotlin/Any?> (kotlinx.coroutines.flow/StateFlow<#A>).kotlinx.coroutines.flow/onSubscription(kotlin.coroutines/SuspendFunction1<kotlinx.coroutines.flow/FlowCollector<#A>, kotlin/Unit>): kotlinx.coroutines.flow/StateFlow<#A> // kotlinx.coroutines.flow/onSubscription|[email protected]<0:0>(kotlin.coroutines.SuspendFunction1<kotlinx.coroutines.flow.FlowCollector<0:0>,kotlin.Unit>){0§<kotlin.Any?>}[0]
 final fun <#A: kotlin/Any?> (kotlinx.coroutines.selects/SelectBuilder<#A>).kotlinx.coroutines.selects/onTimeout(kotlin.time/Duration, kotlin.coroutines/SuspendFunction0<#A>) // kotlinx.coroutines.selects/onTimeout|[email protected]<0:0>(kotlin.time.Duration;kotlin.coroutines.SuspendFunction0<0:0>){0§<kotlin.Any?>}[0]
 final fun <#A: kotlin/Any?> (kotlinx.coroutines.selects/SelectBuilder<#A>).kotlinx.coroutines.selects/onTimeout(kotlin/Long, kotlin.coroutines/SuspendFunction0<#A>) // kotlinx.coroutines.selects/onTimeout|[email protected]<0:0>(kotlin.Long;kotlin.coroutines.SuspendFunction0<0:0>){0§<kotlin.Any?>}[0]
 final fun <#A: kotlin/Any?> (kotlinx.coroutines/CompletableDeferred<#A>).kotlinx.coroutines/completeWith(kotlin/Result<#A>): kotlin/Boolean // kotlinx.coroutines/completeWith|[email protected]<0:0>(kotlin.Result<0:0>){0§<kotlin.Any?>}[0]

kotlinx-coroutines-core/common/src/flow/operators/Share.kt

@@ -412,6 +412,29 @@ private class SubscribedSharedFlow<T>(
         sharedFlow.collect(SubscribedFlowCollector(collector, action))
 }
 
+/**
+ * Returns a flow that invokes the given [action] **after** this state flow starts to be collected
+ * (after the subscription is registered).
+ *
+ * The [action] is called before any value is emitted from the upstream
+ * flow to this subscription but after the subscription is established. It is guaranteed that all emissions to
+ * the upstream flow that happen inside or immediately after this `onSubscription` action will be
+ * collected by this subscription.
+ *
+ * The receiver of the [action] is [FlowCollector], so `onSubscription` can emit additional elements.
+ */
+public fun <T> StateFlow<T>.onSubscription(action: suspend FlowCollector<T>.() -> Unit): StateFlow<T> =
+    SubscribedStateFlow(this, action)
+
+@OptIn(ExperimentalForInheritanceCoroutinesApi::class)
+private class SubscribedStateFlow<T>(
+    private val stateFlow: StateFlow<T>,
+    private val action: suspend FlowCollector<T>.() -> Unit
+) : StateFlow<T> by stateFlow {
+    override suspend fun collect(collector: FlowCollector<T>) =
+        stateFlow.collect(SubscribedFlowCollector(collector, action))
+}
+
 internal class SubscribedFlowCollector<T>(
     private val collector: FlowCollector<T>,
     private val action: suspend FlowCollector<T>.() -> Unit

kotlinx-coroutines-core/common/test/flow/sharing/StateFlowTest.kt

@@ -110,6 +110,67 @@ class StateFlowTest : TestBase() {
         }
     }
 
+    @Test
+    fun testOnSubscription() = runTest {
+        expect(1)
+        val state = MutableStateFlow("initial") // "initial" gets lost, replaced by "D"
+        state
+            .onSubscription {
+                emit("collector->A")
+                state.value = "A" // gets lost, replaced by "B"
+            }
+            .onSubscription {
+                emit("collector->B")
+                state.value = "B"
+            }
+            .onStart {
+                emit("collector->C")
+                state.value = "C" // gets lost, replaced by "A"
+            }
+            .onStart {
+                emit("collector->D")
+                state.value = "D" // gets lost, replaced by "C"
+            }
+            .onEach {
+                when (it) {
+                    "collector->D" -> expect(2)
+                    "collector->C" -> expect(3)
+                    "collector->A" -> expect(4)
+                    "collector->B" -> expect(5)
+                    "B" -> {
+                        expect(6)
+                        currentCoroutineContext().cancel()
+                    }
+                    else -> expectUnreached()
+                }
+            }
+            .launchIn(this)
+            .join()
+        finish(7)
+    }
+
+    @Test
+    @Suppress("DEPRECATION") // 'catch'
+    fun testOnSubscriptionThrows() = runTest {
+        expect(1)
+        val state = MutableStateFlow("initial")
+        state
+            .onSubscription {
+                expect(2)
+                throw TestException()
+            }
+            .catch { e ->
+                assertIs<TestException>(e)
+                expect(3)
+            }
+            .collect {
+                // onSubscription throws before "initial" is emitted, so no value is collected
+                expectUnreached()
+            }
+        assertEquals(0, state.subscriptionCount.value)
+        finish(4)
+    }
+
     @Test
     public fun testOnSubscriptionWithException() = runTest {
         expect(1)