Implement SendChannel#invokeOnClose

Fixes #341

Author: qwwdfsad

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

@@ -504,6 +504,7 @@ public abstract class kotlinx/coroutines/experimental/channels/AbstractSendChann
 	protected final fun getClosedForSend ()Lkotlinx/coroutines/experimental/channels/Closed;
 	public final fun getOnSend ()Lkotlinx/coroutines/experimental/selects/SelectClause2;
 	protected final fun getQueue ()Lkotlinx/coroutines/experimental/internal/LockFreeLinkedListHead;
+	public fun invokeOnClose (Lkotlin/jvm/functions/Function1;)V
 	protected abstract fun isBufferAlwaysFull ()Z
 	protected abstract fun isBufferFull ()Z
 	public final fun isClosedForSend ()Z
@@ -734,6 +735,7 @@ public final class kotlinx/coroutines/experimental/channels/ConflatedBroadcastCh
 	public fun getOnSend ()Lkotlinx/coroutines/experimental/selects/SelectClause2;
 	public final fun getValue ()Ljava/lang/Object;
 	public final fun getValueOrNull ()Ljava/lang/Object;
+	public fun invokeOnClose (Lkotlin/jvm/functions/Function1;)V
 	public fun isClosedForSend ()Z
 	public fun isFull ()Z
 	public fun offer (Ljava/lang/Object;)Z
@@ -832,6 +834,7 @@ public abstract interface class kotlinx/coroutines/experimental/channels/Send {
 public abstract interface class kotlinx/coroutines/experimental/channels/SendChannel {
 	public abstract fun close (Ljava/lang/Throwable;)Z
 	public abstract fun getOnSend ()Lkotlinx/coroutines/experimental/selects/SelectClause2;
+	public abstract fun invokeOnClose (Lkotlin/jvm/functions/Function1;)V
 	public abstract fun isClosedForSend ()Z
 	public abstract fun isFull ()Z
 	public abstract fun offer (Ljava/lang/Object;)Z

common/kotlinx-coroutines-core-common/src/channels/AbstractChannel.kt

@@ -4,6 +4,7 @@
 
 package kotlinx.coroutines.experimental.channels
 
+import kotlinx.atomicfu.*
 import kotlinx.coroutines.experimental.*
 import kotlinx.coroutines.experimental.internal.*
 import kotlinx.coroutines.experimental.internalAnnotations.*
@@ -32,6 +33,9 @@ public abstract class AbstractSendChannel<E> : SendChannel<E> {
      */
     protected abstract val isBufferFull: Boolean
 
+    // State transitions: null -> handler -> HANDLER_INVOKED
+    private val onCloseHandler = atomic<Any?>(null)
+
     // ------ internal functions for override by buffered channels ------
 
     /**
@@ -247,11 +251,40 @@ public abstract class AbstractSendChannel<E> : SendChannel<E> {
         }
 
         helpClose(closed)
+        invokeOnCloseHandler(cause)
+        // TODO We can get rid of afterClose
         onClosed(closed)
         afterClose(cause)
         return true
     }
 
+    private fun invokeOnCloseHandler(cause: Throwable?) {
+        val handler = onCloseHandler.value
+        if (handler !== null && handler !== HANDLER_INVOKED
+            && onCloseHandler.compareAndSet(handler, HANDLER_INVOKED)) {
+            // CAS failed -> concurrent invokeOnClose() invoked handler
+            (handler as Handler)(cause)
+        }
+    }
+
+    override fun invokeOnClose(handler: Handler) {
+        // Intricate dance for concurrent invokeOnClose and close calls
+        if (!onCloseHandler.compareAndSet(null, handler)) {
+            val value = onCloseHandler.value
+            if (value === HANDLER_INVOKED) {
+                throw IllegalStateException("Another handler was already registered and successfully invoked")
+            }
+
+            throw IllegalStateException("Another handler was already registered: $value")
+        } else {
+            val closedToken = closedForSend
+            if (closedToken != null && onCloseHandler.compareAndSet(handler, HANDLER_INVOKED)) {
+                // CAS failed -> close() call invoked handler
+                (handler)(closedToken.closeCause)
+            }
+        }
+    }
+
     private fun helpClose(closed: Closed<*>) {
         /*
          * It's important to traverse list from right to left to avoid races with sender.
@@ -983,6 +1016,9 @@ public abstract class AbstractChannel<E> : AbstractSendChannel<E>(), Channel<E>
 /** @suppress **This is unstable API and it is subject to change.** */
 @JvmField internal val SEND_RESUMED = Symbol("SEND_RESUMED")
 
+internal typealias Handler = (Throwable?) -> Unit
+@JvmField internal val HANDLER_INVOKED = Any()
+
 /**
  * Represents sending waiter in the queue.
  * @suppress **This is unstable API and it is subject to change.**
@@ -1043,4 +1079,3 @@ private abstract class Receive<in E> : LockFreeLinkedListNode(), ReceiveOrClosed
     override val offerResult get() = OFFER_SUCCESS
     abstract fun resumeReceiveClosed(closed: Closed<*>)
 }
-

common/kotlinx-coroutines-core-common/src/channels/Channel.kt

@@ -82,6 +82,40 @@ public interface SendChannel<in E> {
      * receive on a failed channel throw the specified [cause] exception.
      */
     public fun close(cause: Throwable? = null): Boolean
+
+    /**
+     * Registers handler which is synchronously invoked once the channel is [closed][close]
+     * or receiving side of this channel is [cancelled][ReceiveChannel.cancel].
+     * Only one handler can be attached to the channel during channel's lifetime.
+     * Handler is invoked when [isClosedForSend] starts to return `true`.
+     * If channel is already closed, handler is invoked immediately.
+     *
+     * The meaning of `cause` that is passed to the handler:
+     * * `null` if channel was closed or cancelled without corresponding argument
+     * * close or cancel cause otherwise.
+     *
+     * Example of usage (exception handling is omitted):
+     * ```
+     * val events = Channel(UNLIMITED)
+     * callbackBasedApi.registerCallback { event ->
+     *   events.offer(event)
+     * }
+     *
+     * val uiUpdater = launch(UI, parent = UILifecycle) {
+     *    events.consume {}
+     *    events.cancel()
+     * }
+     *
+     * events.invokeOnClose { callbackBasedApi.stop() }
+     *
+     * ```
+     *
+     * @throws UnsupportedOperationException if underlying channel doesn't support [invokeOnClose].
+     * Implementation note: currently, [invokeOnClose] is unsupported only by Rx-like integrations
+     *
+     * @throws IllegalStateException if another handler was already registered
+     */
+    public fun invokeOnClose(handler: (cause: Throwable?) -> Unit)
 }
 
 /**

common/kotlinx-coroutines-core-common/src/channels/ConflatedBroadcastChannel.kt

@@ -38,6 +38,8 @@ public class ConflatedBroadcastChannel<E>() : BroadcastChannel<E> {
 
     private val _state = atomic<Any>(INITIAL_STATE) // State | Closed
     private val _updating = atomic(0)
+    // State transitions: null -> handler -> HANDLER_INVOKED
+    private val onCloseHandler = atomic<Any?>(null)
 
     private companion object {
         @JvmField
@@ -163,6 +165,7 @@ public class ConflatedBroadcastChannel<E>() : BroadcastChannel<E> {
                     val update = if (cause == null) CLOSED else Closed(cause)
                     if (_state.compareAndSet(state, update)) {
                         (state as State<E>).subscribers?.forEach { it.close(cause) }
+                        invokeOnCloseHandler(cause)
                         return true
                     }
                 }
@@ -171,6 +174,31 @@ public class ConflatedBroadcastChannel<E>() : BroadcastChannel<E> {
         }
     }
 
+    private fun invokeOnCloseHandler(cause: Throwable?) {
+        val handler = onCloseHandler.value
+        if (handler !== null && handler !== HANDLER_INVOKED
+            && onCloseHandler.compareAndSet(handler, HANDLER_INVOKED)) {
+            (handler as Handler)(cause)
+        }
+    }
+
+    override fun invokeOnClose(handler: Handler) {
+        // Intricate dance for concurrent invokeOnClose and close
+        if (!onCloseHandler.compareAndSet(null, handler)) {
+            val value = onCloseHandler.value
+            if (value === HANDLER_INVOKED) {
+                throw IllegalStateException("Another handler was already registered and successfully invoked")
+            } else {
+                throw IllegalStateException("Another handler was already registered: $value")
+            }
+        } else {
+            val state = _state.value
+            if (state is Closed && onCloseHandler.compareAndSet(handler, HANDLER_INVOKED)) {
+                (handler)(state.closeCause)
+            }
+        }
+    }
+
     /**
      * Closes this broadcast channel. Same as [close].
      */
@@ -249,4 +277,4 @@ public class ConflatedBroadcastChannel<E>() : BroadcastChannel<E> {
 
         public override fun offerInternal(element: E): Any = super.offerInternal(element)
     }
-}
+}

common/kotlinx-coroutines-core-common/test/TestBase.common.kt

@@ -12,6 +12,7 @@ public expect open class TestBase constructor() {
     public fun expect(index: Int)
     public fun expectUnreached()
     public fun finish(index: Int)
+    public fun reset() // Resets counter and finish flag. Workaround for parametrized tests absence in common
 
     public fun runTest(
         expected: ((Throwable) -> Boolean)? = null,

common/kotlinx-coroutines-core-common/test/channels/BasicOperationsTest.kt

@@ -31,6 +31,45 @@ class BasicOperationsTest : TestBase() {
         TestChannelKind.values().forEach { kind -> testReceiveOrNullException(kind) }
     }
 
+    @Test
+    fun testInvokeOnClose() = TestChannelKind.values().forEach { kind ->
+        reset()
+        val channel = kind.create()
+        channel.invokeOnClose {
+            if (it is AssertionError) {
+                expect(3)
+            }
+        }
+        expect(1)
+        channel.offer(42)
+        expect(2)
+        channel.close(AssertionError())
+        finish(4)
+    }
+
+    @Test
+    fun testInvokeOnClosed() = TestChannelKind.values().forEach { kind ->
+        reset()
+        expect(1)
+        val channel = kind.create()
+        channel.close()
+        channel.invokeOnClose { expect(2) }
+        assertFailsWith<IllegalStateException> { channel.invokeOnClose { expect(3) } }
+        finish(3)
+    }
+
+    @Test
+    fun testMultipleInvokeOnClose() = TestChannelKind.values().forEach { kind ->
+        reset()
+        val channel = kind.create()
+        channel.invokeOnClose { expect(3) }
+        expect(1)
+        assertFailsWith<IllegalStateException> { channel.invokeOnClose { expect(4) } }
+        expect(2)
+        channel.close()
+        finish(4)
+    }
+
     private suspend fun testReceiveOrNull(kind: TestChannelKind) {
         val channel = kind.create()
         val d = async(coroutineContext) {

core/kotlinx-coroutines-core/test/TestBase.kt

@@ -4,11 +4,8 @@
 
 package kotlinx.coroutines.experimental
 
-import org.junit.After
-import org.junit.Before
-import java.util.concurrent.atomic.AtomicBoolean
-import java.util.concurrent.atomic.AtomicInteger
-import java.util.concurrent.atomic.AtomicReference
+import org.junit.*
+import java.util.concurrent.atomic.*
 
 /**
  * Base class for tests, so that tests for predictable scheduling of actions in multiple coroutines sharing a single
@@ -88,6 +85,12 @@ public actual open class TestBase actual constructor() {
         check(!finished.getAndSet(true)) { "Should call 'finish(...)' at most once" }
     }
 
+    public actual fun reset() {
+        check(actionIndex.get() == 0 || finished.get()) { "Expecting that 'finish(...)' was invoked, but it was not" }
+        actionIndex.set(0)
+        finished.set(false)
+    }
+
     private lateinit var threadsBefore: Set<Thread>
     private val SHUTDOWN_TIMEOUT = 10_000L // 10s at most to wait
 

core/kotlinx-coroutines-core/test/channels/ChannelIsClosedLinearizabilityTest.kt

@@ -19,7 +19,7 @@ class ChannelIsClosedLinearizabilityTest : TestBase() {
     private lateinit var channel: Channel<Int>
 
     @Reset
-    fun reset() {
+    fun resetChannel() {
         channel = Channel()
     }
 

core/kotlinx-coroutines-core/test/channels/ChannelLinearizabilityTest.kt

@@ -23,7 +23,7 @@ class ChannelLinearizabilityTest : TestBase() {
     private lateinit var channel: Channel<Int>
 
     @Reset
-    fun reset() {
+    fun resetChannel() {
         channel = Channel(capacity)
     }
 

core/kotlinx-coroutines-core/test/channels/InvokeOnCloseStressTest.kt

@@ -0,0 +1,62 @@
+/*
+ * Copyright 2016-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package channels
+
+import kotlinx.coroutines.experimental.*
+import kotlinx.coroutines.experimental.channels.*
+import org.junit.*
+import org.junit.Test
+import java.util.concurrent.*
+import java.util.concurrent.atomic.*
+import kotlin.test.*
+
+class InvokeOnCloseStressTest : TestBase() {
+
+    private val iterations = 1000 * stressTestMultiplier
+
+    private val pool = newFixedThreadPoolContext(3, "InvokeOnCloseStressTest")
+
+    @After
+    fun tearDown() {
+        pool.close()
+    }
+
+    @Test
+    fun testInvokedExactlyOnce() = runBlocking {
+        runStressTest(TestChannelKind.ARRAY_1)
+    }
+
+    @Test
+    fun testInvokedExactlyOnceBroadcast() = runBlocking {
+        runStressTest(TestChannelKind.CONFLATED_BROADCAST)
+    }
+
+    private suspend fun runStressTest(kind: TestChannelKind) {
+        repeat(iterations) {
+            val counter = AtomicInteger(0)
+            val channel = kind.create()
+
+            val latch = CountDownLatch(1)
+            val j1 = async(pool) {
+                latch.await()
+                channel.close()
+            }
+
+            val j2 = async(pool) {
+                latch.await()
+                channel.invokeOnClose { counter.incrementAndGet() }
+            }
+
+            val j3 = async(pool) {
+                latch.await()
+                channel.invokeOnClose { counter.incrementAndGet() }
+            }
+
+            latch.countDown()
+            joinAll(j1, j2, j3)
+            assertEquals(1, counter.get())
+        }
+    }
+}