Fix SOE with select onSend/onReceive clauses on the same channel

* Instead of StackOverflowError we throw IllegalStateException and leave
  the channel in the original state.

Fixes #1411

Author: elizarov

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

@@ -61,7 +61,7 @@ internal abstract class AbstractSendChannel<E> : SendChannel<E> {
      */
     protected open fun offerSelectInternal(element: E, select: SelectInstance<*>): Any {
         // offer atomically with select
-        val offerOp = describeTryOffer(element)
+        val offerOp = describeTryOffer(element, select)
         val failure = select.performAtomicTrySelect(offerOp)
         if (failure != null) return failure
         val receive = offerOp.result
@@ -322,15 +322,17 @@ internal abstract class AbstractSendChannel<E> : SendChannel<E> {
     /**
      * @suppress **This is unstable API and it is subject to change.**
      */
-    protected fun describeTryOffer(element: E): TryOfferDesc<E> = TryOfferDesc(element, queue)
+    protected fun describeTryOffer(element: E, select: SelectInstance<*>): TryOfferDesc<E> =
+        TryOfferDesc(element, select, queue)
 
     /**
      * @suppress **This is unstable API and it is subject to change.**
      */
     protected class TryOfferDesc<E>(
         @JvmField val element: E,
+        override val select: SelectInstance<*>,
         queue: LockFreeLinkedListHead
-    ) : RemoveFirstDesc<ReceiveOrClosed<E>>(queue) {
+    ) : RemoveFirstDesc<ReceiveOrClosed<E>>(queue), SelectInstanceDesc {
         @JvmField var resumeToken: Any? = null
 
         override fun failure(affected: LockFreeLinkedListNode, next: Any): Any? {
@@ -521,7 +523,7 @@ internal abstract class AbstractChannel<E> : AbstractSendChannel<E>(), Channel<E
      */
     protected open fun pollSelectInternal(select: SelectInstance<*>): Any? {
         // poll atomically with select
-        val pollOp = describeTryPoll()
+        val pollOp = describeTryPoll(select)
         val failure = select.performAtomicTrySelect(pollOp)
         if (failure != null) return failure
         val send = pollOp.result
@@ -650,12 +652,16 @@ internal abstract class AbstractChannel<E> : AbstractSendChannel<E>(), Channel<E
     /**
      * @suppress **This is unstable API and it is subject to change.**
      */
-    protected fun describeTryPoll(): TryPollDesc<E> = TryPollDesc(queue)
+    protected fun describeTryPoll(select: SelectInstance<*>): TryPollDesc<E> =
+        TryPollDesc(select, queue)
 
     /**
      * @suppress **This is unstable API and it is subject to change.**
      */
-    protected class TryPollDesc<E>(queue: LockFreeLinkedListHead) : RemoveFirstDesc<Send>(queue) {
+    protected class TryPollDesc<E>(
+        override val select: SelectInstance<*>,
+        queue: LockFreeLinkedListHead
+    ) : RemoveFirstDesc<Send>(queue), SelectInstanceDesc {
         @JvmField var resumeToken: Any? = null
         @JvmField var pollResult: E? = null
 

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

@@ -88,7 +88,7 @@ internal open class ArrayChannel<E>(
                 // check for receivers that were waiting on empty queue
                 if (size == 0) {
                     loop@ while (true) {
-                        val offerOp = describeTryOffer(element)
+                        val offerOp = describeTryOffer(element, select)
                         val failure = select.performAtomicTrySelect(offerOp)
                         when {
                             failure == null -> { // offered successfully
@@ -175,7 +175,7 @@ internal open class ArrayChannel<E>(
             var replacement: Any? = POLL_FAILED
             if (size == capacity) {
                 loop@ while (true) {
-                    val pollOp = describeTryPoll()
+                    val pollOp = describeTryPoll(select)
                     val failure = select.performAtomicTrySelect(pollOp)
                     when {
                         failure == null -> { // polled successfully

kotlinx-coroutines-core/common/src/selects/Select.kt

@@ -133,6 +133,16 @@ public interface SelectInstance<in R> {
     public fun disposeOnSelect(handle: DisposableHandle)
 }
 
+/**
+ * Descriptor that is a part of the atomic select operator.
+ * 
+ * @suppress **This is unstable API and it is subject to change.**
+ */
+@InternalCoroutinesApi
+public interface SelectInstanceDesc {
+    val select: SelectInstance<*>
+}
+
 /**
  * Waits for the result of multiple suspending functions simultaneously, which are specified using _clauses_
  * in the [builder] scope of this select invocation. The caller is suspended until one of the clauses
@@ -325,7 +335,7 @@ internal class SelectBuilderImpl<in R>(
         while (true) { // lock-free loop on state
             val state = this.state
             if (state === this) {
-                if (addLastIf(node, { this.state === this }))
+                if (addLastIf(node) { this.state === this })
                     return
             } else { // already selected
                 handle.dispose()
@@ -344,9 +354,26 @@ internal class SelectBuilderImpl<in R>(
     // it is just like start(), but support idempotent start
     override fun trySelect(idempotent: Any?): Boolean {
         assert { idempotent !is OpDescriptor } // "cannot use OpDescriptor as idempotent marker"
-        while (true) { // lock-free loop on state
-            val state = this.state
+        _state.loop { state -> // lock-free loop on state
             when {
+                // Found descriptor from the same select instance while registering another clause on the same instance
+                state is AtomicSelectOp<*> && state.impl === this &&
+                    idempotent is SelectInstanceDesc && idempotent.select === this ->
+                {
+                    /*
+                     * We cannot do state.perform(this) here and "help" it since it is the same select and
+                     * we'll get StackOverflowError. See https://github.com/Kotlin/kotlinx.coroutines/issues/1411
+                     * "Properly" supporting this case is complicated in the current select architecture, since we
+                     * have to "pull up" body of two clauses that need to be activated and link them
+                     * together so that they are sequentially executed, followed by the continuation of
+                     * the code after the select expression. This requires massive implementation rewrite
+                     * (instead of immediately resuming clause bodies, they have to be stored somewhere first).
+                     */
+                    error("Cannot use matching select clauses on the same object")
+                }
+                // Found descriptor of any other atomic operation -- help it
+                state is OpDescriptor -> state.perform(this)
+                // Found initial state (not selected yet)
                 state === this -> {
                     if (_state.compareAndSet(this, idempotent)) {
                         doAfterSelect()
@@ -361,10 +388,14 @@ internal class SelectBuilderImpl<in R>(
         }
     }
 
-    override fun performAtomicTrySelect(desc: AtomicDesc): Any? = AtomicSelectOp(desc, true).perform(null)
-    override fun performAtomicIfNotSelected(desc: AtomicDesc): Any? = AtomicSelectOp(desc, false).perform(null)
+    override fun performAtomicTrySelect(desc: AtomicDesc): Any? =
+        AtomicSelectOp(this, desc, true).perform(null)
+
+    override fun performAtomicIfNotSelected(desc: AtomicDesc): Any? =
+        AtomicSelectOp(this, desc, false).perform(null)
 
-    private inner class AtomicSelectOp(
+    private class AtomicSelectOp<R>(
+        @JvmField val impl: SelectBuilderImpl<R>,
         @JvmField val desc: AtomicDesc,
         @JvmField val select: Boolean
     ) : AtomicOp<Any?>() {
@@ -375,34 +406,45 @@ internal class SelectBuilderImpl<in R>(
                 // we are originator (affected reference is not null if helping)
                 prepareIfNotSelected()?.let { return it }
             }
-            return desc.prepare(this)
+            try {
+                return desc.prepare(this)
+            } catch (e: Throwable) {
+                // undo prepareIfNotSelected on crash (for example if IllegalStateException is thrown)
+                if (affected == null) undoPrepare()
+                throw e
+            }
         }
 
         override fun complete(affected: Any?, failure: Any?) {
             completeSelect(failure)
             desc.complete(this, failure)
         }
 
-        fun prepareIfNotSelected(): Any? {
-            _state.loop { state ->
+        private fun prepareIfNotSelected(): Any? {
+            impl._state.loop { state ->
                 when {
-                    state === this@AtomicSelectOp -> return null // already in progress
-                    state is OpDescriptor -> state.perform(this@SelectBuilderImpl) // help
-                    state === this@SelectBuilderImpl -> {
-                        if (_state.compareAndSet(this@SelectBuilderImpl, this@AtomicSelectOp))
+                    state === this -> return null // already in progress
+                    state is OpDescriptor -> state.perform(impl) // help
+                    state === impl -> {
+                        if (impl._state.compareAndSet(impl, this))
                             return null // success
                     }
                     else -> return ALREADY_SELECTED
                 }
             }
         }
 
+        // reverts the change done by prepareIfNotSelected
+        private fun undoPrepare() {
+            impl._state.compareAndSet(this, impl)
+        }
+
         private fun completeSelect(failure: Any?) {
             val selectSuccess = select && failure == null
-            val update = if (selectSuccess) null else this@SelectBuilderImpl
-            if (_state.compareAndSet(this@AtomicSelectOp, update)) {
+            val update = if (selectSuccess) null else impl
+            if (impl._state.compareAndSet(this, update)) {
                 if (selectSuccess)
-                    doAfterSelect()
+                    impl.doAfterSelect()
             }
         }
     }

kotlinx-coroutines-core/common/test/selects/SelectRendezvousChannelTest.kt

@@ -409,4 +409,38 @@ class SelectRendezvousChannelTest : TestBase() {
         if (!trySelect(null)) return
         block.startCoroutineUnintercepted(this)
     }
+
+    @Test
+    fun testSelectSendAndReceive() = runTest {
+        val c = Channel<Int>()
+        assertFailsWith<IllegalStateException> {
+            select<Unit> {
+                c.onSend(1) { expectUnreached() }
+                c.onReceive { expectUnreached() }
+            }
+        }
+        checkNotBroken(c)
+    }
+
+    @Test
+    fun testSelectReceiveAndSend() = runTest {
+        val c = Channel<Int>()
+        assertFailsWith<IllegalStateException> {
+            select<Unit> {
+                c.onReceive { expectUnreached() }
+                c.onSend(1) { expectUnreached() }
+            }
+        }
+        checkNotBroken(c)
+    }
+
+    // makes sure the channel is not broken
+    private suspend fun checkNotBroken(c: Channel<Int>) {
+        coroutineScope {
+            launch {
+                c.send(42)
+            }
+            assertEquals(42, c.receive())
+        }
+    }
 }

kotlinx-coroutines-core/jvm/src/internal/LockFreeLinkedList.kt

@@ -466,9 +466,15 @@ public actual open class LockFreeLinkedListNode {
                 val prepareOp = PrepareOp(next as Node, op as AtomicOp<Node>, this)
                 if (affected._next.compareAndSet(next, prepareOp)) {
                     // prepared -- complete preparations
-                    val prepFail = prepareOp.perform(affected)
-                    if (prepFail === REMOVE_PREPARED) continue // retry
-                    return prepFail
+                    try {
+                        val prepFail = prepareOp.perform(affected)
+                        if (prepFail === REMOVE_PREPARED) continue // retry
+                        return prepFail
+                    } catch (e: Throwable) {
+                        // Crashed during preparation (for example IllegalStateExpception) -- undo & rethrow
+                        affected._next.compareAndSet(prepareOp, next)
+                        throw e
+                    }
                 }
             }
         }