Scheduler: capture TaskContext for afterTask before running it

This is not strictly needed, since the reusable Task instances
(DispatchedTask) should be always reused in the same context, but
this makes the code more resilient to future changes.

Author: elizarov

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

@@ -218,6 +218,7 @@ internal interface DispatchedTask<in T> : SchedulerTask {
         (state as? CompletedExceptionally)?.cause
 
     public override fun run() {
+        val taskContext = this.taskContext
         try {
             val delegate = delegate as DispatchedContinuation<T>
             val continuation = delegate.continuation
@@ -238,7 +239,7 @@ internal interface DispatchedTask<in T> : SchedulerTask {
         } catch (e: Throwable) {
             throw DispatchException("Unexpected exception running $this", e)
         } finally {
-            afterTask()
+            afterTask(taskContext)
         }
     }
 }

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

@@ -8,4 +8,8 @@ internal expect interface SchedulerTask : Runnable
 
 internal expect abstract class SchedulerTaskBase() : SchedulerTask
 
-internal expect inline fun SchedulerTask.afterTask()
+internal expect interface SchedulerTaskContext
+
+internal expect val SchedulerTask.taskContext: SchedulerTaskContext
+
+internal expect inline fun SchedulerTask.afterTask(taskContext: SchedulerTaskContext)

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

@@ -13,5 +13,11 @@ internal actual abstract class SchedulerTaskBase actual constructor() : Schedule
     override var taskContext: TaskContext = NonBlockingContext
 }
 
+internal actual typealias SchedulerTaskContext = TaskContext
+
+@Suppress("EXTENSION_SHADOWED_BY_MEMBER")
+internal actual val SchedulerTask.taskContext: SchedulerTaskContext get() = taskContext
+
 @Suppress("NOTHING_TO_INLINE")
-internal actual inline fun SchedulerTask.afterTask() = taskContext.afterTask()
+internal actual inline fun SchedulerTask.afterTask(taskContext: SchedulerTaskContext) =
+    taskContext.afterTask()

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

@@ -88,6 +88,7 @@ internal interface Task : Runnable {
     val mode: TaskMode get() = taskContext.taskMode
 }
 
+// Non-reusable Task implementation to wrap Runnable instances that do not otherwise implement task
 internal class TaskImpl(
     @JvmField val block: Runnable,
     override var submissionTime: Long,
@@ -97,7 +98,7 @@ internal class TaskImpl(
         try {
             block.run()
         } finally {
-            afterTask()
+            afterTask(taskContext)
         }
     }
 

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

@@ -9,6 +9,11 @@ internal actual typealias SchedulerTask = Runnable
 
 internal actual abstract class SchedulerTaskBase actual constructor() : SchedulerTask
 
+@Suppress("ACTUAL_WITHOUT_EXPECT")
+internal actual typealias SchedulerTaskContext = Unit
+
+internal actual val SchedulerTask.taskContext: SchedulerTaskContext get() = Unit
+
 @Suppress("NOTHING_TO_INLINE")
-internal actual inline fun SchedulerTask.afterTask() {}
+internal actual inline fun SchedulerTask.afterTask(taskContext: SchedulerTaskContext) {}
 

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

@@ -9,5 +9,10 @@ internal actual typealias SchedulerTask = Runnable
 
 internal actual abstract class SchedulerTaskBase actual constructor() : SchedulerTask
 
+@Suppress("ACTUAL_WITHOUT_EXPECT")
+internal actual typealias SchedulerTaskContext = Unit
+
+internal actual val SchedulerTask.taskContext: SchedulerTaskContext get() = kotlin.Unit
+
 @Suppress("NOTHING_TO_INLINE")
-internal actual inline fun SchedulerTask.afterTask() {}
+internal actual inline fun SchedulerTask.afterTask(taskContext: SchedulerTaskContext) {}