use withSerialContext to run coroutines serially on a multithreaded dispatcher

SerialContext.kt

/*
Copyright 2019 Matthew D.G. Timmermans

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package com.nobigsoftware.eventless.coroutines

import kotlinx.coroutines.CoroutineDispatcher
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Runnable
import kotlinx.coroutines.withContext
import java.util.concurrent.ConcurrentLinkedQueue
import java.util.concurrent.atomic.AtomicBoolean
import java.util.concurrent.atomic.AtomicReference
import kotlin.coroutines.CoroutineContext
import kotlin.coroutines.EmptyCoroutineContext

suspend fun <T> withSerialContext(
        context: CoroutineDispatcher,
        block: suspend CoroutineScope.() -> T
): T = withContext(SerialContextDispatcher(context), block)


/**
 * A `SerialContextDispatcher` is a dispatcher wrapper that runs all of its tasks one at a time.  This
 * allows clients to use multithreaded dispatchers without (maybe unintentionally) multithreading
 * their coroutines.
 *
 * @param target The dispatcher used to run tasks
 */
private class SerialContextDispatcher(private val target: CoroutineDispatcher) : CoroutineDispatcher() {

    companion object {
        private val EMPTY_PENDING = Object()
    }

    /**
     * An AtomicReference pointing to the state of our task queue.  To avoid an unnecessary allocation,
     * this doubles as the [Runnable] we schedule in the target dispatcher to run our tasks.
     *
     * The reference contains:
     * - null if there are no scheduled tasks and the state Runnable is not queued in [target]
     * - EMPTY_PENDING if there are no scheduled tasks and the state Runnable is queued in [target]
     * - a [Runnable] if there is one scheduled task and the state Runnable is queued in [target]
     * - A [SerialContextNode] if there is more than one task scheduled and the state Runnable is
     *      queued in [target].  [SerialContextNodes] are linked into a stack terminated by the first
     *      task [Runnable] to run
     * Whoever CASes the ref from null to non-null MUST queue the state Runnable in [target]
     */
    private val state = object : AtomicReference<Any?>(null), Runnable {

        // When tasks are pulled off of the [state] queue, they are turned into a forwrd linked
        // list here.  This is only accesed by one thread at a time, so no other CC is needed
        var taskList: Any? = null

        override fun run() {
            while(true) {
                var tl = taskList
                if (tl != null) {
                    if (tl is SerialContextNode) {
                        taskList = tl.next
                        tl = tl.task
                    } else {
                        taskList = null
                    }
                    if (tl is Runnable) {
                        try {
                            tl.run()
                        } catch (e: Throwable) {
                            try {
                                target.dispatch(EmptyCoroutineContext, Thrower(e))
                            } catch (e2: Throwable) {
                            }
                        }
                    }
                    continue
                }
                tl = getAndSet(EMPTY_PENDING)
                if (tl === EMPTY_PENDING || tl === null) {
                    if (compareAndSet(EMPTY_PENDING, null)) {
                        break
                    }
                    continue
                }
                //reverse stack onto taskList to create forward-linked list
                if (tl is SerialContextNode) {
                    val freeNode = tl
                    taskList = tl.task
                    tl = tl.next
                    while(tl is SerialContextNode) {
                        val ntl = tl.next
                        tl.next = taskList
                        taskList = tl
                        tl = ntl
                    }
                    if (tl is Runnable) {
                        freeNode.next = taskList
                        freeNode.task = tl
                        taskList = freeNode
                    }
                } else if (tl is Runnable) {
                    taskList = tl
                }
            }
        }
    }

    override fun dispatch(context: CoroutineContext, block: Runnable) {
        var node : SerialContextNode? = null
        @Suppress("SENSELESS_COMPARISON")
        if (block == null) {
            return
        }
        var oldState : Any?
        var newState : Any?
        do {
            oldState = state.get()
            newState = when {
                oldState === null -> block
                oldState === EMPTY_PENDING -> block
                node !== null -> {
                    node.next = oldState
                    node
                }
                else -> {
                    node = SerialContextNode(block, oldState)
                    node
                }
            }
        } while(!state.compareAndSet(oldState, newState))
        if (oldState === null) {
            target.dispatch(EmptyCoroutineContext, state)
        }
    }
}

private data class SerialContextNode(@JvmField var task: Runnable, @JvmField var next: Any?)

private class Thrower(private val e: Throwable) : Runnable {
    override fun run() {
        throw e
    }
}