Dispose/cancel reactive subscription exactly once using atomic

Author: elizarov

reactive/kotlinx-coroutines-reactive/src/Channel.kt

@@ -41,8 +41,7 @@ private class SubscriptionChannel<T>(
         require(request >= 0) { "Invalid request size: $request" }
     }
 
-    @Volatile
-    private var subscription: Subscription? = null
+    private val _subscription = atomic<Subscription?>(null)
 
     // requested from subscription minus number of received minus number of enqueued receivers,
     // can be negative if we have receivers, but no subscription yet
@@ -52,7 +51,7 @@ private class SubscriptionChannel<T>(
     @Suppress("CANNOT_OVERRIDE_INVISIBLE_MEMBER")
     override fun onReceiveEnqueued() {
         _requested.loop { wasRequested ->
-            val subscription = this.subscription
+            val subscription = _subscription.value
             val needRequested = wasRequested - 1
             if (subscription != null && needRequested < 0) { // need to request more from subscription
                 // try to fixup by making request
@@ -73,13 +72,12 @@ private class SubscriptionChannel<T>(
 
     @Suppress("CANNOT_OVERRIDE_INVISIBLE_MEMBER")
     override fun onClosedIdempotent(closed: LockFreeLinkedListNode) {
-        subscription?.cancel()
-        subscription = null // optimization -- no need to cancel it again
+        _subscription.getAndSet(null)?.cancel() // cancel exactly once
     }
 
     // Subscriber overrides
     override fun onSubscribe(s: Subscription) {
-        subscription = s
+        _subscription.value = s
         while (true) { // lock-free loop on _requested
             if (isClosedForSend) {
                 s.cancel()

reactive/kotlinx-coroutines-rx2/src/RxChannel.kt

@@ -6,8 +6,9 @@ package kotlinx.coroutines.rx2
 
 import io.reactivex.*
 import io.reactivex.disposables.*
-import kotlinx.coroutines.channels.*
+import kotlinx.atomicfu.*
 import kotlinx.coroutines.*
+import kotlinx.coroutines.channels.*
 import kotlinx.coroutines.internal.*
 
 /**
@@ -56,18 +57,16 @@ public suspend inline fun <T> ObservableSource<T>.consumeEach(action: (T) -> Uni
 private class SubscriptionChannel<T> :
     LinkedListChannel<T>(), Observer<T>, MaybeObserver<T>
 {
-    @Volatile
-    var subscription: Disposable? = null
+    private val _subscription = atomic<Disposable?>(null)
 
     @Suppress("CANNOT_OVERRIDE_INVISIBLE_MEMBER")
     override fun onClosedIdempotent(closed: LockFreeLinkedListNode) {
-        subscription?.dispose()
-        subscription = null // optimization -- no need to dispose it again
+        _subscription.getAndSet(null)?.dispose() // dispose exactly once
     }
 
     // Observer overrider
     override fun onSubscribe(sub: Disposable) {
-        subscription = sub
+        _subscription.value = sub
     }
 
     override fun onSuccess(t: T) {