update subtask function to work around withContext behaviour

See Kotlin/kotlinx.coroutines#3411
See Kotlin/kotlinx.coroutines#3414

GitOrigin-RevId: 234f21de1899b48425f97193ca63769246b54173

Author: dovchinnikov

platform/diagnostic/src/tracer.kt

@@ -28,27 +28,29 @@ import kotlin.coroutines.coroutineContext
 fun rootTask(): CoroutineContext = MeasureCoroutineTime
 
 /**
- * This function forces child activity for [withContext] blocks,
- * which don't fork [CopyableThreadContextElement] under normal conditions.
- *
- * Example:
- * ```
- * launch(CoroutineName("init stuff")) {
- *   ...
- *   withContext(subtask("init stuff step 1") {
- *     ...
- *   }
- *   ...
- *   withContext(Dispatchers.EDT + subtask("init stuff step 2") {
- *     ...
- *   }
- *   ...
- * }
+ * This function is designed to be used instead of `withContext(CoroutineName("subtask")) { ... }`.
  */
-suspend fun subtask(name: String): CoroutineContext {
-  val measurer = coroutineContext[CoroutineTimeMeasurerKey]
-  val childMeasurer = measurer?.copyForChild() ?: EmptyCoroutineContext
-  return CoroutineName(name) + childMeasurer
+suspend fun <X> subtask(
+  name: String,
+  context: CoroutineContext = EmptyCoroutineContext,
+  action: suspend CoroutineScope.() -> X,
+): X {
+  val namedContext = context + CoroutineName(name)
+  if (coroutineContext[CoroutineTimeMeasurerKey] == null) {
+    return withContext(namedContext, 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()
+  }
 }
 
 private val noActivity: Activity = object : Activity {

platform/platform-impl/src/com/intellij/idea/StartupUtil.kt

@@ -176,7 +176,7 @@ fun start(args: Array<String>) {
       else {
         val euaDocument = euaDocumentDeferred.await()
         val baseLaF = initUiJob.await()
-        withContext(subtask("show eua")) {
+        subtask("show eua") {
           showEuaIfNeeded(euaDocument, baseLaF)
         }
       }
@@ -187,7 +187,7 @@ fun start(args: Array<String>) {
       configImportNeededDeferred.join()
 
       val (configPath, systemPath) = pathDeferred.await()
-      withContext(subtask("system dirs checking")) {
+      subtask("system dirs checking") {
         if (!checkSystemDirs(configPath, systemPath)) {
           exitProcess(AppExitCodes.DIR_CHECK_FAILED)
         }
@@ -268,11 +268,11 @@ fun start(args: Array<String>) {
       // this must happen after locking system dirs
       val log = logDeferred.await()
 
-      withContext(subtask("system libs setup")) {
+      subtask("system libs setup") {
         setupSystemLibraries()
       }
 
-      withContext(subtask("system libs loading") + Dispatchers.IO) {
+      subtask("system libs loading", Dispatchers.IO) {
         JnaLoader.load(log)
         if (SystemInfoRt.isWindows) {
           IdeaWin32.isAvailable()
@@ -289,7 +289,7 @@ fun start(args: Array<String>) {
         launch {
           initUiJob.join()
 
-          withContext(subtask("mac app init")) {
+          subtask("mac app init") {
             MacOSApplicationProvider.initApplication(log)
           }
         }
@@ -310,7 +310,7 @@ fun start(args: Array<String>) {
 
     val telemetryInitJob = launch(asyncDispatcher) {
       appInfoDeferred.join()
-      withContext(subtask("opentelemetry configuration")) {
+      subtask("opentelemetry configuration") {
         TraceManager.init()
       }
     }
@@ -873,7 +873,7 @@ private fun CoroutineScope.lockSystemDirs(checkConfig: Boolean,
 
   return async(Dispatchers.IO) {
     val (configPath, systemPath) = pathDeferred.await()
-    withContext(subtask("system dirs locking")) {
+    subtask("system dirs locking") {
       // this check must be performed before system directories are locked
       val importNeeded = checkConfig &&
                          (!Files.exists(configPath) || Files.exists(configPath.resolve(ConfigImportHelper.CUSTOM_MARKER_FILE_NAME)))