Avoid allocation of CoroutineScheduler task object

* DispatchedContinuation and CancellableContinuationImpl have all
  the necessarily field to be scheduled

Author: elizarov

common/kotlinx-coroutines-core-common/src/AbstractContinuation.kt

@@ -19,7 +19,7 @@ private const val RESUMED = 2
 internal abstract class AbstractContinuation<in T>(
     public final override val delegate: Continuation<T>,
     public final override val resumeMode: Int
-) : Continuation<T>, DispatchedTask<T> {
+) : SchedulerTaskBase(), Continuation<T>, DispatchedTask<T> {
 
     /*
      * Implementation notes

common/kotlinx-coroutines-core-common/src/Dispatched.kt

@@ -81,7 +81,7 @@ internal object UndispatchedEventLoop {
 internal class DispatchedContinuation<in T>(
     @JvmField val dispatcher: CoroutineDispatcher,
     @JvmField val continuation: Continuation<T>
-) : Continuation<T> by continuation, DispatchedTask<T> {
+) : SchedulerTaskBase(), Continuation<T> by continuation, DispatchedTask<T> {
     @JvmField
     @Suppress("PropertyName")
     internal var _state: Any? = UNDEFINED
@@ -204,7 +204,7 @@ internal fun <T> Continuation<T>.resumeDirectWithException(exception: Throwable)
     else -> resumeWithException(exception)
 }
 
-internal interface DispatchedTask<in T> : Runnable {
+internal interface DispatchedTask<in T> : SchedulerTask {
     public val delegate: Continuation<T>
     public val resumeMode: Int get() = MODE_CANCELLABLE
 
@@ -237,6 +237,8 @@ internal interface DispatchedTask<in T> : Runnable {
             }
         } catch (e: Throwable) {
             throw DispatchException("Unexpected exception running $this", e)
+        } finally {
+            afterTask()
         }
     }
 }

common/kotlinx-coroutines-core-common/src/SchedulerTask.common.kt

@@ -0,0 +1,11 @@
+/*
+ * Copyright 2016-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.coroutines
+
+internal expect interface SchedulerTask : Runnable 
+
+internal expect abstract class SchedulerTaskBase() : SchedulerTask
+
+internal expect inline fun SchedulerTask.afterTask()

core/kotlinx-coroutines-core/src/SchedulerTask.kt

@@ -0,0 +1,17 @@
+/*
+ * Copyright 2016-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.coroutines
+
+import kotlinx.coroutines.scheduling.*
+
+internal actual typealias SchedulerTask = Task
+
+internal actual abstract class SchedulerTaskBase actual constructor() : SchedulerTask {
+    override var submissionTime: Long = 0
+    override var taskContext: TaskContext = NonBlockingContext
+}
+
+@Suppress("NOTHING_TO_INLINE")
+internal actual inline fun SchedulerTask.afterTask() = taskContext.afterTask()

core/kotlinx-coroutines-core/src/scheduling/CoroutineScheduler.kt

@@ -340,7 +340,6 @@ internal class CoroutineScheduler(
      */
     fun dispatch(block: Runnable, taskContext: TaskContext = NonBlockingContext, fair: Boolean = false) {
         timeSource.trackTask() // this is needed for virtual time support
-        // TODO at some point make DispatchTask extend Task and make its field settable to save an allocation
         val task = createTask(block, taskContext)
         // try to submit the task to the local queue and act depending on the result
         when (submitToLocalQueue(task, fair)) {
@@ -357,7 +356,15 @@ internal class CoroutineScheduler(
         }
     }
 
-    internal fun createTask(block: Runnable, taskContext: TaskContext) = Task(block, schedulerTimeSource.nanoTime(), taskContext)
+    internal fun createTask(block: Runnable, taskContext: TaskContext): Task {
+        val nanoTime = schedulerTimeSource.nanoTime()
+        if (block is Task) {
+            block.submissionTime = nanoTime
+            block.taskContext = taskContext
+            return block
+        }
+        return TaskImpl(block, nanoTime, taskContext)
+    }
 
     /**
      * Unparks or creates a new [Worker] for executing non-blocking tasks if there are idle cores
@@ -724,20 +731,22 @@ internal class CoroutineScheduler(
                     }
                     wasIdle = true
                 } else {
+                    // Note: read task.mode before running the task, because Task object will be reused after run
+                    val taskMode = task.mode
                     if (wasIdle) {
-                        idleReset(task.mode)
+                        idleReset(taskMode)
                         wasIdle = false
                     }
-                    beforeTask(task)
+                    beforeTask(taskMode, task.submissionTime)
                     runSafely(task)
-                    afterTask(task)
+                    afterTask(taskMode)
                 }
             }
             tryReleaseCpu(WorkerState.TERMINATED)
         }
 
-        private fun beforeTask(task: Task) {
-            if (task.mode != TaskMode.NON_BLOCKING) {
+        private fun beforeTask(taskMode: TaskMode, taskSubmissionTime: Long) {
+            if (taskMode != TaskMode.NON_BLOCKING) {
                 /*
                  * We should release CPU *before* checking for CPU starvation,
                  * otherwise requestCpuWorker() will not count current thread as blocking
@@ -756,16 +765,16 @@ internal class CoroutineScheduler(
                 return
             }
             val now = schedulerTimeSource.nanoTime()
-            if (now - task.submissionTime >= WORK_STEALING_TIME_RESOLUTION_NS &&
+            if (now - taskSubmissionTime >= WORK_STEALING_TIME_RESOLUTION_NS &&
                 now - lastExhaustionTime >= WORK_STEALING_TIME_RESOLUTION_NS * 5
             ) {
                 lastExhaustionTime = now
                 requestCpuWorker()
             }
         }
 
-        private fun afterTask(task: Task) {
-            if (task.mode != TaskMode.NON_BLOCKING) {
+        private fun afterTask(taskMode: TaskMode) {
+            if (taskMode != TaskMode.NON_BLOCKING) {
                 decrementBlockingWorkers()
                 val currentState = state
                 // Shutdown sequence of blocking dispatcher

core/kotlinx-coroutines-core/src/scheduling/Tasks.kt

@@ -82,18 +82,22 @@ internal object NonBlockingContext : TaskContext {
     }
 }
 
-internal class Task(
-    @JvmField val block: Runnable,
-    @JvmField val submissionTime: Long,
-    @JvmField val taskContext: TaskContext
-) : Runnable {
+internal interface Task : Runnable {
+    var submissionTime: Long
+    var taskContext: TaskContext
     val mode: TaskMode get() = taskContext.taskMode
-    
+}
+
+internal class TaskImpl(
+    @JvmField val block: Runnable,
+    override var submissionTime: Long,
+    override var taskContext: TaskContext
+) : SchedulerTaskBase() {
     override fun run() {
         try {
             block.run()
         } finally {
-            taskContext.afterTask()
+            afterTask()
         }
     }
 

core/kotlinx-coroutines-core/test/scheduling/WorkQueueTest.kt

@@ -130,7 +130,7 @@ internal fun GlobalQueue.asTimeList(): List<Long> {
     return result
 }
 
-internal fun task(n: Long) = Task(Runnable {}, n, NonBlockingContext)
+internal fun task(n: Long) = TaskImpl(Runnable {}, n, NonBlockingContext)
 
 internal fun WorkQueue.drain(): List<Long> {
     var task: Task? = poll()

js/kotlinx-coroutines-core-js/src/SchedulerTask.kt

@@ -0,0 +1,14 @@
+/*
+ * Copyright 2016-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.coroutines
+
+@Suppress("ACTUAL_WITHOUT_EXPECT")
+internal actual typealias SchedulerTask = Runnable
+
+internal actual abstract class SchedulerTaskBase actual constructor() : SchedulerTask
+
+@Suppress("NOTHING_TO_INLINE")
+internal actual inline fun SchedulerTask.afterTask() {}
+

native/kotlinx-coroutines-core-native/src/SchedulerTask.kt

@@ -0,0 +1,13 @@
+/*
+ * Copyright 2016-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.coroutines
+
+@Suppress("ACTUAL_WITHOUT_EXPECT")
+internal actual typealias SchedulerTask = Runnable
+
+internal actual abstract class SchedulerTaskBase actual constructor() : SchedulerTask
+
+@Suppress("NOTHING_TO_INLINE")
+internal actual inline fun SchedulerTask.afterTask() {}