Almost non-blocking ConflatedBufferedChannel

Signed-off-by: Nikita Koval <[email protected]>

Author: ndkoval

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

@@ -69,7 +69,17 @@ internal open class ConflatedBufferedChannel<E>(
         return success(Unit)
     }
 
-    private fun trySendDropOldest(element: E): ChannelResult<Unit> =
+    private val closeStarted = atomic(false)
+    private val trySendStarted = atomic(1L)
+    private val trySendCompleted = atomic(1L)
+
+    private fun trySendDropOldest(element: E): ChannelResult<Unit> = try {
+        trySendStarted.incrementAndGet()
+        if (closeStarted.value) {
+            trySendCompleted.incrementAndGet()
+            while (!isClosedForSend) {}
+            trySendStarted.incrementAndGet()
+        }
         sendImpl( // <-- this is an inline function
             element = element,
             // Put the element into the logical buffer in any case,
@@ -90,6 +100,9 @@ internal open class ConflatedBufferedChannel<E>(
             // If the channel is closed, return the corresponding result.
             onClosed = { closed(sendException) }
         )
+    } finally {
+        trySendCompleted.incrementAndGet()
+    }
 
     @Suppress("UNCHECKED_CAST")
     override fun registerSelectForSend(select: SelectInstance<*>, element: Any?) {
@@ -109,4 +122,26 @@ internal open class ConflatedBufferedChannel<E>(
     }
 
     override fun shouldSendSuspend() = false // never suspends
+
+    override fun closeOrCancelImpl(cause: Throwable?, cancel: Boolean): Boolean {
+        // Inform `trySend(..)` the the channel closing procedure has been started.
+        // All `trySend(..)`s that start after setting  this flag, should wait
+        // until the channel is properly closed. Otherwise, a non-linearizable
+        // behaviour can be observed.
+        closeStarted.value = true
+        // It is critical to wait until all the started `trySend(..)`s
+        // complete their operation.
+        waitUntilStartedTrySendsComplete()
+        // Finally, close/cancel this channel.
+        return super.closeOrCancelImpl(cause, cancel)
+    }
+
+    private fun waitUntilStartedTrySendsComplete() {
+        while (true) {
+            val started = trySendStarted.value
+            val completed = trySendCompleted.value
+            if (started != trySendStarted.value) continue
+            if (started == completed) break
+        }
+    }
 }