Draft optimized async iterator via CoroutineIntrinsics

Author: elizarov

kotlinx-coroutines-rx/src/main/kotlin/asyncRxIterator.kt

@@ -4,9 +4,14 @@ import rx.Observable
 import rx.Subscriber
 import java.util.concurrent.atomic.AtomicReference
 import kotlin.coroutines.Continuation
-import kotlin.coroutines.suspendCoroutine
+import kotlin.coroutines.CoroutineIntrinsics.SUSPENDED
+import kotlin.coroutines.CoroutineIntrinsics.suspendCoroutineOrReturn
 
-// supports suspending iteration on observables
+/**
+ * Suspending iteration extension. It does not have its own buffer and applies back-pressure to the source.
+ * If iterating coroutine does not have a dispatcher with its own thread, then the iterating coroutine
+ * is resumed in the thread that invokes [Subscriber.onNext].
+ */
 suspend operator fun <V : Any> Observable<V>.iterator(): ObserverIterator<V> {
     val iterator = ObserverIterator<V>()
     subscribe(iterator.subscriber)
@@ -17,101 +22,45 @@ private sealed class Waiter<in T>(val c: Continuation<T>)
 private class HasNextWaiter(c: Continuation<Boolean>) : Waiter<Boolean>(c)
 private class NextWaiter<V>(c: Continuation<V>) : Waiter<V>(c)
 
-private object Complete
-private class Fail(val e: Throwable)
+private object Completed
+private class CompletedWith(val v: Any)
+private class Error(val e: Throwable)
 
 class ObserverIterator<V : Any> {
     internal val subscriber = Sub()
-    // Contains either null, Complete, Fail(exception), Waiter, or next value
+    // Contains either null, Completed, CompletedWith, Error(exception), Waiter, or next value
     private val rendezvous = AtomicReference<Any?>()
 
     @Suppress("UNCHECKED_CAST")
-    private suspend fun awaitHasNext(): Boolean = suspendCoroutine sc@ { c ->
+    suspend operator fun hasNext(): Boolean = suspendCoroutineOrReturn sc@ { c ->
         while (true) { // lock-free loop for rendezvous
             val cur = rendezvous.get()
             when (cur) {
-                null -> {
-                    if (rendezvous.compareAndSet(null, HasNextWaiter(c))) return@sc
-                }
-                Complete -> {
-                    c.resume(false)
-                    return@sc
-                }
-                is Fail -> {
-                    c.resumeWithException(cur.e)
-                    return@sc
-                }
-                is Waiter<*> -> {
-                    c.resumeWithException(IllegalStateException("Concurrent iteration"))
-                    return@sc
-                }
-                else -> {
-                    c.resume(true)
-                    return@sc
-                }
-            }
-        }
-    }
-
-    private suspend fun awaitNext(): V = suspendCoroutine sc@ { c ->
-        while (true) { // lock-free loop for rendezvous
-            val cur = rendezvous.get()
-            when (cur) {
-                null -> {
-                    if (rendezvous.compareAndSet(null, NextWaiter(c))) return@sc
-                }
-                Complete -> {
-                    c.resumeWithException(NoSuchElementException())
-                    return@sc
-                }
-                is Fail -> {
-                    c.resumeWithException(cur.e)
-                    return@sc
-                }
-                is Waiter<*> -> {
-                    c.resumeWithException(IllegalStateException("Concurrent iteration"))
-                    return@sc
-                }
-                else -> {
-                    if (rendezvous.compareAndSet(cur, null)) {
-                        c.resume(consumeValue(cur))
-                        return@sc
-                    }
-                }
+                null -> if (rendezvous.compareAndSet(null, HasNextWaiter(c))) return@sc SUSPENDED
+                Completed -> return@sc false
+                is CompletedWith -> return@sc true
+                is Error -> throw cur.e
+                is Waiter<*> -> throw IllegalStateException("Concurrent iteration")
+                else -> return@sc true
             }
         }
     }
 
-    suspend operator fun hasNext(): Boolean {
-        val cur = rendezvous.get()
-        return when (cur) {
-            null ->  awaitHasNext()
-            Complete -> false
-            is Fail -> throw cur.e
-            is Waiter<*> -> throw IllegalStateException("Concurrent iteration")
-            else -> true
-        }
-    }
-
-    suspend operator fun next(): V {
+    @Suppress("UNCHECKED_CAST")
+    suspend operator fun next(): V = suspendCoroutineOrReturn sc@ { c ->
         while (true) { // lock-free loop for rendezvous
             val cur = rendezvous.get()
             when (cur) {
-                null ->  return awaitNext()
-                Complete -> throw NoSuchElementException()
-                is Fail -> throw cur.e
+                null -> if (rendezvous.compareAndSet(null, NextWaiter(c))) return@sc SUSPENDED
+                Completed -> throw NoSuchElementException()
+                is CompletedWith -> if (rendezvous.compareAndSet(cur, Completed)) return@sc cur.v as V
+                is Error -> throw cur.e
                 is Waiter<*> -> throw IllegalStateException("Concurrent iteration")
-                else -> if (rendezvous.compareAndSet(cur, null)) return consumeValue(cur)
+                else -> if (rendezvous.compareAndSet(cur, null)) return (cur as V).apply { subscriber.requestOne() }
             }
         }
     }
 
-    @Suppress("UNCHECKED_CAST")
-    private fun consumeValue(cur: Any?): V {
-        subscriber.requestOne()
-        return cur as V
-    }
-
     internal inner class Sub : Subscriber<V>() {
         fun requestOne() {
             request(1)
@@ -125,9 +74,10 @@ class ObserverIterator<V : Any> {
             while (true) { // lock-free loop for rendezvous
                 val cur = rendezvous.get()
                 when (cur) {
-                    null ->  if (rendezvous.compareAndSet(null, Fail(e))) return
-                    Complete -> throw IllegalStateException("onError after onComplete")
-                    is Fail -> throw IllegalStateException("onError after onError")
+                    null ->  if (rendezvous.compareAndSet(null, Error(e))) return
+                    Completed -> throw IllegalStateException("onError after onCompleted")
+                    is CompletedWith -> throw IllegalStateException("onError after onCompleted")
+                    is Error -> throw IllegalStateException("onError after onError")
                     is Waiter<*> -> if (rendezvous.compareAndSet(cur, null)) {
                         cur.c.resumeWithException(e)
                         return
@@ -143,8 +93,9 @@ class ObserverIterator<V : Any> {
                 val cur = rendezvous.get()
                 when (cur) {
                     null ->  if (rendezvous.compareAndSet(null, v)) return
-                    Complete -> throw IllegalStateException("onNext after onComplete")
-                    is Fail -> throw IllegalStateException("onNext after onError")
+                    Completed -> throw IllegalStateException("onNext after onCompleted")
+                    is CompletedWith -> throw IllegalStateException("onNext after onCompleted")
+                    is Error -> throw IllegalStateException("onNext after onError")
                     is HasNextWaiter -> if (rendezvous.compareAndSet(cur, v)) {
                         cur.c.resume(true)
                         return
@@ -162,18 +113,19 @@ class ObserverIterator<V : Any> {
             while (true) { // lock-free loop for rendezvous
                 val cur = rendezvous.get()
                 when (cur) {
-                    null ->  if (rendezvous.compareAndSet(null, Complete)) return
-                    Complete -> throw IllegalStateException("onComplete after onComplete")
-                    is Fail -> throw IllegalStateException("onComplete after onError")
-                    is HasNextWaiter -> if (rendezvous.compareAndSet(cur, Complete)) {
+                    null ->  if (rendezvous.compareAndSet(null, Completed)) return
+                    Completed -> throw IllegalStateException("onCompleted after onCompleted")
+                    is CompletedWith -> throw IllegalStateException("onCompleted after onCompleted")
+                    is Error -> throw IllegalStateException("onCompleted after onError")
+                    is HasNextWaiter -> if (rendezvous.compareAndSet(cur, Completed)) {
                         cur.c.resume(false)
                         return
                     }
-                    is NextWaiter<*> -> if (rendezvous.compareAndSet(cur, Complete)) {
+                    is NextWaiter<*> -> if (rendezvous.compareAndSet(cur, Completed)) {
                         cur.c.resumeWithException(NoSuchElementException())
                         return
                     }
-                    else -> throw IllegalStateException("onComplete after onNext before request(1) was called")
+                    else -> if (rendezvous.compareAndSet(cur, CompletedWith(cur))) return
                 }
             }
         }

kotlinx-coroutines-rx/src/test/kotlin/AsyncRxTest.kt

@@ -130,17 +130,34 @@ class AsyncRxTest {
     @Test
     fun testAsyncIterator() {
         val observable = asyncRx {
-            val sb = StringBuilder()
-            for (s in Observable.just("O", "K"))
-                sb.append(s)
-            sb.toString()
+            var result = ""
+            for (s in Observable.just("O", "K")) {
+                result += s
+            }
+            result
         }
 
         checkObservableWithSingleValue(observable) {
             assertEquals("OK", it)
         }
     }
 
+    @Test
+    fun testAsyncIteratorException() {
+        val observable = asyncRx {
+            var result = ""
+            for (s in Observable.error<String>(RuntimeException("OK"))) {
+                result += s
+            }
+            result
+        }
+
+        checkErroneousObservable(observable) {
+            assert(it is RuntimeException)
+            assertEquals("OK", it.message)
+        }
+    }
+
     private fun checkErroneousObservable(
             observable: Observable<*>,
             checker: (Throwable) -> Unit