~Properly handle cancellation in FlowSubscription, handle exceptions from onComplete and onError, make context injector R8-friendly, do not use flowOn for empty context

Author: qwwdfsad

reactive/kotlinx-coroutines-reactive/src/Await.kt

@@ -85,7 +85,7 @@ public suspend fun <T> Publisher<T>.awaitSingle(): T = awaitOne(Mode.SINGLE)
 // ContextInjector service is implemented in `kotlinx-coroutines-reactor` module only.
 // If `kotlinx-coroutines-reactor` module is not included, the list is empty.
 private val contextInjectors: Array<ContextInjector> =
-    ServiceLoader.load(ContextInjector::class.java, ContextInjector::class.java.classLoader).toList().toTypedArray()
+    ServiceLoader.load(ContextInjector::class.java, ContextInjector::class.java.classLoader).iterator().asSequence().toList().toTypedArray() // R8 opto
 
 private fun <T> Publisher<T>.injectCoroutineContext(coroutineContext: CoroutineContext) =
     contextInjectors.fold(this) { pub, contextInjector ->

reactive/kotlinx-coroutines-reactive/src/ContextInjector.kt

@@ -8,7 +8,8 @@ import kotlin.coroutines.CoroutineContext
 @InternalCoroutinesApi
 public interface ContextInjector {
     /**
-     * Injects the coroutine context into the context of the publisher.
+     * Injects `ReactorContext` element from the given context into the `SubscriberContext` of the publisher.
+     * This API used as an indirection layer between `reactive` and `reactor` modules.
      */
     public fun <T> injectCoroutineContext(publisher: Publisher<T>, coroutineContext: CoroutineContext): Publisher<T>
 }

reactive/kotlinx-coroutines-reactive/src/FlowAsPublisher.kt

@@ -7,10 +7,11 @@
 
 package kotlinx.coroutines.reactive
 
+import kotlinx.atomicfu.*
 import kotlinx.coroutines.*
 import kotlinx.coroutines.flow.*
 import org.reactivestreams.*
-import java.util.concurrent.atomic.*
+import java.util.concurrent.atomic.AtomicLong
 import kotlin.coroutines.*
 
 /**
@@ -40,35 +41,40 @@ public class FlowSubscription<T>(
     @Volatile
     private var canceled: Boolean = false
     private val requested = AtomicLong(0L)
-    private val producer: AtomicReference<CancellableContinuation<Unit>?> = AtomicReference()
+    private val producer = atomic<CancellableContinuation<Unit>?>(null)
 
     // This is actually optimizable
     private var job = GlobalScope.launch(Dispatchers.Unconfined, start = CoroutineStart.LAZY) {
         try {
             consumeFlow()
             subscriber.onComplete()
         } catch (e: Throwable) {
-            // Failed with real exception, not due to cancellation
-            if (!coroutineContext[Job]!!.isCancelled) {
-                subscriber.onError(e)
+            try {
+                if (e is CancellationException) {
+                    subscriber.onComplete()
+                } else {
+                    subscriber.onError(e)
+                }
+            } catch (e: Throwable) {
+                // Last ditch report
+                handleCoroutineException(coroutineContext, e)
             }
         }
     }
 
     private suspend fun consumeFlow() {
         flow.collect { value ->
-            if (!coroutineContext.isActive) {
-                subscriber.onComplete()
-                coroutineContext.ensureActive()
-            }
-
+            /*
+             * Flow is scopeless, thus if it's not active, its subscription was cancelled.
+             * No intermediate "child failed, but flow coroutine is not" states are allowed.
+             */
+            coroutineContext.ensureActive()
             if (requested.get() == 0L) {
                 suspendCancellableCoroutine<Unit> {
-                    producer.set(it)
+                    producer.value = it
                     if (requested.get() != 0L) it.resumeSafely()
                 }
             }
-
             requested.decrementAndGet()
             subscriber.onNext(value)
         }
@@ -80,12 +86,9 @@ public class FlowSubscription<T>(
     }
 
     override fun request(n: Long) {
-        if (n <= 0) {
+        if (n <= 0 || canceled) {
             return
         }
-
-        if (canceled) return
-
         job.start()
         var snapshot: Long
         var newValue: Long
@@ -95,7 +98,7 @@ public class FlowSubscription<T>(
             if (newValue <= 0L) newValue = Long.MAX_VALUE
         } while (!requested.compareAndSet(snapshot, newValue))
 
-        val prev = producer.get()
+        val prev = producer.value
         if (prev == null || !producer.compareAndSet(prev, null)) return
         prev.resumeSafely()
     }
@@ -106,4 +109,4 @@ public class FlowSubscription<T>(
             completeResume(token)
         }
     }
-}
+}

reactive/kotlinx-coroutines-reactive/test/FlowAsPublisherTest.kt

@@ -0,0 +1,79 @@
+/*
+ * Copyright 2016-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.coroutines.reactive
+
+import kotlinx.coroutines.*
+import kotlinx.coroutines.flow.*
+import org.junit.Test
+import org.reactivestreams.*
+import kotlin.test.*
+
+class FlowAsPublisherTest : TestBase() {
+
+    @Test
+    fun testErrorOnCancellationIsReported() {
+        expect(1)
+        flow<Int> {
+            emit(2)
+            try {
+                hang { expect(3) }
+            } finally {
+                throw TestException()
+            }
+        }.asPublisher().subscribe(object : Subscriber<Int> {
+            private lateinit var subscription: Subscription
+
+            override fun onComplete() {
+                expectUnreached()
+            }
+
+            override fun onSubscribe(s: Subscription?) {
+                subscription = s!!
+                subscription.request(2)
+            }
+
+            override fun onNext(t: Int) {
+                expect(t)
+                subscription.cancel()
+            }
+
+            override fun onError(t: Throwable?) {
+                assertTrue(t is TestException)
+                expect(4)
+            }
+        })
+        finish(5)
+    }
+
+    @Test
+    fun testCancellationIsNotReported() {
+        expect(1)
+        flow {
+            emit(2)
+            hang { expect(3) }
+        }.asPublisher().subscribe(object : Subscriber<Int> {
+            private lateinit var subscription: Subscription
+
+            override fun onComplete() {
+                expect(4)
+            }
+
+            override fun onSubscribe(s: Subscription?) {
+                subscription = s!!
+                subscription.request(2)
+            }
+
+            override fun onNext(t: Int) {
+                expect(t)
+                subscription.cancel()
+            }
+
+            override fun onError(t: Throwable?) {
+                expectUnreached()
+            }
+        })
+        finish(5)
+    }
+}

reactive/kotlinx-coroutines-reactor/src/FlowAsFlux.kt

@@ -19,11 +19,8 @@ public fun <T: Any> Flow<T>.asFlux(): Flux<T> = FlowAsFlux(this)
 private class FlowAsFlux<T : Any>(private val flow: Flow<T>) : Flux<T>() {
     override fun subscribe(subscriber: CoreSubscriber<in T>?) {
         if (subscriber == null) throw NullPointerException()
-        subscriber.onSubscribe(
-            FlowSubscription(
-                flow.flowOn(subscriber.currentContext().asCoroutineContext()),
-                subscriber
-            )
-        )
+        val hasContext = subscriber.currentContext().isEmpty
+        val source = if (hasContext) flow.flowOn(subscriber.currentContext().asCoroutineContext()) else flow
+        subscriber.onSubscribe(FlowSubscription(source, subscriber))
     }
 }

reactive/kotlinx-coroutines-reactor/src/ReactorContext.kt

@@ -33,7 +33,6 @@ import kotlin.coroutines.*
  *
  * [CoroutineContext] of a suspendable function that awaits a value from [Mono] or [Flux] instance
  * is propagated into [mono] and [flux] Reactor builders:
- *
  * ```
  * launch(Context.of("key", "value").asCoroutineContext()) {
  *   assertEquals(bar().awaitFirst(), "value")
@@ -43,7 +42,6 @@ import kotlin.coroutines.*
  *   coroutineContext[ReactorContext]!!.context.get("key")
  * }
  * ```
-}
  */
 @ExperimentalCoroutinesApi
 public class ReactorContext(val context: Context) : AbstractCoroutineContextElement(ReactorContext) {