Handle cancelled scope and empty flow in Flow.stateIn (#4327)

Fixes #4322

Author: francescotescari

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

@@ -317,29 +317,33 @@ public fun <T> Flow<T>.stateIn(
  * with multiple downstream subscribers. See the [StateFlow] documentation for the general concepts of state flows.
  *
  * @param scope the coroutine scope in which sharing is started.
+ * @throws NoSuchElementException if the upstream flow does not emit any value.
  */
 public suspend fun <T> Flow<T>.stateIn(scope: CoroutineScope): StateFlow<T> {
     val config = configureSharing(1)
-    val result = CompletableDeferred<StateFlow<T>>()
+    val result = CompletableDeferred<Result<StateFlow<T>>>(scope.coroutineContext[Job])
     scope.launchSharingDeferred(config.context, config.upstream, result)
-    return result.await()
+    return result.await().getOrThrow()
 }
 
 private fun <T> CoroutineScope.launchSharingDeferred(
     context: CoroutineContext,
     upstream: Flow<T>,
-    result: CompletableDeferred<StateFlow<T>>
+    result: CompletableDeferred<Result<StateFlow<T>>>,
 ) {
     launch(context) {
         try {
             var state: MutableStateFlow<T>? = null
             upstream.collect { value ->
                 state?.let { it.value = value } ?: run {
                     state = MutableStateFlow(value).also {
-                        result.complete(ReadonlyStateFlow(it, coroutineContext.job))
+                        result.complete(Result.success(ReadonlyStateFlow(it, coroutineContext.job)))
                     }
                 }
             }
+            if (state == null) {
+                result.complete(Result.failure(NoSuchElementException("Flow is empty")))
+            }
         } catch (e: Throwable) {
             // Notify the waiter that the flow has failed
             result.completeExceptionally(e)

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

@@ -3,6 +3,7 @@ package kotlinx.coroutines.flow
 import kotlinx.coroutines.testing.*
 import kotlinx.coroutines.*
 import kotlinx.coroutines.channels.*
+import kotlin.coroutines.*
 import kotlin.test.*
 
 /**
@@ -88,4 +89,23 @@ class StateInTest : TestBase() {
     fun testSubscriptionByFirstSuspensionInStateFlow() = runTest {
         testSubscriptionByFirstSuspensionInCollect(flowOf(1).stateIn(this@runTest)) { }
     }
+
+    @Test
+    fun testRethrowsCEOnCancelledScope() = runTest {
+        val cancelledScope = CoroutineScope(EmptyCoroutineContext).apply { cancel("CancelMessageToken") }
+        val flow = flowOf(1, 2, 3)
+        assertFailsWith<CancellationException>("CancelMessageToken") {
+            flow.stateIn(cancelledScope)
+        }
+    }
+
+    @Test
+    fun testThrowsNoSuchElementExceptionOnEmptyFlow() = runTest {
+        val flow = emptyFlow<Any>()
+        assertFailsWith<NoSuchElementException> {
+            flow.stateIn(this)
+        }
+        // Ensure that the collecting scope is not cancelled by the NoSuchElementException
+        assertEquals(true, coroutineContext[Job]?.isActive)
+    }
 }