Fix race condition in pair select

This bug was introduced by #1524. The crux of problem is that
TryOffer/PollDesc.onPrepare method is no longer allowed to update
fields in these classes (like "resumeToken" and "pollResult") after call
to tryResumeSend/Receive method, because the latter will complete
the ongoing atomic operation and helper method might find it complete
and try reading "resumeToken" which was not initialized yet.

This change removes "pollResult" field which was not really needed
("result.pollResult" field is used) and removes "resumeToken" by
exploiting the fact that current implementation of
CancellableContinuationImpl does not need a token anymore. However,
CancellableContinuation.tryResume/completeResume ABI is left intact,
because it is used by 3rd party code.

This fix lead to overall simplification of the code. A number of fields
and an auxiliary IdempotentTokenValue class are removed, tokens used to
indicate various results are consolidated, so that resume success
is now consistently indicated by a single RESUME_TOKEN symbol.

Fixes #1561

Author: elizarov

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

@@ -14,6 +14,10 @@ private const val UNDECIDED = 0
 private const val SUSPENDED = 1
 private const val RESUMED = 2
 
+@JvmField
+@SharedImmutable
+internal val RESUME_TOKEN = Symbol("RESUME_TOKEN")
+
 /**
  * @suppress **This is unstable API and it is subject to change.**
  */
@@ -285,20 +289,21 @@ internal open class CancellableContinuationImpl<in T>(
         }
     }
 
+    // Note: Always returns RESUME_TOKEN | null
     override fun tryResume(value: T, idempotent: Any?): Any? {
         _state.loop { state ->
             when (state) {
                 is NotCompleted -> {
                     val update: Any? = if (idempotent == null) value else
-                        CompletedIdempotentResult(idempotent, value, state)
+                        CompletedIdempotentResult(idempotent, value)
                     if (!_state.compareAndSet(state, update)) return@loop // retry on cas failure
                     disposeParentHandle()
-                    return state
+                    return RESUME_TOKEN
                 }
                 is CompletedIdempotentResult -> {
                     return if (state.idempotentResume === idempotent) {
                         assert { state.result === value } // "Non-idempotent resume"
-                        state.token
+                        RESUME_TOKEN
                     } else {
                         null
                     }
@@ -315,15 +320,16 @@ internal open class CancellableContinuationImpl<in T>(
                     val update = CompletedExceptionally(exception)
                     if (!_state.compareAndSet(state, update)) return@loop // retry on cas failure
                     disposeParentHandle()
-                    return state
+                    return RESUME_TOKEN
                 }
                 else -> return null // cannot resume -- not active anymore
             }
         }
     }
 
+    // note: token is always RESUME_TOKEN
     override fun completeResume(token: Any) {
-        // note: We don't actually use token anymore, because handler needs to be invoked on cancellation only
+        assert { token === RESUME_TOKEN }
         dispatchResume(resumeMode)
     }
 
@@ -375,8 +381,7 @@ private class InvokeOnCancel( // Clashes with InvokeOnCancellation
 
 private class CompletedIdempotentResult(
     @JvmField val idempotentResume: Any?,
-    @JvmField val result: Any?,
-    @JvmField val token: NotCompleted
+    @JvmField val result: Any?
 ) {
     override fun toString(): String = "CompletedIdempotentResult[$result]"
 }

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

@@ -48,7 +48,8 @@ internal abstract class AbstractSendChannel<E> : SendChannel<E> {
             val receive = takeFirstReceiveOrPeekClosed() ?: return OFFER_FAILED
             val token = receive.tryResumeReceive(element, null)
             if (token != null) {
-                receive.completeResumeReceive(token)
+                assert { token === RESUME_TOKEN }
+                receive.completeResumeReceive(element)
                 return receive.offerResult
             }
         }
@@ -65,7 +66,7 @@ internal abstract class AbstractSendChannel<E> : SendChannel<E> {
         val failure = select.performAtomicTrySelect(offerOp)
         if (failure != null) return failure
         val receive = offerOp.result
-        receive.completeResumeReceive(offerOp.resumeToken!!)
+        receive.completeResumeReceive(element)
         return receive.offerResult
     }
 
@@ -354,8 +355,6 @@ internal abstract class AbstractSendChannel<E> : SendChannel<E> {
         @JvmField val element: E,
         queue: LockFreeLinkedListHead
     ) : RemoveFirstDesc<ReceiveOrClosed<E>>(queue) {
-        @JvmField var resumeToken: Any? = null
-
         override fun failure(affected: LockFreeLinkedListNode): Any? = when (affected) {
             is Closed<*> -> affected
             !is ReceiveOrClosed<*> -> OFFER_FAILED
@@ -367,7 +366,7 @@ internal abstract class AbstractSendChannel<E> : SendChannel<E> {
             val affected = prepareOp.affected as ReceiveOrClosed<E> // see "failure" impl
             val token = affected.tryResumeReceive(element, prepareOp) ?: return REMOVE_PREPARED
             if (token === RETRY_ATOMIC) return RETRY_ATOMIC
-            resumeToken = token
+            assert { token === RESUME_TOKEN }
             return null
         }
     }
@@ -454,8 +453,7 @@ internal abstract class AbstractSendChannel<E> : SendChannel<E> {
         override fun tryResumeSend(otherOp: PrepareOp?): Any? =
             select.trySelectOther(otherOp)
 
-        override fun completeResumeSend(token: Any) {
-            assert { token === SELECT_STARTED }
+        override fun completeResumeSend() {
             block.startCoroutine(receiver = channel, completion = select.completion)
         }
 
@@ -475,8 +473,8 @@ internal abstract class AbstractSendChannel<E> : SendChannel<E> {
         @JvmField val element: E
     ) : Send() {
         override val pollResult: Any? get() = element
-        override fun tryResumeSend(otherOp: PrepareOp?): Any?  = SEND_RESUMED.also { otherOp?.finishPrepare() }
-        override fun completeResumeSend(token: Any) { assert { token === SEND_RESUMED } }
+        override fun tryResumeSend(otherOp: PrepareOp?): Any?  = RESUME_TOKEN.also { otherOp?.finishPrepare() }
+        override fun completeResumeSend() {}
         override fun resumeSendClosed(closed: Closed<*>) {}
     }
 }
@@ -511,7 +509,8 @@ internal abstract class AbstractChannel<E> : AbstractSendChannel<E>(), Channel<E
             val send = takeFirstSendOrPeekClosed() ?: return POLL_FAILED
             val token = send.tryResumeSend(null)
             if (token != null) {
-                send.completeResumeSend(token)
+                assert { token === RESUME_TOKEN }
+                send.completeResumeSend()
                 return send.pollResult
             }
         }
@@ -528,8 +527,8 @@ internal abstract class AbstractChannel<E> : AbstractSendChannel<E>(), Channel<E
         val failure = select.performAtomicTrySelect(pollOp)
         if (failure != null) return failure
         val send = pollOp.result
-        send.completeResumeSend(pollOp.resumeToken!!)
-        return pollOp.pollResult
+        send.completeResumeSend()
+        return pollOp.result.pollResult
     }
 
     // ------ state functions & helpers for concrete implementations ------
@@ -673,9 +672,6 @@ internal abstract class AbstractChannel<E> : AbstractSendChannel<E>(), Channel<E
      * @suppress **This is unstable API and it is subject to change.**
      */
     protected class TryPollDesc<E>(queue: LockFreeLinkedListHead) : RemoveFirstDesc<Send>(queue) {
-        @JvmField var resumeToken: Any? = null
-        @JvmField var pollResult: E? = null
-
         override fun failure(affected: LockFreeLinkedListNode): Any? = when (affected) {
             is Closed<*> -> affected
             !is Send -> POLL_FAILED
@@ -687,8 +683,7 @@ internal abstract class AbstractChannel<E> : AbstractSendChannel<E>(), Channel<E
             val affected = prepareOp.affected as Send // see "failure" impl
             val token = affected.tryResumeSend(prepareOp) ?: return REMOVE_PREPARED
             if (token === RETRY_ATOMIC) return RETRY_ATOMIC
-            resumeToken = token
-            pollResult = affected.pollResult as E
+            assert { token === RESUME_TOKEN }
             return null
         }
     }
@@ -908,7 +903,8 @@ internal abstract class AbstractChannel<E> : AbstractSendChannel<E>(), Channel<E
             return cont.tryResume(resumeValue(value), otherOp?.desc)
         }
 
-        override fun completeResumeReceive(token: Any) = cont.completeResume(token)
+        override fun completeResumeReceive(value: E) = cont.completeResume(RESUME_TOKEN)
+
         override fun resumeReceiveClosed(closed: Closed<*>) {
             when {
                 receiveMode == RECEIVE_NULL_ON_CLOSE && closed.closeCause == null -> cont.resume(null)
@@ -925,25 +921,16 @@ internal abstract class AbstractChannel<E> : AbstractSendChannel<E>(), Channel<E
     ) : Receive<E>() {
         override fun tryResumeReceive(value: E, otherOp: PrepareOp?): Any? {
             otherOp?.finishPrepare()
-            val token = cont.tryResume(true, otherOp?.desc)
-            if (token != null) {
-                /*
-                   When otherOp != null this invocation can be stale and we cannot directly update iterator.result
-                   Instead, we save both token & result into a temporary IdempotentTokenValue object and
-                   set iterator result only in completeResumeReceive that is going to be invoked just once
-                 */
-                if (otherOp != null) return IdempotentTokenValue(token, value)
-                iterator.result = value
-            }
-            return token
+            return cont.tryResume(true, otherOp?.desc)
         }
 
-        override fun completeResumeReceive(token: Any) {
-            if (token is IdempotentTokenValue<*>) {
-                iterator.result = token.value
-                cont.completeResume(token.token)
-            } else
-                cont.completeResume(token)
+        override fun completeResumeReceive(value: E) {
+            /*
+               When otherOp != null invocation of tryResumeReceive can happen multiple times and much later,
+               but completeResumeReceive is called once so we set iterator result here.
+             */
+            iterator.result = value
+            cont.completeResume(RESUME_TOKEN)
         }
 
         override fun resumeReceiveClosed(closed: Closed<*>) {
@@ -966,14 +953,11 @@ internal abstract class AbstractChannel<E> : AbstractSendChannel<E>(), Channel<E
         @JvmField val block: suspend (Any?) -> R,
         @JvmField val receiveMode: Int
     ) : Receive<E>(), DisposableHandle {
-        override fun tryResumeReceive(value: E, otherOp: PrepareOp?): Any? {
-            val result = select.trySelectOther(otherOp)
-            return if (result === SELECT_STARTED) value ?: NULL_VALUE else result
-        }
+        override fun tryResumeReceive(value: E, otherOp: PrepareOp?): Any? =
+            select.trySelectOther(otherOp)
 
         @Suppress("UNCHECKED_CAST")
-        override fun completeResumeReceive(token: Any) {
-            val value: E = NULL_VALUE.unbox<E>(token)
+        override fun completeResumeReceive(value: E) {
             block.startCoroutine(if (receiveMode == RECEIVE_RESULT) ValueOrClosed.value(value) else value, select.completion)
         }
 
@@ -997,11 +981,6 @@ internal abstract class AbstractChannel<E> : AbstractSendChannel<E>(), Channel<E
 
         override fun toString(): String = "ReceiveSelect[$select,receiveMode=$receiveMode]"
     }
-
-    private class IdempotentTokenValue<out E>(
-        @JvmField val token: Any,
-        @JvmField val value: E
-    )
 }
 
 // receiveMode values
@@ -1025,18 +1004,6 @@ internal val POLL_FAILED: Any = Symbol("POLL_FAILED")
 @SharedImmutable
 internal val ENQUEUE_FAILED: Any = Symbol("ENQUEUE_FAILED")
 
-@JvmField
-@SharedImmutable
-internal val NULL_VALUE: Symbol = Symbol("NULL_VALUE")
-
-@JvmField
-@SharedImmutable
-internal val CLOSE_RESUMED: Any = Symbol("CLOSE_RESUMED")
-
-@JvmField
-@SharedImmutable
-internal val SEND_RESUMED: Any = Symbol("SEND_RESUMED")
-
 @JvmField
 @SharedImmutable
 internal val HANDLER_INVOKED: Any = Symbol("ON_CLOSE_HANDLER_INVOKED")
@@ -1050,10 +1017,10 @@ internal abstract class Send : LockFreeLinkedListNode() {
     abstract val pollResult: Any? // E | Closed
     // Returns: null - failure,
     //          RETRY_ATOMIC for retry (only when otherOp != null),
-    //          otherwise token for completeResumeSend
+    //          RESUME_TOKEN on success (call completeResumeSend)
     // Must call otherOp?.finishPrepare() before deciding on result other than RETRY_ATOMIC
     abstract fun tryResumeSend(otherOp: PrepareOp?): Any?
-    abstract fun completeResumeSend(token: Any)
+    abstract fun completeResumeSend()
     abstract fun resumeSendClosed(closed: Closed<*>)
 }
 
@@ -1064,10 +1031,10 @@ internal interface ReceiveOrClosed<in E> {
     val offerResult: Any // OFFER_SUCCESS | Closed
     // Returns: null - failure,
     //          RETRY_ATOMIC for retry (only when otherOp != null),
-    //          otherwise token for completeResumeReceive
+    //          RESUME_TOKEN on success (call completeResumeReceive)
     // Must call otherOp?.finishPrepare() before deciding on result other than RETRY_ATOMIC
     fun tryResumeReceive(value: E, otherOp: PrepareOp?): Any?
-    fun completeResumeReceive(token: Any)
+    fun completeResumeReceive(value: E)
 }
 
 /**
@@ -1082,7 +1049,7 @@ internal class SendElement(
         otherOp?.finishPrepare()
         return cont.tryResume(Unit, otherOp?.desc)
     }
-    override fun completeResumeSend(token: Any) = cont.completeResume(token)
+    override fun completeResumeSend() = cont.completeResume(RESUME_TOKEN)
     override fun resumeSendClosed(closed: Closed<*>) = cont.resumeWithException(closed.sendException)
     override fun toString(): String = "SendElement($pollResult)"
 }
@@ -1098,10 +1065,10 @@ internal class Closed<in E>(
 
     override val offerResult get() = this
     override val pollResult get() = this
-    override fun tryResumeSend(otherOp: PrepareOp?): Any? = CLOSE_RESUMED.also { otherOp?.finishPrepare() }
-    override fun completeResumeSend(token: Any) { assert { token === CLOSE_RESUMED } }
-    override fun tryResumeReceive(value: E, otherOp: PrepareOp?): Any? = CLOSE_RESUMED.also { otherOp?.finishPrepare() }
-    override fun completeResumeReceive(token: Any) { assert { token === CLOSE_RESUMED } }
+    override fun tryResumeSend(otherOp: PrepareOp?): Any? = RESUME_TOKEN.also { otherOp?.finishPrepare() }
+    override fun completeResumeSend() {}
+    override fun tryResumeReceive(value: E, otherOp: PrepareOp?): Any? = RESUME_TOKEN.also { otherOp?.finishPrepare() }
+    override fun completeResumeReceive(value: E) {}
     override fun resumeSendClosed(closed: Closed<*>) = assert { false } // "Should be never invoked"
     override fun toString(): String = "Closed[$closeCause]"
 }

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

@@ -134,7 +134,6 @@ internal class ArrayBroadcastChannel<E>(
     private tailrec fun updateHead(addSub: Subscriber<E>? = null, removeSub: Subscriber<E>? = null) {
         // update head in a tail rec loop
         var send: Send? = null
-        var token: Any? = null
         bufferLock.withLock {
             if (addSub != null) {
                 addSub.subHead = tail // start from last element
@@ -163,8 +162,9 @@ internal class ArrayBroadcastChannel<E>(
                     while (true) {
                         send = takeFirstSendOrPeekClosed() ?: break // when when no sender
                         if (send is Closed<*>) break // break when closed for send
-                        token = send!!.tryResumeSend(null)
+                        val token = send!!.tryResumeSend(null)
                         if (token != null) {
+                            assert { token === RESUME_TOKEN }
                             // put sent element to the buffer
                             buffer[(tail % capacity).toInt()] = (send as Send).pollResult
                             this.size = size + 1
@@ -177,7 +177,7 @@ internal class ArrayBroadcastChannel<E>(
             return // done updating here -> return
         }
         // we only get out of the lock normally when there is a sender to resume
-        send!!.completeResumeSend(token!!)
+        send!!.completeResumeSend()
         // since we've just sent an element, we might need to resume some receivers
         checkSubOffers()
         // tailrec call to recheck
@@ -229,9 +229,9 @@ internal class ArrayBroadcastChannel<E>(
                 // it means that `checkOffer` must be retried after every `unlock`
                 if (!subLock.tryLock()) break
                 val receive: ReceiveOrClosed<E>?
-                val token: Any?
+                var result: Any?
                 try {
-                    val result = peekUnderLock()
+                    result = peekUnderLock()
                     when {
                         result === POLL_FAILED -> continue@loop // must retest `needsToCheckOfferWithoutLock` outside of the lock
                         result is Closed<*> -> {
@@ -242,15 +242,15 @@ internal class ArrayBroadcastChannel<E>(
                     // find a receiver for an element
                     receive = takeFirstReceiveOrPeekClosed() ?: break // break when no one's receiving
                     if (receive is Closed<*>) break // noting more to do if this sub already closed
-                    token = receive.tryResumeReceive(result as E, null)
-                    if (token == null) continue // bail out here to next iteration (see for next receiver)
+                    val token = receive.tryResumeReceive(result as E, null) ?: continue
+                    assert { token === RESUME_TOKEN }
                     val subHead = this.subHead
                     this.subHead = subHead + 1 // retrieved element for this subscriber
                     updated = true
                 } finally {
                     subLock.unlock()
                 }
-                receive!!.completeResumeReceive(token!!)
+                receive!!.completeResumeReceive(result as E)
             }
             // do close outside of lock if needed
             closed?.also { close(cause = it.closeCause) }