RxJava2: “TakeWhenObservable”

Updated: Jun 28, 2018

The code snippet is based on RxJava 2.1.13

RxJava2 has many convenient operators, where one of my favorite is takeUntil.

The takeUntil operator essentially says: “I am interested in this stream/Observable until something happens on that stream/Observable”.

It leads me to an interesting question, what if I want it temporarily to throttle the stream/Observable: Collect the items emitted during the filter is on and dump the collection when the filter is off? Which looks like this:

like a transistor with buffer

The Rx marble diagram is:

Therefore, I create TakeWhenObservable by referring to zip operator. The usage is quite simple,

TakeWhenObservable(src = $YOUR_SOURCE_OBSERVABLE,                   whenSrc = $YOUR_RELAY_SIGNAL,                   bufferSize = $AS_IT_DESCRIBED)

Internally, this TakeWhenObservable Observable looks like:

Internal look

(Optional) This is the full code that you could directly use:

/**
 * The Observable that emits the items emitted by the [src] Observable when a
 * second ObservableSource, [whenSrc], emits an true; It buffers all them items
 * emitted by the [src] when [whenSrc] emits an false.
 */
open class TakeWhenObservable<T>(private val src: ObservableSource<T>,
                                 private val whenSrc: ObservableSource<Boolean>,
                                 private val bufferSize: Int = Flowable.bufferSize())
    : Observable<T>() {

    override fun subscribeActual(observer: Observer<in T>) {
        val coordinator = Coordinator(src = src,
                                      whenSrc = whenSrc,
                                      actualObserver = observer,
                                      bufferSize = bufferSize)
        coordinator.subscribe()
    }

    internal class Coordinator<T>(val src: ObservableSource<T>,
                                  val whenSrc: ObservableSource<Boolean>,
                                  val actualObserver: Observer<in T>,
                                  val bufferSize: Int)
        : AtomicBoolean(false),
          Disposable {

        @Volatile
        private var canDrain = false
        private val queue = SpscLinkedArrayQueue<T>(bufferSize)

        private val srcObserver = SrcObserver(parent = this)
        private val whenObserver = WhenObserver(parent = this)

        fun subscribe() {
            // Establish a connections between this coordinator and the external
            // observer.
            actualObserver.onSubscribe(this)

            // Establish two connections between this coordinator and the two
            // inner observers observing the two external input signals.
            src.subscribe(srcObserver)
            whenSrc.subscribe(whenObserver)
        }

        fun bufferIt(t: T) {
            queue.offer(t)
            // Drop item when size is over the buffer size
            while (queue.size() > bufferSize) {
                queue.poll()
            }
        }

        fun canDrain(t: Boolean) {
            canDrain = t
        }

        fun drain() {
            synchronized(this) {
                while (!isDisposed &&
                       canDrain &&
                       !srcObserver.done &&
                       !queue.isEmpty) {
                    val t = queue.poll()!!
                    this.actualObserver.onNext(t)
                }
            }
        }

        fun onComplete() {
            if (!srcObserver.done) {
                actualObserver.onComplete()
            }
        }

        fun onError(err: Throwable) {
            actualObserver.onError(err)
        }

        override fun isDisposed(): Boolean {
            return get()
        }

        override fun dispose() {
            // Mark disposed!
            set(true)

            srcObserver.dispose()
            whenObserver.dispose()

            synchronized(this) {
                queue.clear()
            }
        }
    }

    internal class SrcObserver<T>(val parent: Coordinator<in T>)
        : Observer<T> {

        @Volatile
        var done = false

        val disposable = AtomicReference<Disposable>()

        override fun onComplete() {
            parent.onComplete()
            done = true
        }

        override fun onSubscribe(d: Disposable) {
            DisposableHelper.setOnce(disposable, d)
        }

        override fun onNext(t: T) {
            if (!done) {
                parent.bufferIt(t)
                parent.drain()
            }
        }

        override fun onError(e: Throwable) {
            if (!done) {
                parent.onError(e)
            }
        }

        fun dispose() {
            DisposableHelper.dispose(disposable)
        }
    }

    internal class WhenObserver<T>(val parent: Coordinator<in T>)
        : Observer<Boolean> {

        @Volatile
        var done = false

        val disposable = AtomicReference<Disposable>()

        override fun onComplete() {
            // Close the throttle permanently
            parent.canDrain(false)
        }

        override fun onSubscribe(d: Disposable) {
            DisposableHelper.setOnce(disposable, d)
        }

        override fun onNext(t: Boolean) {
            if (!done) {
                parent.canDrain(t)
                parent.drain()
            }
        }

        override fun onError(e: Throwable) {
            if (!done) {
                parent.onError(e)
            }
        }

        fun dispose() {
            DisposableHelper.dispose(disposable)
        }
    }
}

(Optional) Unit test: link.

// Activity or Fragment

private val mResumeSignal = PublishSubject.create<Boolean>()
private val mResultSignal = PublishSubject.create<Any>()

override void onResume() {
    super.onResume();
    mResumeSignal.onNext(true);
}

override void onPause() {
    super.onPause();
    mResumeSignal.onNext(false);
}

override void onActivityResult(requestCode: Int, 
                               resultCode: Int, 
                               data: Intent) {
    mResultSignal.onNext(...);
}

// Postpone the moment of sending the result to your Presenter or domain 
// component regardless of handling save-restore
fun onGetResult(): Observable<Any> {
    return TakeWhenObservable(src = mResultSignal,
                              whenSrc = mResumeSignal,
                              bufferSize = 1)
}