From 10a0628c594dac1089d0100030fd9dcef5ea52cb Mon Sep 17 00:00:00 2001
From: Daniil Ovchinnikov <daniil.ovchinnikov@jetbrains.com>
Date: Mon, 13 Mar 2023 15:45:21 +0100
Subject: [PATCH] Revert "simplify `subtask` in `tracer.kt`"

This reverts commit c08a93f58c22ee3933cd080dfbd6d931de87447f.

GitOrigin-RevId: 728a32cdf77bc83458d8bf7b7fc7f7abf320bae4
---
 platform/diagnostic/src/tracer.kt | 19 +++++++++++++------
 1 file changed, 13 insertions(+), 6 deletions(-)

diff --git a/platform/diagnostic/src/tracer.kt b/platform/diagnostic/src/tracer.kt
index 601c09e374c9..3b607f1dce4e 100644
--- a/platform/diagnostic/src/tracer.kt
+++ b/platform/diagnostic/src/tracer.kt
@@ -29,7 +29,6 @@ fun rootTask(): CoroutineContext = MeasureCoroutineTime
 
 /**
  * This function is designed to be used instead of `withContext(CoroutineName("subtask")) { ... }`.
- * See https://github.com/Kotlin/kotlinx.coroutines/issues/3414
  */
 suspend fun <X> subtask(
   name: String,
@@ -37,12 +36,20 @@ suspend fun <X> subtask(
   action: suspend CoroutineScope.() -> X,
 ): X {
   val namedContext = context + CoroutineName(name)
-  val measurer = coroutineContext[CoroutineTimeMeasurerKey]
-  return if (measurer == null) {
-    withContext(namedContext, action)
+  if (coroutineContext[CoroutineTimeMeasurerKey] == null) {
+    return withContext(namedContext, action)
   }
-  else {
-    withContext(namedContext + measurer.copyForChild(), action)
+  return coroutineScope {
+    @OptIn(ExperimentalStdlibApi::class)
+    val start = if (coroutineContext[CoroutineDispatcher] == context[CoroutineDispatcher]) {
+      // mimic withContext behaviour with its UndispatchedCoroutine
+      CoroutineStart.UNDISPATCHED
+    }
+    else {
+      CoroutineStart.DEFAULT
+    }
+    // async forces the framework to call CopyableThreadContextElement#copyForChild
+    async(namedContext, start, action).await()
   }
 }