use new telemetry API in StartupManager

GitOrigin-RevId: e943f6b6c9f9d8ab39b4386d46e77cb0aa7b73c2

platform/diagnostic/telemetry-impl/src/OtlpService.kt

@@ -138,7 +138,7 @@ internal class OtlpService(private val coroutineScope: CoroutineScope) {
     return byteBuffer.array()
   }
 
-  private suspend fun flush(scopeToSpans: Map<Scope, ScopeSpans>, resource: Resource, endpoint: String) {
+  private suspend fun flush(scopeToSpans: MutableMap<Scope, ScopeSpans>, resource: Resource, endpoint: String) {
     if (scopeToSpans.isEmpty()) {
       return
     }
@@ -152,6 +152,7 @@ internal class OtlpService(private val coroutineScope: CoroutineScope) {
       ),
     )
     httpPost(url = endpoint, contentType = ContentType.XProtobuf, body = ProtoBuf.encodeToByteArray(data))
+    scopeToSpans.clear()
   }
 
   fun add(activity: ActivityImpl) {

platform/diagnostic/telemetry-impl/src/TelemetryManagerImpl.kt

@@ -1,10 +1,7 @@
 // Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
 package com.intellij.platform.diagnostic.telemetry.impl
 
-import com.intellij.diagnostic.ActivityImpl
-import com.intellij.diagnostic.DefaultTraceReporter
-import com.intellij.diagnostic.PluginException
-import com.intellij.diagnostic.rootTask
+import com.intellij.diagnostic.*
 import com.intellij.openapi.application.Application
 import com.intellij.openapi.application.ApplicationInfo
 import com.intellij.openapi.application.ApplicationManager
@@ -13,6 +10,7 @@ import com.intellij.openapi.application.impl.ApplicationInfoImpl
 import com.intellij.openapi.components.service
 import com.intellij.openapi.diagnostic.logger
 import com.intellij.openapi.extensions.ExtensionPointName
+import com.intellij.openapi.util.Ref
 import com.intellij.platform.diagnostic.telemetry.*
 import com.intellij.platform.diagnostic.telemetry.impl.otExporters.OpenTelemetryExporterProvider
 import com.intellij.util.childScope
@@ -81,14 +79,40 @@ class TelemetryManagerImpl(app: Application) : TelemetryManager {
 }
 
 private class IntelliJTracerImpl(private val scope: Scope, private val otlpService: OtlpService) : IntelliJTracer {
-  override fun createSpan(name: String): CoroutineContext {
-    return rootTask(traceReporter = object : DefaultTraceReporter(reportScheduleTimeForRoot = true) {
+  private val traceReporter = object : DefaultTraceReporter(reportScheduleTimeForRoot = true) {
+    override fun setEndAndAdd(activity: ActivityImpl, end: Long) {
+      activity.setEnd(end)
+      activity.scope = scope
+      otlpService.add(activity)
+    }
+  }
+
+  override suspend fun span(name: String): CoroutineContext {
+    return createSpan(traceReporter = traceReporter) + CoroutineName(name)
+  }
+
+  override suspend fun span(name: String, attributes: Array<String>): CoroutineContext {
+    return CoroutineName(name) +
+           createSpan(traceReporter = createAttributeAwareReporter(name, Ref(attributes.takeIf { it.isNotEmpty() })))
+  }
+
+  override fun rootSpan(name: String, attributes: Array<String>): CoroutineContext {
+    return CoroutineName(name) +
+           createRootSpan(traceReporter = createAttributeAwareReporter(name, Ref(attributes.takeIf { it.isNotEmpty() })))
+  }
+
+  private fun createAttributeAwareReporter(name: String, ref: Ref<Array<String>>): DefaultTraceReporter {
+    return object : DefaultTraceReporter(reportScheduleTimeForRoot = true) {
       override fun setEndAndAdd(activity: ActivityImpl, end: Long) {
-        activity.setEnd(end)
-        activity.scope = scope
-        otlpService.add(activity)
+        // yes, compare by identity
+        if (activity.name === name) {
+          activity.attributes = ref.get()
+          ref.set(null)
+        }
+
+        traceReporter.setEndAndAdd(activity, end)
       }
-    }) + CoroutineName(name)
+    }
   }
 }
 

platform/diagnostic/telemetry-impl/src/tracer.kt

@@ -5,11 +5,9 @@ package com.intellij.diagnostic
 
 import com.intellij.platform.diagnostic.telemetry.impl.TraceReporter
 import kotlinx.coroutines.*
-import org.jetbrains.annotations.ApiStatus.Experimental
 import org.jetbrains.annotations.ApiStatus.Internal
 import java.lang.invoke.MethodHandles
 import java.lang.invoke.VarHandle
-import kotlin.coroutines.AbstractCoroutineContextElement
 import kotlin.coroutines.CoroutineContext
 import kotlin.coroutines.EmptyCoroutineContext
 import kotlin.coroutines.coroutineContext
@@ -33,8 +31,13 @@ private val defaultTracer = MeasureCoroutineTime(reporter = DefaultTraceReporter
  */
 fun rootTask(): CoroutineContext = defaultTracer
 
-@Internal
-fun rootTask(traceReporter: TraceReporter): CoroutineContext = MeasureCoroutineTime(reporter = traceReporter)
+internal suspend fun createSpan(traceReporter: TraceReporter): CoroutineContext {
+  return CoroutineTimeMeasurer(reporter = traceReporter, parentActivity = coroutineContext[CoroutineTimeMeasurerKey]?.getActivity())
+}
+
+internal fun createRootSpan(traceReporter: TraceReporter): CoroutineContext {
+  return CoroutineTimeMeasurer(reporter = traceReporter, parentActivity = null)
+}
 
 /**
  * This function is designed to be used instead of `withContext(CoroutineName("subtask")) { ... }`.
@@ -52,17 +55,6 @@ suspend fun <T> span(name: String, context: CoroutineContext = EmptyCoroutineCon
   }
 }
 
-@Experimental
-class SpanAttributes(@JvmField val attributes: Array<String>) : AbstractCoroutineContextElement(SpanAttributes) {
-  companion object Key : CoroutineContext.Key<SpanAttributes>
-
-  override fun toString(): String = "SpanAttributes($attributes)"
-
-  override fun equals(other: Any?): Boolean = this === other || (other is SpanAttributes && attributes.contentEquals(other.attributes))
-
-  override fun hashCode(): Int = attributes.contentHashCode()
-}
-
 @Internal
 suspend fun getTraceActivity(): Activity? = coroutineContext[CoroutineTimeMeasurerKey]?.getActivity()
 
@@ -153,7 +145,6 @@ private class CoroutineTimeMeasurer(private val parentActivity: ActivityImpl?,
     val activity = reporter.start(coroutineName = coroutineName, scheduleTime = creationTime, parentActivity = parentActivity)
     CURRENT_ACTIVITY.setVolatile(this, activity)
 
-    activity.attributes = context[SpanAttributes]?.attributes
     val job = context.job
     job.invokeOnCompletion {
       val end = System.nanoTime()
@@ -168,11 +159,9 @@ private class CoroutineTimeMeasurer(private val parentActivity: ActivityImpl?,
   }
 
   override fun copyForChild(): CopyableThreadContextElement<Unit> {
-    val activity = checkNotNull(CURRENT_ACTIVITY.getVolatile(this)) {
-      "updateThreadContext was never called"
-    }
+    val activity = CURRENT_ACTIVITY.getVolatile(this)
     // don't format into one line - complicates debugging
-    if (activity === noActivity) {
+    if (activity === noActivity || activity == null) {
       return CoroutineTimeMeasurer(parentActivity = null, reporter = reporter)
     }
     else {
@@ -182,7 +171,7 @@ private class CoroutineTimeMeasurer(private val parentActivity: ActivityImpl?,
 
   override fun mergeForChild(overwritingElement: CoroutineContext.Element): CoroutineContext = overwritingElement
 
-  fun getActivity() = CURRENT_ACTIVITY.getVolatile(this) as Activity?
+  fun getActivity() = CURRENT_ACTIVITY.getVolatile(this) as ActivityImpl?
 }
 
 open class DefaultTraceReporter(private val reportScheduleTimeForRoot: Boolean) : TraceReporter {

platform/diagnostic/telemetry/src/com/intellij/platform/diagnostic/telemetry/TelemetryManager.kt

@@ -5,7 +5,8 @@ import com.intellij.openapi.diagnostic.logger
 import com.intellij.util.concurrency.SynchronizedClearableLazy
 import io.opentelemetry.api.OpenTelemetry
 import io.opentelemetry.api.metrics.Meter
-import org.jetbrains.annotations.ApiStatus
+import kotlinx.coroutines.CoroutineName
+import org.jetbrains.annotations.ApiStatus.Experimental
 import org.jetbrains.annotations.ApiStatus.Internal
 import java.util.*
 import kotlin.coroutines.CoroutineContext
@@ -15,7 +16,7 @@ import kotlin.coroutines.EmptyCoroutineContext
  * See [Span](https://opentelemetry.io/docs/reference/specification),
  * [Manual Instrumentation](https://opentelemetry.io/docs/instrumentation/java/manual/#create-spans-with-events).
  */
-@ApiStatus.Experimental
+@Experimental
 @Internal
 interface TelemetryManager {
   companion object {
@@ -96,11 +97,22 @@ internal class NoopTelemetryManager : TelemetryManager {
   }
 }
 
+// suspend here is required to get parent span from coroutine context; that's why a version without `suspend` is called `rootSpan`.
+@Internal
+@Experimental
 interface IntelliJTracer {
-  fun createSpan(name: String): CoroutineContext
+  suspend fun span(name: String): CoroutineContext
+
+  suspend fun span(name: String, attributes: Array<String>): CoroutineContext
+
+  fun rootSpan(name: String, attributes: Array<String>): CoroutineContext
 }
 
 @Internal
 object NoopIntelliJTracer : IntelliJTracer {
-  override fun createSpan(name: String): CoroutineContext = EmptyCoroutineContext
+  override suspend fun span(name: String) = CoroutineName(name)
+
+  override suspend fun span(name: String, attributes: Array<String>) = span(name)
+
+  override fun rootSpan(name: String, attributes: Array<String>) = EmptyCoroutineContext
 }

platform/platform-impl/src/com/intellij/ide/startup/impl/StartupManagerImpl.kt

@@ -15,6 +15,7 @@ import com.intellij.openapi.application.EDT
 import com.intellij.openapi.application.ModalityState
 import com.intellij.openapi.application.readActionBlocking
 import com.intellij.openapi.diagnostic.ControlFlowException
+import com.intellij.openapi.diagnostic.getOrLogException
 import com.intellij.openapi.diagnostic.logger
 import com.intellij.openapi.extensions.ExtensionPointListener
 import com.intellij.openapi.extensions.ExtensionPointName
@@ -33,9 +34,9 @@ import com.intellij.openapi.startup.StartupActivity
 import com.intellij.openapi.util.registry.Registry
 import com.intellij.platform.diagnostic.telemetry.Scope
 import com.intellij.platform.diagnostic.telemetry.TelemetryManager
-import com.intellij.platform.diagnostic.telemetry.helpers.useWithScope
 import com.intellij.serviceContainer.ComponentManagerImpl
 import com.intellij.util.ModalityUiUtil
+import com.intellij.util.ui.EDT
 import io.opentelemetry.api.common.AttributeKey
 import io.opentelemetry.api.common.Attributes
 import io.opentelemetry.api.trace.Span
@@ -48,22 +49,17 @@ import org.jetbrains.annotations.VisibleForTesting
 import java.awt.event.InvocationEvent
 import java.util.*
 import java.util.concurrent.CancellationException
-import java.util.concurrent.TimeUnit
-import java.util.concurrent.atomic.AtomicBoolean
 import java.util.concurrent.atomic.AtomicInteger
-import java.util.concurrent.atomic.AtomicReference
 import kotlin.coroutines.coroutineContext
 
 private val LOG = logger<StartupManagerImpl>()
-private val tracer by lazy { TelemetryManager.getTracer(Scope("startupManager")) }
-private val tracer2 by lazy { TelemetryManager.getSimpleTracer(Scope("startupManager")) }
+private val tracer by lazy { TelemetryManager.getSimpleTracer(Scope("startupManager")) }
 
 /**
  * Acts as [StartupActivity.POST_STARTUP_ACTIVITY], but executed with 5-seconds delay after project opening.
  */
 private val BACKGROUND_POST_STARTUP_ACTIVITY: ExtensionPointName<Any> = ExtensionPointName("com.intellij.backgroundPostStartupActivity")
 
-private val EDT_WARN_THRESHOLD_IN_NANO = TimeUnit.MILLISECONDS.toNanos(100)
 private const val DUMB_AWARE_PASSED = 1
 private const val ALL_PASSED = 2
 
@@ -88,13 +84,13 @@ open class StartupManagerImpl(private val project: Project, private val coroutin
             is DumbAware -> {
               startupManager.coroutineScope.launch {
                 @Suppress("UsagesOfObsoleteApi")
-                startupManager.runActivityAndMeasureDuration(extension as StartupActivity, pluginDescriptor.pluginId)
+                startupManager.runOldActivity(extension as StartupActivity)
               }
             }
             else -> {
               DumbService.getInstance(project).runWhenSmart {
                 @Suppress("UsagesOfObsoleteApi")
-                startupManager.runActivityAndMeasureDuration(extension as StartupActivity, pluginDescriptor.pluginId)
+                startupManager.runOldActivity(extension as StartupActivity)
               }
             }
           }
@@ -171,11 +167,11 @@ open class StartupManagerImpl(private val project: Project, private val coroutin
     }
     else {
       // doesn't block project opening
-      coroutineScope.launch {
+      coroutineScope.launch(tracer.span("runPostStartupActivities")) {
         runPostStartupActivities(async = true)
       }
       if (app.isUnitTestMode) {
-        ApplicationManager.getApplication().assertIsDispatchThread()
+        EDT.assertIsEdt()
         waitAndProcessInvocationEventsInIdeEventQueue(this)
       }
     }
@@ -183,9 +179,9 @@ open class StartupManagerImpl(private val project: Project, private val coroutin
 
   private suspend fun runInitProjectActivities() {
     runActivities(initProjectStartupActivities)
-    val app = ApplicationManager.getApplication()
-    val extensionPoint = (app.extensionArea as ExtensionsAreaImpl).getExtensionPoint<InitProjectActivity>(
-      "com.intellij.initProjectActivity")
+
+    val extensionPoint = (ApplicationManager.getApplication().extensionArea as ExtensionsAreaImpl)
+      .getExtensionPoint<InitProjectActivity>("com.intellij.initProjectActivity")
     // do not create an extension if not allow-listed
     for (adapter in extensionPoint.sortedAdapters) {
       coroutineContext.ensureActive()
@@ -206,14 +202,11 @@ open class StartupManagerImpl(private val project: Project, private val coroutin
 
       val activity = adapter.createInstance<InitProjectActivity>(project) ?: continue
       val startTime = System.nanoTime()
-      tracer.spanBuilder("run activity")
-        .setAttribute(AttributeKey.stringKey("class"), activity.javaClass.name)
-        .setAttribute(AttributeKey.stringKey("plugin"), pluginId.idString)
-        .useWithScope {
-          if (project !is LightEditCompatible || activity is LightEditCompatible) {
-            activity.run(project)
-          }
+      withContext(tracer.span("run activity", arrayOf("class", activity.javaClass.name, "plugin", pluginId.idString))) {
+        if (project !is LightEditCompatible || activity is LightEditCompatible) {
+          activity.run(project)
         }
+      }
 
       addCompletedActivity(startTime = startTime, runnableClass = activity.javaClass, pluginId = pluginId)
     }
@@ -228,8 +221,6 @@ open class StartupManagerImpl(private val project: Project, private val coroutin
       // (depending on dumb-awareness), but because there is no other concurrent phase, we measure it as a sequential activity
       // to put it on the timeline and make clear what's going at the end (avoiding the last "unknown" phase)
       val dumbAwareActivity = StartUpMeasurer.startActivity(StartUpMeasurer.Activities.PROJECT_DUMB_POST_START_UP_ACTIVITIES)
-      val edtActivity = AtomicReference<Activity?>()
-      val uiFreezeWarned = AtomicBoolean()
       val counter = AtomicInteger()
       val dumbService = DumbService.getInstance(project)
       val isProjectLightEditCompatible = project is LightEditCompatible
@@ -256,46 +247,45 @@ open class StartupManagerImpl(private val project: Project, private val coroutin
 
         @Suppress("SSBasedInspection", "UsagesOfObsoleteApi")
         if (activity is DumbAware) {
+          //LOG.warn(PluginException("Migrate ${item.implementationClassName} to ProjectActivity", pluginDescriptor.pluginId))
           dumbService.runWithWaitForSmartModeDisabled().use {
             blockingContext {
-              runActivityAndMeasureDuration(activity as StartupActivity, pluginDescriptor.pluginId)
+              runOldActivity(activity as StartupActivity)
             }
           }
           continue
         }
         else if (!isProjectLightEditCompatible) {
+          //LOG.warn(PluginException("Migrate ${item.implementationClassName} to ProjectActivity", pluginDescriptor.pluginId))
           // DumbService.unsafeRunWhenSmart throws an assertion in LightEdit mode, see LightEditDumbService.unsafeRunWhenSmart
           counter.incrementAndGet()
           blockingContext {
             dumbService.runWhenSmart {
               traceContext.makeCurrent()
 
-              if (edtActivity.get() == null) {
-                edtActivity.set(StartUpMeasurer.startActivity("project post-startup edt activities"))
-              }
-
-              val duration = runActivityAndMeasureDuration(activity as StartupActivity, pluginDescriptor.pluginId)
-              if (duration > EDT_WARN_THRESHOLD_IN_NANO) {
-                reportUiFreeze(uiFreezeWarned)
-              }
-              dumbUnawarePostActivitiesPassed(edtActivity, counter.decrementAndGet())
+              runOldActivity(activity as StartupActivity)
             }
           }
         }
       }
 
-      dumbUnawarePostActivitiesPassed(edtActivity, counter.get())
-
-      runPostStartupActivitiesRegisteredDynamically()
-      dumbAwareActivity.end()
-      coroutineScope.launch {
-        StartUpPerformanceService.getInstance().projectDumbAwareActivitiesFinished()
+      withContext(tracer.span("run post-startup dynamically registered activities")) {
+        runActivities(postStartupActivities)
       }
-      snapshot.logResponsivenessSinceCreation("Post-startup activities under progress")
+      postStartupActivitiesPassed = DUMB_AWARE_PASSED
+      postStartupActivitiesPassed = ALL_PASSED
+      allActivitiesPassed.complete(value = null)
+      dumbAwareActivity.end()
 
       coroutineContext.ensureActive()
 
+      StartUpPerformanceService.getInstance().projectDumbAwareActivitiesFinished()
+
       if (!ApplicationManager.getApplication().isUnitTestMode) {
+        coroutineContext.ensureActive()
+
+        snapshot.logResponsivenessSinceCreation("Post-startup activities under progress")
+
         scheduleBackgroundPostStartupActivities(project, coroutineScope)
         addActivityEpListener(project)
       }
@@ -313,17 +303,11 @@ open class StartupManagerImpl(private val project: Project, private val coroutin
     }
   }
 
-  private fun runActivityAndMeasureDuration(activity: StartupActivity, pluginId: PluginId): Long {
-    val startTime = System.nanoTime()
+  private fun runOldActivity(@Suppress("UsagesOfObsoleteApi") activity: StartupActivity) {
     try {
-      tracer.spanBuilder("run activity")
-        .setAttribute(AttributeKey.stringKey("class"), activity.javaClass.name)
-        .setAttribute(AttributeKey.stringKey("plugin"), pluginId.idString)
-        .useWithScope {
-          if (project !is LightEditCompatible || activity is LightEditCompatible) {
-            activity.runActivity(project)
-          }
-        }
+      if (project !is LightEditCompatible || activity is LightEditCompatible) {
+        activity.runActivity(project)
+      }
     }
     catch (e: Throwable) {
       if (e is ControlFlowException || e is CancellationException) {
@@ -331,18 +315,6 @@ open class StartupManagerImpl(private val project: Project, private val coroutin
       }
       LOG.error(e)
     }
-
-    return addCompletedActivity(startTime = startTime, runnableClass = activity.javaClass, pluginId = pluginId)
-  }
-
-  private suspend fun runPostStartupActivitiesRegisteredDynamically() {
-    span("run post-startup dynamically registered activities") {
-      runActivities(postStartupActivities)
-    }
-
-    postStartupActivitiesPassed = DUMB_AWARE_PASSED
-    postStartupActivitiesPassed = ALL_PASSED
-    allActivitiesPassed.complete(value = null)
   }
 
   private suspend fun runActivities(activities: Deque<Runnable>) {
@@ -352,31 +324,21 @@ open class StartupManagerImpl(private val project: Project, private val coroutin
       }
     }
 
-    while (true) {
-      coroutineContext.ensureActive()
+    coroutineScope {
+      while (true) {
+        coroutineContext.ensureActive()
 
-      val runnable = synchronized(lock, activities::pollFirst) ?: break
-      val startTime = System.nanoTime()
-      val runnableClass = runnable.javaClass
-      val pluginId = (runnableClass.classLoader as? PluginAwareClassLoader)?.pluginId ?: PluginManagerCore.CORE_ID
-      try {
-        tracer.spanBuilder("run activity")
-          .setAttribute(AttributeKey.stringKey("class"), runnableClass.name)
-          .setAttribute(AttributeKey.stringKey("plugin"), pluginId.idString)
-          .useWithScope {
+        val runnable = synchronized(lock, activities::pollFirst) ?: break
+        val runnableClass = runnable.javaClass
+        val pluginId = (runnableClass.classLoader as? PluginAwareClassLoader)?.pluginId ?: PluginManagerCore.CORE_ID
+        launch(tracer.span("run activity", arrayOf("class", runnableClass.name, "plugin", pluginId.idString))) {
+          runCatching {
             blockingContext {
               runnable.run()
             }
-          }
+          }.getOrLogException(LOG)
+        }
       }
-      catch (e: CancellationException) {
-        throw e
-      }
-      catch (e: Throwable) {
-        LOG.error(e)
-      }
-
-      addCompletedActivity(startTime = startTime, runnableClass = runnableClass, pluginId = pluginId)
     }
   }
 
@@ -493,24 +455,9 @@ private fun addCompletedActivity(startTime: Long, runnableClass: Class<*>, plugi
   )
 }
 
-private fun dumbUnawarePostActivitiesPassed(edtActivity: AtomicReference<Activity?>, count: Int) {
-  if (count == 0) {
-    edtActivity.getAndSet(null)?.end()
-  }
-}
-
-private fun reportUiFreeze(uiFreezeWarned: AtomicBoolean) {
-  val app = ApplicationManager.getApplication()
-  if (!app.isUnitTestMode && app.isDispatchThread && uiFreezeWarned.compareAndSet(false, true)) {
-    LOG.info("Some post-startup activities freeze UI for noticeable time. " +
-             "Please consider making them DumbAware to run them in background under modal progress," +
-             " or just making them faster to speed up project opening.")
-  }
-}
-
 private fun launchActivity(activity: ProjectActivity, project: Project, pluginId: PluginId) {
   (project as ComponentManagerImpl).pluginCoroutineScope(activity.javaClass.classLoader).launch(
-    tracer2.createSpan(name = "run activity") + SpanAttributes(arrayOf("class", activity.javaClass.name, "plugin", pluginId.idString))
+    tracer.rootSpan(name = "run activity", arrayOf("class", activity.javaClass.name, "plugin", pluginId.idString))
   ) {
     activity.execute(project)
   }