MVP of new coroutine scheduler

Author: qwwdfsad

benchmarks/src/jmh/kotlin/benchmarks/ParametrizedDispatcherBase.kt

@@ -1,8 +1,10 @@
 package benchmarks
 
+import benchmarks.actors.CORES_COUNT
 import kotlinx.coroutines.experimental.CommonPool
 import kotlinx.coroutines.experimental.ThreadPoolDispatcher
 import kotlinx.coroutines.experimental.newFixedThreadPoolContext
+import kotlinx.coroutines.experimental.scheduling.ExperimentalCoroutineDispatcher
 import org.openjdk.jmh.annotations.Param
 import org.openjdk.jmh.annotations.Setup
 import org.openjdk.jmh.annotations.TearDown
@@ -23,6 +25,9 @@ abstract class ParametrizedDispatcherBase {
     open fun setup() {
         benchmarkContext = when {
             dispatcher == "fjp" -> CommonPool
+            dispatcher == "experimental" -> {
+                ExperimentalCoroutineDispatcher(CORES_COUNT)
+            }
             dispatcher.startsWith("ftp") -> {
                 newFixedThreadPoolContext(dispatcher.substring(4).toInt(), dispatcher).also { closeable = it }
             }

core/kotlinx-coroutines-core/src/main/kotlin/kotlinx/coroutines/experimental/scheduling/CoroutineScheduler.kt

@@ -0,0 +1,107 @@
+package kotlinx.coroutines.experimental.scheduling
+
+import kotlinx.coroutines.experimental.Runnable
+import java.io.Closeable
+import java.util.*
+import java.util.concurrent.ConcurrentLinkedQueue
+import java.util.concurrent.Executor
+import java.util.concurrent.locks.LockSupport
+
+/**
+ * TODO design rationale
+ */
+class CoroutineScheduler(private val corePoolSize: Int) : Executor, Closeable {
+
+    private val workers: Array<PoolWorker>
+    private val globalWorkQueue: Queue<Task> = ConcurrentLinkedQueue<Task>()
+    @Volatile
+    private var isClosed = false
+
+    init {
+        require(corePoolSize >= 1, { "Expected positive core pool size, but was $corePoolSize" })
+        workers = Array(corePoolSize, { PoolWorker(it) })
+        workers.forEach { it.start() }
+    }
+
+    override fun execute(command: Runnable) = dispatch(command)
+
+    override fun close() {
+        isClosed = true
+    }
+
+    fun dispatch(command: Runnable, intensive: Boolean = false) {
+        val task = TimedTask(System.nanoTime(), command)
+        if (!submitToLocalQueue(task, intensive)) {
+            globalWorkQueue.add(task)
+        }
+    }
+
+    private fun submitToLocalQueue(task: Task, intensive: Boolean): Boolean {
+        val worker = Thread.currentThread() as? PoolWorker ?: return false
+        if (intensive && worker.localQueue.bufferSize > FORKED_TASK_OFFLOAD_THRESHOLD) return false
+        worker.localQueue.offer(task, globalWorkQueue)
+        return true
+    }
+
+    private inner class PoolWorker(index: Int) : Thread("CoroutinesScheduler-worker-$index") {
+        init {
+            isDaemon = true
+        }
+
+        val localQueue: WorkQueue = WorkQueue()
+
+        @Volatile
+        var yields = 0
+
+        override fun run() {
+            while (!isClosed) {
+                try {
+                    val job = findTask()
+                    if (job == null) {
+                        awaitWork()
+                    } else {
+                        yields = 0
+                        job.task.run()
+                    }
+                } catch (e: Throwable) {
+                    println(e) // TODO handler
+                }
+            }
+        }
+
+        private fun awaitWork() {
+            // Temporary solution
+            if (++yields > 100000) {
+                LockSupport.parkNanos(WORK_STEALING_TIME_RESOLUTION / 2)
+            }
+        }
+
+        private fun findTask(): Task? {
+            // TODO explain, probabilistic check with park counter?
+            var task: Task? = globalWorkQueue.poll()
+            if (task != null) return task
+
+            task = localQueue.poll()
+            if (task != null) return task
+
+            return trySteal()
+        }
+
+        private fun trySteal(): Task? {
+            if (corePoolSize == 1) {
+                return null
+            }
+
+            while (true) {
+                val worker = workers[RANDOM_PROVIDER().nextInt(workers.size)]
+                if (worker !== this) {
+                    worker.localQueue.offloadWork(true) {
+                        localQueue.offer(it, globalWorkQueue)
+                    }
+
+                    return localQueue.poll()
+                }
+            }
+        }
+    }
+}

core/kotlinx-coroutines-core/src/main/kotlin/kotlinx/coroutines/experimental/scheduling/ExperimentalCoroutineDispatcher.kt

@@ -0,0 +1,31 @@
+package kotlinx.coroutines.experimental.scheduling
+
+import kotlinx.coroutines.experimental.*
+import java.io.Closeable
+import java.util.concurrent.TimeUnit
+import kotlin.coroutines.experimental.AbstractCoroutineContextElement
+import kotlin.coroutines.experimental.CoroutineContext
+
+/**
+ * Unstable API and subject to change.
+ * Context marker which gives scheduler a hint that submitted jobs can be distributed among cores aggressively.
+ * Usually it's useful for massive jobs submission produced by single coroutine, e.g. data intensive fork-join tasks
+ * or fan-out notifications for a large number of listeners.
+ */
+object ForkedMarker : AbstractCoroutineContextElement(ForkedKey)
+
+private object ForkedKey : CoroutineContext.Key<ForkedMarker>
+
+class ExperimentalCoroutineDispatcher(threads: Int = Runtime.getRuntime().availableProcessors()) : CoroutineDispatcher(), Delay, Closeable {
+
+    private val coroutineScheduler = CoroutineScheduler(threads)
+
+    override fun dispatch(context: CoroutineContext, block: Runnable) {
+        coroutineScheduler.dispatch(block, context[ForkedKey] != null)
+    }
+
+    override fun scheduleResumeAfterDelay(time: Long, unit: TimeUnit, continuation: CancellableContinuation<Unit>) =
+            DefaultExecutor.scheduleResumeAfterDelay(time, unit, continuation)
+
+    override fun close() = coroutineScheduler.close()
+}

core/kotlinx-coroutines-core/src/main/kotlin/kotlinx/coroutines/experimental/scheduling/RandomUtils.kt

@@ -0,0 +1,10 @@
+package kotlinx.coroutines.experimental.scheduling
+
+import java.util.*
+
+private val RANDOM = object : ThreadLocal<Random>() {
+    override fun initialValue() = Random()
+}
+
+// Dynamic discovery is not yet supported
+val RANDOM_PROVIDER = { RANDOM.get() }

core/kotlinx-coroutines-core/src/main/kotlin/kotlinx/coroutines/experimental/scheduling/Tasks.kt

@@ -0,0 +1,35 @@
+package kotlinx.coroutines.experimental.scheduling
+
+import java.util.*
+
+internal typealias Task = TimedTask
+internal typealias GlobalQueue = Queue<Task>
+
+internal val WORK_STEALING_TIME_RESOLUTION = readFromSystemProperties(
+        "kotlinx.coroutines.scheduler.resolution.us", 500L, String::toLongOrNull)
+
+internal val FORKED_TASK_OFFLOAD_THRESHOLD = readFromSystemProperties(
+        "kotlinx.coroutines.scheduler.fork.threshold", 64L, String::toLongOrNull)
+
+internal var schedulerTimeSource: TimeSource = NanoTimeSource
+
+internal data class TimedTask(val submissionTime: Long, val task: Runnable)
+
+internal abstract class TimeSource {
+    abstract fun nanoTime(): Long
+}
+
+internal object NanoTimeSource : TimeSource() {
+    override fun nanoTime() = System.nanoTime()
+}
+
+private fun <T> readFromSystemProperties(propertyName: String, defaultValue: T, parser: (String) -> T?): T {
+    val value = try {
+        System.getProperty(propertyName)
+    } catch (e: SecurityException) {
+        null
+    } ?: return defaultValue
+
+    val parsed = parser(value)
+    return parsed ?: error("System property '$propertyName' has unrecognized value '$value'")
+}

core/kotlinx-coroutines-core/src/main/kotlin/kotlinx/coroutines/experimental/scheduling/WorkQueue.kt

@@ -0,0 +1,136 @@
+package kotlinx.coroutines.experimental.scheduling
+
+import java.util.concurrent.atomic.AtomicInteger
+import java.util.concurrent.atomic.AtomicReference
+import java.util.concurrent.atomic.AtomicReferenceArray
+
+internal const val BUFFER_CAPACITY_BASE = 7
+internal const val BUFFER_CAPACITY = 1 shl BUFFER_CAPACITY_BASE
+internal const val MASK = BUFFER_CAPACITY - 1 // 128 by default
+
+/**
+ * Unstable API and subject to change.
+ * Tightly coupled with [CoroutineScheduler] queue of pending tasks, but extracted to separate file for simplicity.
+ * At any moment queue is used only by [CoroutineScheduler.PoolWorker] threads, has only one producer (worker owning this queue)
+ * and any amount of consumers, other pool workers which are trying to steal work.
+ *
+ * Fairness
+ * [WorkQueue] provides semi-FIFO order, but with priority for most recently submitted task assuming
+ * that these two (current and submitted) are communicating and sharing state thus making such communication extremely fast.
+ * E.g. submitted jobs [1, 2, 3, 4] will be executed in [4, 1, 2, 3] order.
+ *
+ * Work offloading
+ * When queue is full, half of existing tasks is offloaded to global queue which is regularly polled by other pool workers.
+ * Offloading occurs in LIFO order for the sake of implementation simplicity: offload should be extremely rare and occurs only in specific use-cases
+ * (e.g. when coroutine starts heavy fork-join-like computation), so fairness is not important.
+ * As an alternative, offloading directly to some [CoroutineScheduler.PoolWorker] may be used, but then strategy of selecting most idle worker
+ * should be implemented and implementation should be aware multiple producers.
+ */
+internal class WorkQueue {
+
+    internal val bufferSize: Int get() = producerIndex.get() - consumerIndex.get()
+    private val buffer: AtomicReferenceArray<Task?> = AtomicReferenceArray(BUFFER_CAPACITY)
+    private val lastScheduledTask: AtomicReference<Task?> = AtomicReference(null)
+
+    private val producerIndex: AtomicInteger = AtomicInteger(0)
+    private val consumerIndex: AtomicInteger = AtomicInteger(0)
+
+    /**
+     * Retrieves and removes task from head of the queue
+     * Invariant: this method is called only by owner of the queue ([pollExternal] is not)
+     */
+    fun poll(): Task? {
+        return lastScheduledTask.getAndSet(null) ?: pollExternal()
+    }
+
+    /**
+     * Invariant: this method is called only by owner of the queue
+     * @param task task to put into local queue
+     * @param globalQueue fallback queue which is used when local queue is overflown
+     */
+    fun offer(task: Task, globalQueue: GlobalQueue) {
+        while (true) {
+            val previous = lastScheduledTask.get()
+            if (lastScheduledTask.compareAndSet(previous, task)) {
+                if (previous != null) {
+                    addLast(previous, globalQueue)
+                }
+                return
+            }
+        }
+    }
+
+    /**
+     * Offloads half of the current buffer to [sink]
+     * @param byTimer whether task deadline should be checked before offloading
+     */
+    inline fun offloadWork(byTimer: Boolean, sink: (Task) -> Unit) {
+        repeat((bufferSize / 2).coerceAtLeast(1)) {
+            if (bufferSize == 0) { // try to steal head if buffer is empty
+                val lastScheduled = lastScheduledTask.get() ?: return
+                if (!byTimer || schedulerTimeSource.nanoTime() - lastScheduled.submissionTime < WORK_STEALING_TIME_RESOLUTION) {
+                    return
+                }
+
+                if (lastScheduledTask.compareAndSet(lastScheduled, null)) {
+                    sink(lastScheduled)
+                    return
+                }
+            }
+
+            // TODO use batch drain and (if target queue allows) batch insert
+            val task = pollExternal { !byTimer || schedulerTimeSource.nanoTime() - it.submissionTime >= WORK_STEALING_TIME_RESOLUTION }
+                    ?: return
+            sink(task)
+        }
+    }
+
+    /**
+     * [poll] for external (not owning this queue) workers
+     */
+    private inline fun pollExternal(predicate: (Task) -> Boolean = { true }): Task? {
+        while (true) {
+            val tailLocal = consumerIndex.get()
+            if (tailLocal - producerIndex.get() == 0) return null
+            val index = tailLocal and MASK
+            val element = buffer[index] ?: continue
+            if (!predicate(element)) {
+                return null
+            }
+
+            if (consumerIndex.compareAndSet(tailLocal, tailLocal + 1)) {
+                // 1) Help GC 2) Signal producer that this slot is consumed and may be used
+                return buffer.getAndSet(index, null)
+            }
+        }
+    }
+
+    // Called only by owner
+    private fun addLast(task: Task, globalQueue: GlobalQueue) {
+        while (!tryAddLast(task)) {
+            offloadWork(false) {
+                globalQueue.add(it)
+            }
+        }
+    }
+
+    // Called only by owner
+    private fun tryAddLast(task: Task): Boolean {
+        if (bufferSize == BUFFER_CAPACITY - 1) return false
+        val headLocal = producerIndex.get()
+        val nextIndex = headLocal and MASK
+
+        /*
+         * If current element is not null then we're racing with consumers for tail. If we skip this check then
+         * consumer can null out current element and it will be lost. If we're racing for tail then
+         * queue is close to overflow => it's fine to offload work to global queue
+         */
+        if (buffer[nextIndex] != null) {
+            return false
+        }
+
+        buffer.lazySet(nextIndex, task)
+        producerIndex.incrementAndGet()
+        return true
+    }
+}

core/kotlinx-coroutines-core/src/test/kotlin/kotlinx/coroutines/experimental/scheduling/CoroutineSchedulerStressTest.kt

@@ -0,0 +1,53 @@
+package kotlinx.coroutines.experimental.scheduling
+
+import kotlinx.coroutines.experimental.TestBase
+import org.junit.After
+import org.junit.Test
+import java.util.concurrent.ConcurrentHashMap
+import java.util.concurrent.CountDownLatch
+import java.util.concurrent.atomic.AtomicInteger
+import kotlin.coroutines.experimental.CoroutineContext
+import kotlin.coroutines.experimental.EmptyCoroutineContext
+import kotlin.test.assertEquals
+
+class CoroutineSchedulerStressTest : TestBase() {
+
+    private var dispatcher: ExperimentalCoroutineDispatcher = ExperimentalCoroutineDispatcher()
+    private val observedThreads = ConcurrentHashMap<Thread, MutableSet<Int>>()
+    private val tasksNum = 1_000_000
+    private val processed = AtomicInteger(0)
+
+    @After
+    fun tearDown() {
+        dispatcher.close()
+    }
+
+    @Test
+    fun submitTasks() {
+        stressTest(ForkedMarker)
+    }
+
+    @Test
+    fun submitTasksForked() {
+        stressTest(EmptyCoroutineContext)
+    }
+
+    private fun stressTest(ctx: CoroutineContext) {
+        val finishLatch = CountDownLatch(1)
+
+        for (i in 1..tasksNum) {
+            dispatcher.dispatch(ctx, Runnable {
+                val numbers = observedThreads.computeIfAbsent(Thread.currentThread(), { _ -> hashSetOf() })
+                require(numbers.add(i))
+                if (processed.incrementAndGet() == tasksNum) {
+                    finishLatch.countDown()
+                }
+            })
+        }
+
+        finishLatch.await()
+        assertEquals(Runtime.getRuntime().availableProcessors(), observedThreads.size)
+        val result = observedThreads.values.flatMap { it }.toSet()
+        assertEquals((1..tasksNum).toSet(), result)
+    }
+}