tst

Author: rli

plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/services/telemetry/otel/OTelService.kt

@@ -11,6 +11,7 @@ import com.intellij.openapi.diagnostic.thisLogger
 import com.intellij.openapi.util.SystemInfoRt
 import com.intellij.platform.util.http.ContentType
 import com.intellij.platform.util.http.httpPost
+import com.intellij.serviceContainer.NonInjectable
 import io.opentelemetry.api.common.AttributeKey
 import io.opentelemetry.api.common.Attributes
 import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator
@@ -129,12 +130,19 @@ private object StdoutSpanProcessor : SpanProcessor {
 }
 
 @Service
-class OTelService : Disposable {
+class OTelService @NonInjectable internal constructor(spanProcessors: List<SpanProcessor>): Disposable {
+    @Suppress("unused")
+    constructor() : this(listOf(StdoutSpanProcessor))
+
     private val sdkDelegate = lazy {
         OpenTelemetrySdk.builder()
             .setTracerProvider(
                 SdkTracerProvider.builder()
-                    .addSpanProcessor(StdoutSpanProcessor)
+                    .apply {
+                        spanProcessors.forEach {
+                            addSpanProcessor(it)
+                        }
+                    }
                     .setResource(
                         Resource.create(
                             Attributes.builder()

plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/services/telemetry/otel/OtelBase.kt

@@ -4,6 +4,7 @@
 package software.aws.toolkits.jetbrains.services.telemetry.otel
 
 import com.intellij.platform.diagnostic.telemetry.helpers.use as ijUse
+import com.intellij.platform.diagnostic.telemetry.helpers.useWithScope as ijUseWithScope
 import io.opentelemetry.api.common.AttributeKey
 import io.opentelemetry.api.common.Attributes
 import io.opentelemetry.api.trace.Span
@@ -13,16 +14,21 @@ import io.opentelemetry.api.trace.SpanKind
 import io.opentelemetry.context.Context
 import io.opentelemetry.context.ContextKey
 import io.opentelemetry.context.Scope
+import kotlinx.coroutines.CoroutineScope
 import software.amazon.awssdk.services.toolkittelemetry.model.AWSProduct
+import software.aws.toolkits.core.utils.getLogger
+import software.aws.toolkits.core.utils.warn
 import software.aws.toolkits.jetbrains.services.telemetry.PluginResolver
 import java.time.Instant
 import java.util.concurrent.TimeUnit
+import kotlin.coroutines.CoroutineContext
+import kotlin.coroutines.EmptyCoroutineContext
 
 val AWS_PRODUCT_CONTEXT_KEY = ContextKey.named<AWSProduct>("pluginDescriptor")
-private val PLUGIN_ATTRIBUTE_KEY = AttributeKey.stringKey("plugin")
+internal val PLUGIN_ATTRIBUTE_KEY = AttributeKey.stringKey("plugin")
 
 class DefaultSpanBuilder(delegate: SpanBuilder) : AbstractSpanBuilder<DefaultSpanBuilder, AbstractBaseSpan>(delegate) {
-    override fun doStartSpan() = BaseSpan(parent!!, delegate.startSpan())
+    override fun doStartSpan() = BaseSpan(parent, delegate.startSpan())
 }
 
 abstract class AbstractSpanBuilder<Builder : AbstractSpanBuilder<Builder, Span>, Span : AbstractBaseSpan>(protected val delegate: SpanBuilder) : SpanBuilder {
@@ -36,6 +42,19 @@ abstract class AbstractSpanBuilder<Builder : AbstractSpanBuilder<Builder, Span>,
             operation(span as Span)
         }
 
+    /**
+     * Same as [com.intellij.platform.diagnostic.telemetry.helpers.useWithScope] except downcasts to specific subclass of [BaseSpan]
+     *
+     * @inheritdoc
+     */
+    suspend inline fun<T> useWithScope(
+        context: CoroutineContext = EmptyCoroutineContext,
+        crossinline operation: suspend CoroutineScope.(Span) -> T
+    ): T =
+        ijUseWithScope(context) { span ->
+            operation(span as Span)
+        }
+
     protected var parent: Context? = null
     override fun setParent(context: Context): Builder {
         parent = context
@@ -123,25 +142,31 @@ abstract class AbstractSpanBuilder<Builder : AbstractSpanBuilder<Builder, Span>,
         if (contextValue == null) {
             // FIX_WHEN_MIN_IS_243: Kotlin compiler can't figure out difference between class/type parameter until 2.x
             val s = io.opentelemetry.api.trace.Span.fromContextOrNull(parent)
-            if (s is AbstractBaseSpan) {
-                setParent(s.context.with(io.opentelemetry.api.trace.Span.fromContext(parent)))
+            parent = if (s is AbstractBaseSpan && s.context != null) {
+                s.context.with(io.opentelemetry.api.trace.Span.fromContext(parent))
             } else {
-                setParent(parent.with(AWS_PRODUCT_CONTEXT_KEY, resolvePluginName()))
+                parent.with(AWS_PRODUCT_CONTEXT_KEY, resolvePluginName())
             }
+            setParent(parent)
         }
+        requireNotNull(parent)
 
-        setAttribute(
-            PLUGIN_ATTRIBUTE_KEY,
-            (parent.get(AWS_PRODUCT_CONTEXT_KEY) ?: resolvePluginName()).name
-        )
+        parent.get(AWS_PRODUCT_CONTEXT_KEY)?.toString()?.let {
+            setAttribute(PLUGIN_ATTRIBUTE_KEY, it)
+        } ?: run {
+            LOG.warn { "Reached setAttribute with null AWS_PRODUCT_CONTEXT_KEY, but should not be possible" }
+        }
 
         return doStartSpan()
     }
 
-    private fun resolvePluginName() = PluginResolver.Companion.fromStackTrace(Thread.currentThread().stackTrace).product
+    private companion object {
+        val LOG = getLogger<AbstractSpanBuilder<*, *>>()
+        fun resolvePluginName() = PluginResolver.Companion.fromStackTrace(Thread.currentThread().stackTrace).product
+    }
 }
 
-abstract class AbstractBaseSpan(internal val context: Context, private val delegate: Span) : Span by delegate {
+abstract class AbstractBaseSpan(internal val context: Context?, private val delegate: Span) : Span by delegate {
     /**
      * Same as [com.intellij.platform.diagnostic.telemetry.helpers.use] except downcasts to specific subclass of [BaseSpan]
      *
@@ -155,12 +180,12 @@ abstract class AbstractBaseSpan(internal val context: Context, private val deleg
     fun metadata(key: String, value: String) = setAttribute(key, value)
 
     override fun makeCurrent(): Scope =
-        context.with(this).makeCurrent()
+        context?.with(this)?.makeCurrent() ?: super.makeCurrent()
 }
 
 /**
  * Placeholder; will be generated
  */
-class BaseSpan(context: Context, delegate: Span) : AbstractBaseSpan(context, delegate) {
+class BaseSpan(context: Context?, delegate: Span) : AbstractBaseSpan(context, delegate) {
     fun reason(reason: String) = metadata("reason", reason)
 }