Add OpenTelemetry base classes

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

@@ -0,0 +1,163 @@
+// 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 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 : Disposable {
+    private val sdkDelegate = lazy {
+        OpenTelemetrySdk.builder()
+            .setTracerProvider(
+                SdkTracerProvider.builder()
+                    .addSpanProcessor(StdoutSpanProcessor)
+                    .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,166 @@
+// 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 com.intellij.platform.diagnostic.telemetry.helpers.use as ijUse
+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 software.amazon.awssdk.services.toolkittelemetry.model.AWSProduct
+import software.aws.toolkits.jetbrains.services.telemetry.PluginResolver
+import java.time.Instant
+import java.util.concurrent.TimeUnit
+
+val AWS_PRODUCT_CONTEXT_KEY = ContextKey.named<AWSProduct>("pluginDescriptor")
+private 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)
+        }
+
+    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)
+            if (s is AbstractBaseSpan) {
+                setParent(s.context.with(io.opentelemetry.api.trace.Span.fromContext(parent)))
+            } else {
+                setParent(parent.with(AWS_PRODUCT_CONTEXT_KEY, resolvePluginName()))
+            }
+        }
+
+        setAttribute(
+            PLUGIN_ATTRIBUTE_KEY,
+            (parent.get(AWS_PRODUCT_CONTEXT_KEY) ?: resolvePluginName()).name
+        )
+
+        return doStartSpan()
+    }
+
+    private 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()
+}
+
+/**
+ * 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()
     }
 }

plugins/core/jetbrains-community/tst/software/aws/toolkits/jetbrains/core/credentials/sso/SsoCredentialProviderTest.kt

@@ -3,6 +3,7 @@
 
 package software.aws.toolkits.jetbrains.core.credentials.sso
 
+import com.intellij.openapi.application.ApplicationManager
 import com.intellij.testFramework.ApplicationRule
 import org.assertj.core.api.Assertions.assertThat
 import org.assertj.core.api.Assertions.assertThatThrownBy
@@ -95,6 +96,13 @@ class SsoCredentialProviderTest {
         verify(ssoAccessTokenProvider).invalidate()
     }
 
+    @Test
+    fun aaaaa() {
+        ApplicationManager.getApplication().executeOnPooledThread {
+            println(Thread.currentThread().getContextClassLoader())
+        }.get()
+    }
+
     private fun createSsoResponse(expirationTime: Instant) {
         ssoClient.stub {
             on(