Fixed produce builder to close the channel on completion instead of

cancelling it, which lead to lost elements with buffered channels.

Fixes #256

Author: elizarov

core/kotlinx-coroutines-core/src/main/kotlin/kotlinx/coroutines/experimental/channels/Actor.kt

@@ -148,7 +148,12 @@ private open class ActorCoroutine<E>(
     parentContext: CoroutineContext,
     channel: Channel<E>,
     active: Boolean
-) : ChannelCoroutine<E>(parentContext, channel, active), ActorScope<E>, ActorJob<E>
+) : ChannelCoroutine<E>(parentContext, channel, active), ActorScope<E>, ActorJob<E> {
+    override fun onCancellation(cause: Throwable?) {
+        if (!_channel.cancel(cause) && cause != null)
+            handleCoroutineException(context, cause)
+    }
+}
 
 private class LazyActorCoroutine<E>(
     parentContext: CoroutineContext,

core/kotlinx-coroutines-core/src/main/kotlin/kotlinx/coroutines/experimental/channels/ChannelCoroutine.kt

@@ -17,20 +17,15 @@
 package kotlinx.coroutines.experimental.channels
 
 import kotlinx.coroutines.experimental.*
-import kotlin.coroutines.experimental.CoroutineContext
+import kotlin.coroutines.experimental.*
 
 internal open class ChannelCoroutine<E>(
     parentContext: CoroutineContext,
-    private val _channel: Channel<E>,
+    protected val _channel: Channel<E>,
     active: Boolean
 ) : AbstractCoroutine<Unit>(parentContext, active), Channel<E> by _channel {
     val channel: Channel<E>
         get() = this
 
-    override fun onCancellation(cause: Throwable?) {
-        if (!_channel.cancel(cause) && cause != null)
-            handleCoroutineException(context, cause)
-    }
-
     override fun cancel(cause: Throwable?): Boolean = super.cancel(cause)
 }

core/kotlinx-coroutines-core/src/main/kotlin/kotlinx/coroutines/experimental/channels/Produce.kt

@@ -105,5 +105,16 @@ public fun <E> buildChannel(
 ): ProducerJob<E> =
     produce(context, capacity, block = block) as ProducerJob<E>
 
-private class ProducerCoroutine<E>(parentContext: CoroutineContext, channel: Channel<E>) :
-    ChannelCoroutine<E>(parentContext, channel, active = true), ProducerScope<E>, ProducerJob<E>
+private class ProducerCoroutine<E>(
+    parentContext: CoroutineContext, channel: Channel<E>
+) : ChannelCoroutine<E>(parentContext, channel, active = true), ProducerScope<E>, ProducerJob<E> {
+    override fun onCancellationInternal(exceptionally: CompletedExceptionally?) {
+        val cause = exceptionally?.cause
+        val processed = when (exceptionally) {
+            is Cancelled -> _channel.cancel(cause) // producer coroutine was cancelled -- cancel channel
+            else -> _channel.close(cause) // producer coroutine has completed -- close channel
+        }
+        if (!processed && cause != null)
+            handleCoroutineException(context, cause)
+    }
+}

core/kotlinx-coroutines-core/src/test/kotlin/kotlinx/coroutines/experimental/channels/ProduceConsumeTest.kt

@@ -0,0 +1,74 @@
+/*
+ * Copyright 2016-2017 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package kotlinx.coroutines.experimental.channels
+
+import kotlinx.coroutines.experimental.*
+import org.junit.*
+import org.junit.Assert.*
+import org.junit.runner.*
+import org.junit.runners.*
+
+@RunWith(Parameterized::class)
+class ProduceConsumeTest(
+    private val capacity: Int,
+    private val number: Int
+) : TestBase() {
+    companion object {
+        @Parameterized.Parameters(name = "capacity={0}, number={1}")
+        @JvmStatic
+        fun params(): Collection<Array<Any>> =
+            listOf(0, 1, 10, 1000, Channel.UNLIMITED).flatMap { capacity ->
+                listOf(1, 10, 1000).map { number ->
+                    arrayOf<Any>(capacity, number)
+                }
+            }
+    }
+
+    @Test
+    fun testProducer() = runTest {
+        var sentAll = false
+        val producer = produce(coroutineContext, capacity = capacity) {
+            for(i in 1..number) {
+                send(i)
+            }
+            sentAll = true
+        }
+        var consumed = 0
+        for (x in producer) {
+            consumed++
+        }
+        assertTrue(sentAll)
+        assertEquals(number, consumed)
+    }
+
+    @Test
+    fun testActor() = runTest {
+        val received = CompletableDeferred<Int>()
+        val actor = actor<Int>(coroutineContext, capacity = capacity) {
+            var n = 0
+            for(i in channel) {
+                n++
+            }
+            received.complete(n)
+        }
+        for(i in 1..number) {
+            actor.send(i)
+        }
+        actor.close()
+        assertEquals(number, received.await())
+    }
+}