Add OpenTelemetry base classes

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

@@ -0,0 +1,171 @@
+// Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+@file:Suppress("UnusedPrivateClass")
+
+package software.aws.toolkits.jetbrains.services.telemetry.otel
+
+import com.intellij.openapi.Disposable
+import com.intellij.openapi.components.Service
+import com.intellij.openapi.components.service
+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
+import io.opentelemetry.context.Context
+import io.opentelemetry.context.propagation.ContextPropagators
+import io.opentelemetry.exporter.internal.otlp.traces.TraceRequestMarshaler
+import io.opentelemetry.sdk.OpenTelemetrySdk
+import io.opentelemetry.sdk.resources.Resource
+import io.opentelemetry.sdk.trace.ReadWriteSpan
+import io.opentelemetry.sdk.trace.ReadableSpan
+import io.opentelemetry.sdk.trace.SdkTracerProvider
+import io.opentelemetry.sdk.trace.SpanProcessor
+import kotlinx.coroutines.CancellationException
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.launch
+import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider
+import software.amazon.awssdk.http.ContentStreamProvider
+import software.amazon.awssdk.http.HttpExecuteRequest
+import software.amazon.awssdk.http.SdkHttpMethod
+import software.amazon.awssdk.http.SdkHttpRequest
+import software.amazon.awssdk.http.apache.ApacheHttpClient
+import software.amazon.awssdk.http.auth.aws.signer.AwsV4HttpSigner
+import java.io.ByteArrayOutputStream
+import java.net.ConnectException
+
+private class BasicOtlpSpanProcessor(
+    private val coroutineScope: CoroutineScope,
+    private val traceUrl: String = "http://127.0.0.1:4318/v1/traces",
+) : SpanProcessor {
+    override fun onStart(parentContext: Context, span: ReadWriteSpan) {}
+    override fun isStartRequired() = false
+    override fun isEndRequired() = true
+
+    override fun onEnd(span: ReadableSpan) {
+        val data = span.toSpanData()
+        coroutineScope.launch {
+            try {
+                val item = TraceRequestMarshaler.create(listOf(data))
+
+                httpPost(traceUrl, contentLength = item.binarySerializedSize.toLong(), contentType = ContentType.XProtobuf) {
+                    item.writeBinaryTo(this)
+                }
+            } catch (e: CancellationException) {
+                throw e
+            } catch (e: ConnectException) {
+                thisLogger().warn("Cannot export (url=$traceUrl): ${e.message}")
+            } catch (e: Throwable) {
+                thisLogger().error("Cannot export (url=$traceUrl)", e)
+            }
+        }
+    }
+}
+
+private class SigV4OtlpSpanProcessor(
+    private val coroutineScope: CoroutineScope,
+    private val traceUrl: String,
+    private val creds: AwsCredentialsProvider,
+) : SpanProcessor {
+    override fun onStart(parentContext: Context, span: ReadWriteSpan) {}
+    override fun isStartRequired() = false
+    override fun isEndRequired() = true
+
+    private val client = ApacheHttpClient.create()
+
+    override fun onEnd(span: ReadableSpan) {
+        coroutineScope.launch {
+            val data = span.toSpanData()
+            try {
+                val item = TraceRequestMarshaler.create(listOf(data))
+                // calculate the sigv4 header
+                val signer = AwsV4HttpSigner.create()
+                val httpRequest =
+                    SdkHttpRequest.builder()
+                        .uri(traceUrl)
+                        .method(SdkHttpMethod.POST)
+                        .putHeader("Content-Type", "application/x-protobuf")
+                        .build()
+
+                val baos = ByteArrayOutputStream()
+                item.writeBinaryTo(baos)
+                val payload = ContentStreamProvider.fromByteArray(baos.toByteArray())
+                val signedRequest = signer.sign {
+                    it.identity(creds.resolveIdentity().get())
+                    it.request(httpRequest)
+                    it.payload(payload)
+                    it.putProperty(AwsV4HttpSigner.SERVICE_SIGNING_NAME, "osis")
+                    it.putProperty(AwsV4HttpSigner.REGION_NAME, "us-west-2")
+                }
+
+                // Create and HTTP client and send the request. ApacheHttpClient requires the 'apache-client' module.
+                client.prepareRequest(
+                    HttpExecuteRequest.builder()
+                        .request(signedRequest.request())
+                        .contentStreamProvider(signedRequest.payload().orElse(null))
+                        .build()
+                ).call()
+            } catch (e: CancellationException) {
+                throw e
+            } catch (e: ConnectException) {
+                thisLogger().warn("Cannot export (url=$traceUrl): ${e.message}")
+            } catch (e: Throwable) {
+                thisLogger().error("Cannot export (url=$traceUrl)", e)
+            }
+        }
+    }
+}
+
+private object StdoutSpanProcessor : SpanProcessor {
+    override fun onStart(parentContext: Context, span: ReadWriteSpan) {}
+    override fun isStartRequired() = false
+    override fun isEndRequired() = true
+
+    override fun onEnd(span: ReadableSpan) {
+        println(span.toSpanData())
+    }
+}
+
+@Service
+class OTelService @NonInjectable internal constructor(spanProcessors: List<SpanProcessor>) : Disposable {
+    @Suppress("unused")
+    constructor() : this(listOf(StdoutSpanProcessor))
+
+    private val sdkDelegate = lazy {
+        OpenTelemetrySdk.builder()
+            .setTracerProvider(
+                SdkTracerProvider.builder()
+                    .apply {
+                        spanProcessors.forEach {
+                            addSpanProcessor(it)
+                        }
+                    }
+                    .setResource(
+                        Resource.create(
+                            Attributes.builder()
+                                .put(AttributeKey.stringKey("os.type"), SystemInfoRt.OS_NAME)
+                                .put(AttributeKey.stringKey("os.version"), SystemInfoRt.OS_VERSION)
+                                .put(AttributeKey.stringKey("host.arch"), System.getProperty("os.arch"))
+                                .build()
+                        )
+                    )
+                    .build()
+            )
+            .setPropagators(ContextPropagators.create(W3CTraceContextPropagator.getInstance()))
+            .build()
+    }
+    internal val sdk: OpenTelemetrySdk by sdkDelegate
+
+    override fun dispose() {
+        if (sdkDelegate.isInitialized()) {
+            sdk.close()
+        }
+    }
+
+    companion object {
+        fun getSdk() = service<OTelService>().sdk
+    }
+}

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

@@ -0,0 +1,191 @@
+// Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package software.aws.toolkits.jetbrains.services.telemetry.otel
+
+import io.opentelemetry.api.common.AttributeKey
+import io.opentelemetry.api.common.Attributes
+import io.opentelemetry.api.trace.Span
+import io.opentelemetry.api.trace.SpanBuilder
+import io.opentelemetry.api.trace.SpanContext
+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
+import com.intellij.platform.diagnostic.telemetry.helpers.use as ijUse
+import com.intellij.platform.diagnostic.telemetry.helpers.useWithScope as ijUseWithScope
+
+val AWS_PRODUCT_CONTEXT_KEY = ContextKey.named<AWSProduct>("pluginDescriptor")
+internal val PLUGIN_ATTRIBUTE_KEY = AttributeKey.stringKey("plugin")
+
+class DefaultSpanBuilder(delegate: SpanBuilder) : AbstractSpanBuilder<DefaultSpanBuilder, AbstractBaseSpan>(delegate) {
+    override fun doStartSpan() = BaseSpan(parent, delegate.startSpan())
+}
+
+abstract class AbstractSpanBuilder<Builder : AbstractSpanBuilder<Builder, Span>, Span : AbstractBaseSpan>(protected val delegate: SpanBuilder) : SpanBuilder {
+    /**
+     * Same as [com.intellij.platform.diagnostic.telemetry.helpers.use] except downcasts to specific subclass of [BaseSpan]
+     *
+     * @inheritdoc
+     */
+    inline fun<T> use(operation: (Span) -> T): T =
+        startSpan().ijUse { 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
+        delegate.setParent(context)
+        return this as Builder
+    }
+
+    override fun setNoParent(): Builder {
+        parent = null
+        delegate.setNoParent()
+        return this as Builder
+    }
+
+    override fun addLink(spanContext: SpanContext): Builder {
+        delegate.addLink(spanContext)
+        return this as Builder
+    }
+
+    override fun addLink(
+        spanContext: SpanContext,
+        attributes: Attributes,
+    ): Builder {
+        delegate.addLink(spanContext, attributes)
+        return this as Builder
+    }
+
+    override fun setAttribute(key: String, value: String): Builder {
+        delegate.setAttribute(key, value)
+        return this as Builder
+    }
+
+    override fun setAttribute(key: String, value: Long): Builder {
+        delegate.setAttribute(key, value)
+        return this as Builder
+    }
+
+    override fun setAttribute(key: String, value: Double): Builder {
+        delegate.setAttribute(key, value)
+        return this as Builder
+    }
+
+    override fun setAttribute(key: String, value: Boolean): Builder {
+        delegate.setAttribute(key, value)
+        return this as Builder
+    }
+
+    override fun <V : Any?> setAttribute(
+        key: AttributeKey<V?>,
+        value: V & Any,
+    ): Builder {
+        delegate.setAttribute(key, value)
+        return this as Builder
+    }
+
+    override fun setAllAttributes(attributes: Attributes): Builder {
+        delegate.setAllAttributes(attributes)
+        return this as Builder
+    }
+
+    override fun setSpanKind(spanKind: SpanKind): Builder {
+        delegate.setSpanKind(spanKind)
+        return this as Builder
+    }
+
+    override fun setStartTimestamp(startTimestamp: Long, unit: TimeUnit): Builder {
+        delegate.setStartTimestamp(startTimestamp, unit)
+        return this as Builder
+    }
+
+    override fun setStartTimestamp(startTimestamp: Instant): Builder {
+        delegate.setStartTimestamp(startTimestamp)
+        return this as Builder
+    }
+
+    protected abstract fun doStartSpan(): Span
+
+    override fun startSpan(): Span {
+        var parent = parent
+        if (parent == null) {
+            parent = Context.current()
+        }
+        requireNotNull(parent)
+
+        val contextValue = parent.get(AWS_PRODUCT_CONTEXT_KEY)
+        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)
+            parent = if (s is AbstractBaseSpan && s.context != null) {
+                s.context.with(io.opentelemetry.api.trace.Span.fromContext(parent))
+            } else {
+                parent.with(AWS_PRODUCT_CONTEXT_KEY, resolvePluginName())
+            }
+            setParent(parent)
+        }
+        requireNotNull(parent)
+
+        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 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 {
+    /**
+     * Same as [com.intellij.platform.diagnostic.telemetry.helpers.use] except downcasts to specific subclass of [BaseSpan]
+     *
+     * @inheritdoc
+     */
+    inline fun<T> use(operation: (Span) -> T): T =
+        ijUse { span ->
+            operation(span as Span)
+        }
+
+    fun metadata(key: String, value: String) = setAttribute(key, value)
+
+    override fun makeCurrent(): Scope =
+        context?.with(this)?.makeCurrent() ?: super.makeCurrent()
+}
+
+/**
+ * Placeholder; will be generated
+ */
+class BaseSpan(context: Context?, delegate: Span) : AbstractBaseSpan(context, delegate) {
+    fun reason(reason: String) = metadata("reason", reason)
+}

plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/utils/ThreadingUtils.kt

@@ -14,6 +14,7 @@ import com.intellij.openapi.util.ThrowableComputable
 import com.intellij.util.ExceptionUtil
 import com.intellij.util.concurrency.AppExecutorUtil
 import com.intellij.util.concurrency.Semaphore
+import io.opentelemetry.context.Context
 import software.aws.toolkits.jetbrains.services.telemetry.PluginResolver
 import java.time.Duration
 import java.util.concurrent.Future
@@ -81,8 +82,9 @@ fun <T> pluginAwareExecuteOnPooledThread(action: () -> T): Future<T> {
      * worker thread will not contain original call stack. Necessary for telemetry.
      */
     val pluginResolver = PluginResolver.fromCurrentThread()
+    val context = Context.current()
     return ApplicationManager.getApplication().executeOnPooledThread<T> {
         PluginResolver.setThreadLocal(pluginResolver)
-        action()
+        context.wrap(action).call()
     }
 }