Fixed exception handling for publish builder

* Test that ensures proper parent cancellation
* Stress test to ensure that exception in publisher always handled
  by parent
* Also fixed race between request & signal completion that did not
  reproduce before

Author: elizarov

binary-compatibility-validator/reference-public-api/kotlinx-coroutines-core.txt

@@ -144,6 +144,7 @@ public final class kotlinx/coroutines/CoroutineExceptionHandlerKt {
 	public static final fun CoroutineExceptionHandler (Lkotlin/jvm/functions/Function2;)Lkotlinx/coroutines/CoroutineExceptionHandler;
 	public static final fun handleCoroutineException (Lkotlin/coroutines/CoroutineContext;Ljava/lang/Throwable;Lkotlinx/coroutines/Job;)V
 	public static synthetic fun handleCoroutineException$default (Lkotlin/coroutines/CoroutineContext;Ljava/lang/Throwable;Lkotlinx/coroutines/Job;ILjava/lang/Object;)V
+	public static final fun handleExceptionViaHandler (Lkotlin/coroutines/CoroutineContext;Ljava/lang/Throwable;)V
 }
 
 public final class kotlinx/coroutines/CoroutineName : kotlin/coroutines/AbstractCoroutineContextElement {

common/kotlinx-coroutines-core-common/src/CoroutineExceptionHandler.kt

@@ -31,7 +31,11 @@ public fun handleCoroutineException(context: CoroutineContext, exception: Throwa
     handleExceptionViaHandler(context, exception)
 }
 
-internal fun handleExceptionViaHandler(context: CoroutineContext, exception: Throwable) {
+/**
+ * @suppress This is an internal API and it is subject to change.
+ */
+@InternalCoroutinesApi
+public fun handleExceptionViaHandler(context: CoroutineContext, exception: Throwable) {
     // Invoke exception handler from the context if present
     try {
         context[CoroutineExceptionHandler]?.let {

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

@@ -59,6 +59,9 @@ private class PublisherCoroutine<in T>(
     private val subscriber: Subscriber<T>
 ) : AbstractCoroutine<Unit>(parentContext, true), ProducerScope<T>, Subscription, SelectClause2<T, SendChannel<T>> {
     override val channel: SendChannel<T> get() = this
+
+    // cancelsParent == true ensure that error is always reported to the parent, so that parent cannot complete
+    // without receiving reported error.
     override val cancelsParent: Boolean get() = true
 
     // Mutex is locked when either nRequested == 0 or while subscriber.onXXX is being invoked
@@ -69,6 +72,8 @@ private class PublisherCoroutine<in T>(
     @Volatile
     private var cancelled = false // true when Subscription.cancel() is invoked
 
+    private var handleException = false // when handleJobException is invoked
+
     override val isClosedForSend: Boolean get() = isCompleted
     override val isFull: Boolean = mutex.isLocked
     override fun close(cause: Throwable?): Boolean = cancel(cause)
@@ -105,68 +110,79 @@ private class PublisherCoroutine<in T>(
         }
     }
 
+    /*
+       This code is not trivial because of the two properties:
+       1. It ensures conformance to the reactive specification that mandates that onXXX invocations should not
+          be concurrent. It uses Mutex to protect all onXXX invocation and ensure conformance even when multiple
+          coroutines are invoking `send` function.
+       2. Normally, `onComplete/onError` notification is sent only when coroutine and all its children are complete.
+          However, nothing prevents `publish` coroutine from leaking reference to it send channel to some
+          globally-scoped coroutine that is invoking `send` outside of this context. Without extra precaution this may
+          lead to `onNext` that is concurrent with `onComplete/onError`, so that is why signalling for
+          `onComplete/onError` is also done under the same mutex.
+     */
+
     // assert: mutex.isLocked()
     private fun doLockedNext(elem: T) {
-        // check if already closed for send
+        // check if already closed for send, note, that isActive become false as soon as cancel() is invoked,
+        // because the job is cancelled, so this check also ensure conformance to the reactive specification's
+        // requirement that after cancellation requested we don't call onXXX
         if (!isActive) {
-            doLockedSignalCompleted()
+            unlockAndCheckCompleted()
             throw getCancellationException()
         }
         // notify subscriber
-        try {
-            subscriber.onNext(elem)
-        } catch (e: Throwable) {
-            try {
-                if (!cancel(e))
-                    handleCoroutineException(context, e, this)
-            } finally {
-                doLockedSignalCompleted()
-            }
-            throw getCancellationException()
-        }
+        subscriber.onNext(elem)
         // now update nRequested
         while (true) { // lock-free loop on nRequested
             val cur = _nRequested.value
             if (cur < 0) break // closed from inside onNext => unlock
             if (cur == Long.MAX_VALUE) break // no back-pressure => unlock
             val upd = cur - 1
             if (_nRequested.compareAndSet(cur, upd)) {
-                if (upd == 0L) return // return to keep locked due to back-pressure
+                if (upd == 0L) {
+                    // return to keep locked due to back-pressure
+                    return
+                }
                 break // unlock if upd > 0
             }
         }
-        /*
-           There is no sense to check for `isActive` before doing `unlock`, because cancellation/completion might
-           happen after this check and before `unlock` (see `onCancellation` that does not do anything
-           if it fails to acquire the lock that we are still holding).
-           We have to recheck `isActive` after `unlock` anyway.
-         */
+        unlockAndCheckCompleted()
+    }
+
+    private fun unlockAndCheckCompleted() {
+       /*
+         There is no sense to check completion before doing `unlock`, because completion might
+         happen after this check and before `unlock` (see `signalCompleted` that does not do anything
+         if it fails to acquire the lock that we are still holding).
+         We have to recheck `isCompleted` after `unlock` anyway.
+        */
         mutex.unlock()
-        // recheck isActive
-        if (!isActive && mutex.tryLock())
-            doLockedSignalCompleted()
+        // check isCompleted and and try to regain lock to signal completion
+        if (isCompleted && mutex.tryLock()) doLockedSignalCompleted()
     }
 
-    // assert: mutex.isLocked()
+    // assert: mutex.isLocked() & isCompleted
     private fun doLockedSignalCompleted() {
         try {
             if (_nRequested.value >= CLOSED) {
                 _nRequested.value = SIGNALLED // we'll signal onError/onCompleted (that the final state -- no CAS needed)
                 val cause = getCompletionCause()
                 // Specification requires that after cancellation requested we don't call onXXX
                 if (cancelled) {
-                    // but we cannot just ignore exception so we handle it
-                    if (cause != null) handleCoroutineException(context, cause, this)
-                    return
-                }
-                try {
-                    if (cause != null && cause !is CancellationException)
-                        subscriber.onError(cause)
-                    else
-                        subscriber.onComplete()
-                } catch (e: Throwable) {
-                    handleCoroutineException(context, e, this)
-                }
+                    // If the parent had failed to handle our exception (handleJobException was invoked), then
+                    // we must not loose this exception
+                    if (handleException && cause != null) handleExceptionViaHandler(parentContext, cause)
+                } else {
+                    try {
+                        if (cause != null && cause !is CancellationException)
+                            subscriber.onError(cause)
+                        else
+                            subscriber.onComplete()
+                    } catch (e: Throwable) {
+                        handleExceptionViaHandler(parentContext, e)
+                    }
+                }                                   
             }
         } finally {
             mutex.unlock()
@@ -189,37 +205,48 @@ private class PublisherCoroutine<in T>(
             if (_nRequested.compareAndSet(cur, upd)) {
                 // unlock the mutex when we don't have back-pressure anymore
                 if (cur == 0L) {
-                    mutex.unlock()
-                    // recheck isActive
-                    if (!isActive && mutex.tryLock())
-                        doLockedSignalCompleted()
+                    unlockAndCheckCompleted()
                 }
                 return
             }
         }
     }
 
-    override fun onCompletedExceptionally(exception: Throwable) = onCompleted(Unit)
-
-    override fun onCompleted(value: Unit) {
+    // assert: isCompleted
+    private fun signalCompleted() {
         while (true) { // lock-free loop for nRequested
             val cur = _nRequested.value
             if (cur == SIGNALLED) return // some other thread holding lock already signalled cancellation/completion
-            check(cur >= 0) // no other thread could have marked it as CLOSED, because onCancellation is invoked once
+            check(cur >= 0) // no other thread could have marked it as CLOSED, because onCompleted[Exceptionally] is invoked once
             if (!_nRequested.compareAndSet(cur, CLOSED)) continue // retry on failed CAS
             // Ok -- marked as CLOSED, now can unlock the mutex if it was locked due to backpressure
             if (cur == 0L) {
                 doLockedSignalCompleted()
             } else {
                 // otherwise mutex was either not locked or locked in concurrent onNext... try lock it to signal completion
-                if (mutex.tryLock())
-                    doLockedSignalCompleted()
+                if (mutex.tryLock()) doLockedSignalCompleted()
                 // Note: if failed `tryLock`, then `doLockedNext` will signal after performing `unlock`
             }
             return // done anyway
         }
     }
 
+    // Note: It is invoked when parent fails to handle an exception and strictly before onCompleted[Exception]
+    // so here we just raise a flag (and it need NOT be volatile!) to handle this exception.
+    // This way we defer decision to handle this exception based on our ability to send this exception
+    // to the subscriber (see doLockedSignalCompleted)
+    override fun handleJobException(exception: Throwable) {
+        handleException = true
+    }
+    
+    override fun onCompletedExceptionally(exception: Throwable) {
+        signalCompleted()
+    }
+
+    override fun onCompleted(value: Unit) {
+        signalCompleted()
+    }
+
     override fun cancel() {
         // Specification requires that after cancellation publisher stops signalling
         // This flag distinguishes subscription cancellation request from the job crash

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

@@ -0,0 +1,65 @@
+/*
+ * Copyright 2016-2018 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 org.junit.*
+import org.junit.Test
+import org.reactivestreams.*
+import java.util.concurrent.*
+import kotlin.test.*
+
+public class PublishParentCancelStressTest : TestBase() {
+    private val dispatcher = newFixedThreadPoolContext(3, "PublishParentCancelStressTest")
+    private val N_TIMES = 5000 * stressTestMultiplier
+
+    @After
+    fun tearDown() {
+        dispatcher.close()
+    }
+
+    @Test
+    fun testStress() = runTest {
+        var unhandled: Throwable? = null
+        val handler = CoroutineExceptionHandler { _, ex -> unhandled = ex }
+        repeat(N_TIMES) {
+            val barrier = CyclicBarrier(4)
+            // launch parent job for publisher
+            val parent = GlobalScope.async<Unit>(dispatcher + handler) {
+                val publisher = publish<Unit> {
+                    // BARRIER #1 - child publisher crashes
+                    barrier.await()
+                    throw TestException()
+                }
+                var sub: Subscription? = null
+                publisher.subscribe(object : Subscriber<Unit> {
+                    override fun onComplete() { error("Cannot be reached") }
+                    override fun onSubscribe(s: Subscription?) { sub = s }
+                    override fun onNext(t: Unit?) { error("Cannot be reached" ) }
+                    override fun onError(t: Throwable?) {
+                        assertTrue(t is TestException)
+                    }
+                })
+                launch {
+                    // BARRIER #3 -- cancel subscription
+                    barrier.await()
+                    sub!!.cancel()
+                }
+                // BARRIER #2 -- parent completes
+                barrier.await()
+                Unit
+            }
+            // BARRIE #4 - go 1-3 together
+            barrier.await()
+            // Make sure exception is not lost, but incorporated into parent
+            val result = kotlin.runCatching { parent.await() }
+            assertTrue(result.exceptionOrNull() is TestException)
+            // Make sure unhandled exception handler was not invoked
+            assertNull(unhandled)
+        }
+    }
+
+    private class TestException : Exception()
+}

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

@@ -130,4 +130,43 @@ class PublishTest : TestBase() {
         sub!!.cancel()
         finish(6)
     }
+
+    @Test
+    fun testPublishFailureCancelsParent() = runTest(
+        expected = { it is TestException }
+    ) {
+        expect(1)
+        val publisher = publish<Unit> {
+            expect(5)
+            throw TestException()
+        }
+        expect(2)
+        publisher.subscribe(object : Subscriber<Unit> {
+            override fun onComplete() {
+                expectUnreached()
+            }
+
+            override fun onSubscribe(s: Subscription) {
+                expect(3)
+            }
+
+            override fun onNext(t: Unit?) {
+                expectUnreached()
+            }
+
+            override fun onError(t: Throwable?) {
+                assertTrue(t is TestException)
+                expect(6)
+            }
+        })
+        expect(4)
+        try {
+            yield() // to coroutine, will crash because it is a cancelled parent coroutine
+        } finally {
+            finish(7)
+        }
+        expectUnreached()
+    }
+
+    private class TestException : Exception()
 }