Cancellation in consumeEach should dispose Rx Observable

Fixed bugs in MaybeSource/ObservableSource.consumeEach implementation
so that observable is disposed on cancellation.

Also optimized implementation of bridge function to avoid extra dispose
calls if possible (this is permissible by specification, though)

Fixes #1008

Author: elizarov

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

@@ -30,9 +30,8 @@ public fun <T> Publisher<T>.openSubscription(request: Int = 0): ReceiveChannel<T
 /**
  * Subscribes to this [Publisher] and performs the specified action for each received element.
  */
-public suspend inline fun <T> Publisher<T>.consumeEach(action: (T) -> Unit) {
+public suspend inline fun <T> Publisher<T>.consumeEach(action: (T) -> Unit) =
     openSubscription().consumeEach(action)
-}
 
 @Suppress("INVISIBLE_REFERENCE", "INVISIBLE_MEMBER")
 private class SubscriptionChannel<T>(
@@ -75,6 +74,7 @@ private class SubscriptionChannel<T>(
     @Suppress("CANNOT_OVERRIDE_INVISIBLE_MEMBER")
     override fun onClosedIdempotent(closed: LockFreeLinkedListNode) {
         subscription?.cancel()
+        subscription = null // optimization -- no need to cancel it again
     }
 
     // Subscriber overrides

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

@@ -74,24 +74,6 @@ class IntegrationTest(
         assertThat(cnt, IsEqual(1))
     }
 
-    @Test
-    fun testFailingConsumer() = runTest {
-        val pub = publish {
-            repeat(3) {
-                expect(it + 1) // expect(1), expect(2) *should* be invoked
-                send(it)
-            }
-        }
-
-        try {
-            pub.consumeEach {
-                throw TestException()
-            }
-        } catch (e: TestException) {
-            finish(3)
-        }
-    }
-
     @Test
     fun testNumbers() = runBlocking<Unit> {
         val n = 100 * stressTestMultiplier

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

@@ -252,4 +252,21 @@ class PublishTest : TestBase() {
         latch.await()
         finish(8)
     }
+
+    @Test
+    fun testFailingConsumer() = runTest {
+        val pub = publish {
+            repeat(3) {
+                expect(it + 1) // expect(1), expect(2) *should* be invoked
+                send(it)
+            }
+        }
+        try {
+            pub.consumeEach {
+                throw TestException()
+            }
+        } catch (e: TestException) {
+            finish(3)
+        }
+    }
 }

reactive/kotlinx-coroutines-reactor/test/FluxTest.kt

@@ -5,8 +5,11 @@
 package kotlinx.coroutines.reactor
 
 import kotlinx.coroutines.*
+import kotlinx.coroutines.reactive.*
 import org.hamcrest.core.*
 import org.junit.*
+import org.junit.Test
+import kotlin.test.*
 
 class FluxTest : TestBase() {
     @Test
@@ -80,4 +83,59 @@ class FluxTest : TestBase() {
             { assert(it is RuntimeException) }
         )
     }
+
+    @Test
+    fun testNotifyOnceOnCancellation() = runTest {
+        expect(1)
+        val observable =
+            flux {
+                expect(5)
+                send("OK")
+                try {
+                    delay(Long.MAX_VALUE)
+                } catch (e: CancellationException) {
+                    expect(11)
+                }
+            }
+            .doOnNext {
+                expect(6)
+                assertEquals("OK", it)
+            }
+            .doOnCancel {
+                expect(10) // notified once!
+            }
+        expect(2)
+        val job = launch(start = CoroutineStart.UNDISPATCHED) {
+            expect(3)
+            observable.consumeEach {
+                expect(8)
+                assertEquals("OK", it)
+            }
+        }
+        expect(4)
+        yield() // to observable code
+        expect(7)
+        yield() // to consuming coroutines
+        expect(9)
+        job.cancel()
+        job.join()
+        finish(12)
+    }
+
+    @Test
+    fun testFailingConsumer() = runTest {
+        val pub = flux {
+            repeat(3) {
+                expect(it + 1) // expect(1), expect(2) *should* be invoked
+                send(it)
+            }
+        }
+        try {
+            pub.consumeEach {
+                throw TestException()
+            }
+        } catch (e: TestException) {
+            finish(3)
+        }
+    }
 }

reactive/kotlinx-coroutines-rx2/src/RxChannel.kt

@@ -43,20 +43,14 @@ public fun <T> ObservableSource<T>.openSubscription(): ReceiveChannel<T> {
 /**
  * Subscribes to this [MaybeSource] and performs the specified action for each received element.
  */
-public suspend inline fun <T> MaybeSource<T>.consumeEach(action: (T) -> Unit) {
-    val channel = openSubscription()
-    for (x in channel) action(x)
-    channel.cancel()
-}
+public suspend inline fun <T> MaybeSource<T>.consumeEach(action: (T) -> Unit) =
+    openSubscription().consumeEach(action)
 
 /**
  * Subscribes to this [ObservableSource] and performs the specified action for each received element.
  */
-public suspend inline fun <T> ObservableSource<T>.consumeEach(action: (T) -> Unit) {
-    val channel = openSubscription()
-    for (x in channel) action(x)
-    channel.cancel()
-}
+public suspend inline fun <T> ObservableSource<T>.consumeEach(action: (T) -> Unit) =
+    openSubscription().consumeEach(action)
 
 @Suppress("INVISIBLE_REFERENCE", "INVISIBLE_MEMBER")
 private class SubscriptionChannel<T> :
@@ -68,6 +62,7 @@ private class SubscriptionChannel<T> :
     @Suppress("CANNOT_OVERRIDE_INVISIBLE_MEMBER")
     override fun onClosedIdempotent(closed: LockFreeLinkedListNode) {
         subscription?.dispose()
+        subscription = null // optimization -- no need to dispose it again
     }
 
     // Observer overrider

reactive/kotlinx-coroutines-rx2/test/FlowableTest.kt

@@ -8,6 +8,8 @@ import kotlinx.coroutines.*
 import kotlinx.coroutines.reactive.*
 import org.hamcrest.core.*
 import org.junit.*
+import org.junit.Test
+import kotlin.test.*
 
 class FlowableTest : TestBase() {
     @Test
@@ -81,4 +83,59 @@ class FlowableTest : TestBase() {
             { assert(it is RuntimeException) }
         )
     }
+
+    @Test
+    fun testNotifyOnceOnCancellation() = runTest {
+        expect(1)
+        val observable =
+            rxFlowable {
+                expect(5)
+                send("OK")
+                try {
+                    delay(Long.MAX_VALUE)
+                } catch (e: CancellationException) {
+                    expect(11)
+                }
+            }
+            .doOnNext {
+                expect(6)
+                assertEquals("OK", it)
+            }
+            .doOnCancel {
+                expect(10) // notified once!
+            }
+        expect(2)
+        val job = launch(start = CoroutineStart.UNDISPATCHED) {
+            expect(3)
+            observable.consumeEach{
+                expect(8)
+                assertEquals("OK", it)
+            }
+        }
+        expect(4)
+        yield() // to observable code
+        expect(7)
+        yield() // to consuming coroutines
+        expect(9)
+        job.cancel()
+        job.join()
+        finish(12)
+    }
+
+    @Test
+    fun testFailingConsumer() = runTest {
+        val pub = rxFlowable {
+            repeat(3) {
+                expect(it + 1) // expect(1), expect(2) *should* be invoked
+                send(it)
+            }
+        }
+        try {
+            pub.consumeEach {
+                throw TestException()
+            }
+        } catch (e: TestException) {
+            finish(3)
+        }
+    }
 }

reactive/kotlinx-coroutines-rx2/test/MaybeTest.kt

@@ -12,6 +12,7 @@ import org.hamcrest.core.*
 import org.junit.*
 import org.junit.Assert.*
 import java.util.concurrent.*
+import java.util.concurrent.CancellationException
 
 class MaybeTest : TestBase() {
     @Before
@@ -211,4 +212,30 @@ class MaybeTest : TestBase() {
             { assert(it is RuntimeException) }
         )
     }
+
+    @Test
+    fun testCancelledConsumer() = runTest {
+        expect(1)
+        val maybe = rxMaybe<Int> {
+            expect(4)
+            try {
+                delay(Long.MAX_VALUE)
+            } catch (e: CancellationException) {
+                expect(6)
+            }
+            42
+        }
+        expect(2)
+        val timeout = withTimeoutOrNull(100) {
+            expect(3)
+            maybe.consumeEach {
+                expectUnreached()
+            }
+            expectUnreached()
+        }
+        assertNull(timeout)
+        expect(5)
+        yield() // must cancel code inside maybe!!!
+        finish(7)
+    }
 }

reactive/kotlinx-coroutines-rx2/test/ObservableTest.kt

@@ -7,6 +7,8 @@ package kotlinx.coroutines.rx2
 import kotlinx.coroutines.*
 import org.hamcrest.core.*
 import org.junit.*
+import org.junit.Test
+import kotlin.test.*
 
 class ObservableTest : TestBase() {
     @Test
@@ -80,4 +82,64 @@ class ObservableTest : TestBase() {
             { assert(it is RuntimeException) }
         )
     }
+
+    @Test
+    fun testNotifyOnceOnCancellation() = runTest {
+        expect(1)
+        val observable =
+            rxObservable {
+                expect(5)
+                send("OK")
+                try {
+                    delay(Long.MAX_VALUE)
+                } catch (e: CancellationException) {
+                    expect(11)
+                }
+            }
+            .doOnNext {
+                expect(6)
+                assertEquals("OK", it)
+            }
+            .doOnDispose {
+                expect(10) // notified once!
+            }
+        expect(2)
+        val job = launch(start = CoroutineStart.UNDISPATCHED) {
+            expect(3)
+            observable.consumeEach{
+                expect(8)
+                assertEquals("OK", it)
+            }
+        }
+        expect(4)
+        yield() // to observable code
+        expect(7)
+        yield() // to consuming coroutines
+        expect(9)
+        job.cancel()
+        job.join()
+        finish(12)
+    }
+
+    @Test
+    fun testFailingConsumer() = runTest {
+        expect(1)
+        val pub = rxObservable {
+            expect(2)
+            send("OK")
+            try {
+                delay(Long.MAX_VALUE)
+            } catch (e: CancellationException) {
+                finish(5)
+            }
+        }
+        try {
+            pub.consumeEach {
+                expect(3)
+                throw TestException()
+            }
+        } catch (e: TestException) {
+            expect(4)
+        }
+    }
 }