pausable dispatcher for kotlin coroutines

pausable_coroutines_dispatcher.kt

import kotlinx.coroutines.*
import java.util.concurrent.LinkedBlockingQueue
import java.util.concurrent.atomic.AtomicInteger
import kotlin.coroutines.CoroutineContext

fun main() {
    val dispatcherMgr = PauseDispatcherMgr()

    GlobalScope.launch {
        withContext(dispatcherMgr.pausableDispatcher(Dispatchers.IO)) {
            println("running 1")
        }

        withContext(dispatcherMgr.pausableDispatcher(Dispatchers.Default)) {
            println("running 2")
        }
    }

    dispatcherMgr.dispatcher.pause()
    GlobalScope.launch {
        delay(500)
        dispatcherMgr.dispatcher.resume()
    }

    Thread.sleep(3000)
}

class PauseDispatcherMgr {
    val dispatcher get() = pausableDispatcher

    private val pausableDispatcher by lazy { PausableDispatcher() }

    fun pausableDispatcher(original: CoroutineDispatcher) = pausableDispatcher.apply {
        this.original = original
    }
}

class PausableDispatcher : CoroutineDispatcher() {

    var original: CoroutineDispatcher = Dispatchers.Default

    private var count = AtomicInteger()

    @Volatile
    private var isPaused = false
    private val pausedQueue = LinkedBlockingQueue<Task>()

    override fun dispatch(context: CoroutineContext, block: Runnable) {
        val task = Task(count.incrementAndGet(), context, block, original)
        if (isPaused) {
            pausedQueue.offer(task)
        } else {
            task.dispatch()
        }
    }

    fun pause() {
        isPaused = true
    }

    fun resume() {
        isPaused = false
        val iterator = pausedQueue.iterator()
        while (iterator.hasNext()) {
            val task = iterator.next()
            iterator.remove()
            task.dispatch()
        }
    }

    inner class Task(
            val id: Int,
            private val context: CoroutineContext,
            private val runnable: Runnable,
            private val original: CoroutineDispatcher
    ) {
        fun dispatch() {
            println("$original dispatch $id")
            original.dispatch(context, runnable)
        }
    }
}