Skip to content

Commit 9a3743d

Browse files
rlirli
committed
Add OpenTelemetry base classes
1 parent 4319d6d commit 9a3743d

File tree

3 files changed

+305
-1
lines changed

3 files changed

+305
-1
lines changed

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

Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,143 @@
1+
// Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
package software.aws.toolkits.jetbrains.services.telemetry.otel
5+
6+
import com.intellij.openapi.Disposable
7+
import com.intellij.openapi.components.Service
8+
import com.intellij.openapi.diagnostic.thisLogger
9+
import com.intellij.platform.diagnostic.telemetry.impl.OpenTelemetryConfigurator
10+
import com.intellij.platform.util.http.ContentType
11+
import com.intellij.platform.util.http.httpPost
12+
import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator
13+
import io.opentelemetry.context.Context
14+
import io.opentelemetry.context.propagation.ContextPropagators
15+
import io.opentelemetry.exporter.internal.otlp.traces.TraceRequestMarshaler
16+
import io.opentelemetry.sdk.OpenTelemetrySdk
17+
import io.opentelemetry.sdk.trace.ReadWriteSpan
18+
import io.opentelemetry.sdk.trace.ReadableSpan
19+
import io.opentelemetry.sdk.trace.SdkTracerProvider
20+
import io.opentelemetry.sdk.trace.SpanProcessor
21+
import kotlinx.coroutines.CancellationException
22+
import kotlinx.coroutines.CoroutineScope
23+
import kotlinx.coroutines.launch
24+
import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider
25+
import software.amazon.awssdk.http.ContentStreamProvider
26+
import software.amazon.awssdk.http.HttpExecuteRequest
27+
import software.amazon.awssdk.http.SdkHttpMethod
28+
import software.amazon.awssdk.http.SdkHttpRequest
29+
import software.amazon.awssdk.http.apache.ApacheHttpClient
30+
import software.amazon.awssdk.http.auth.aws.signer.AwsV4HttpSigner
31+
import java.io.ByteArrayOutputStream
32+
import java.net.ConnectException
33+
34+
private class BasicOtlpSpanProcessor(private val coroutineScope: CoroutineScope, private val traceUrl: String = "http://127.0.0.1:4318/v1/traces") : SpanProcessor {
35+
override fun onStart(parentContext: Context, span: ReadWriteSpan) {}
36+
override fun isStartRequired() = false
37+
override fun isEndRequired() = true
38+
39+
override fun onEnd(span: ReadableSpan) {
40+
val data = span.toSpanData()
41+
coroutineScope.launch {
42+
try {
43+
val item = TraceRequestMarshaler.create(listOf(data))
44+
45+
httpPost(traceUrl, contentLength = item.binarySerializedSize.toLong(), contentType = ContentType.XProtobuf) {
46+
item.writeBinaryTo(this)
47+
}
48+
} catch (e: CancellationException) {
49+
throw e
50+
} catch (e: ConnectException) {
51+
thisLogger().warn("Cannot export (url=$traceUrl): ${e.message}")
52+
} catch (e: Throwable) {
53+
thisLogger().error("Cannot export (url=$traceUrl)", e)
54+
}
55+
}
56+
}
57+
}
58+
59+
private class SigV4OtlpSpanProcessor(private val coroutineScope: CoroutineScope, private val traceUrl: String, private val creds: AwsCredentialsProvider) : SpanProcessor {
60+
override fun onStart(parentContext: Context, span: ReadWriteSpan) {}
61+
override fun isStartRequired() = false
62+
override fun isEndRequired() = true
63+
64+
private val client = ApacheHttpClient.create()
65+
66+
override fun onEnd(span: ReadableSpan) {
67+
coroutineScope.launch {
68+
val data = span.toSpanData()
69+
try {
70+
val item = TraceRequestMarshaler.create(listOf(data))
71+
// calculate the sigv4 header
72+
val signer = AwsV4HttpSigner.create()
73+
val httpRequest =
74+
SdkHttpRequest.builder()
75+
.uri(traceUrl)
76+
.method(SdkHttpMethod.POST)
77+
.putHeader("Content-Type", "application/x-protobuf")
78+
.build()
79+
80+
val baos = ByteArrayOutputStream()
81+
item.writeBinaryTo(baos)
82+
val payload = ContentStreamProvider.fromByteArray(baos.toByteArray())
83+
val signedRequest = signer.sign {
84+
it.identity(creds.resolveIdentity().get())
85+
it.request(httpRequest)
86+
it.payload(payload)
87+
it.putProperty(AwsV4HttpSigner.SERVICE_SIGNING_NAME, "osis")
88+
it.putProperty(AwsV4HttpSigner.REGION_NAME, "us-west-2")
89+
}
90+
91+
// Create and HTTP client and send the request. ApacheHttpClient requires the 'apache-client' module.
92+
client.prepareRequest(
93+
HttpExecuteRequest.builder()
94+
.request(signedRequest.request())
95+
.contentStreamProvider(signedRequest.payload().orElse(null))
96+
.build()
97+
).call()
98+
} catch (e: CancellationException) {
99+
throw e
100+
} catch (e: ConnectException) {
101+
thisLogger().warn("Cannot export (url=$traceUrl): ${e.message}")
102+
} catch (e: Throwable) {
103+
thisLogger().error("Cannot export (url=$traceUrl)", e)
104+
}
105+
}
106+
}
107+
}
108+
109+
private object StdoutSpanProcessor : SpanProcessor {
110+
override fun onStart(parentContext: Context, span: ReadWriteSpan) {}
111+
override fun isStartRequired() = false
112+
override fun isEndRequired() = true
113+
114+
override fun onEnd(span: ReadableSpan) {
115+
println(span.toSpanData())
116+
}
117+
}
118+
119+
@Service
120+
class OTelService(private val cs: CoroutineScope) : Disposable {
121+
private val sdkDelegate = lazy {
122+
val configurator = OpenTelemetryConfigurator(
123+
sdkBuilder = OpenTelemetrySdk.builder(),
124+
)
125+
126+
configurator.getConfiguredSdkBuilder()
127+
.setTracerProvider(
128+
SdkTracerProvider.builder()
129+
.addSpanProcessor(StdoutSpanProcessor)
130+
.setResource(configurator.resource)
131+
.build()
132+
)
133+
.setPropagators(ContextPropagators.create(W3CTraceContextPropagator.getInstance()))
134+
.build()
135+
}
136+
val sdk: OpenTelemetrySdk by sdkDelegate
137+
138+
override fun dispose() {
139+
if (sdkDelegate.isInitialized()) {
140+
sdk.close()
141+
}
142+
}
143+
}

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

Lines changed: 159 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,159 @@
1+
// Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
package software.aws.toolkits.jetbrains.services.telemetry.otel
5+
6+
import com.intellij.platform.diagnostic.telemetry.helpers.useWithoutActiveScope
7+
import io.opentelemetry.api.common.AttributeKey
8+
import io.opentelemetry.api.common.Attributes
9+
import io.opentelemetry.api.trace.Span
10+
import io.opentelemetry.api.trace.SpanBuilder
11+
import io.opentelemetry.api.trace.SpanContext
12+
import io.opentelemetry.api.trace.SpanKind
13+
import io.opentelemetry.context.Context
14+
import io.opentelemetry.context.ContextKey
15+
import io.opentelemetry.context.Scope
16+
import software.amazon.awssdk.services.toolkittelemetry.model.AWSProduct
17+
import software.aws.toolkits.jetbrains.services.telemetry.PluginResolver
18+
import java.time.Instant
19+
import java.util.concurrent.TimeUnit
20+
import kotlin.use
21+
22+
val AWS_PRODUCT_CONTEXT_KEY = ContextKey.named<AWSProduct>("pluginDescriptor")
23+
private val PLUGIN_ATTRIBUTE_KEY = AttributeKey.stringKey("plugin")
24+
25+
class DefaultSpanBuilder(delegate: SpanBuilder) : AbstractSpanBuilder<DefaultSpanBuilder, AbstractBaseSpan>(delegate) {
26+
override fun doStartSpan() = BaseSpan(parent!!, delegate.startSpan())
27+
}
28+
29+
abstract class AbstractSpanBuilder<Builder : AbstractSpanBuilder<Builder, Span>, Span : AbstractBaseSpan>(protected val delegate: SpanBuilder) : SpanBuilder {
30+
/**
31+
* Same as [com.intellij.platform.diagnostic.telemetry.helpers.use] except downcasts to specific subclass of [BaseSpan]
32+
*
33+
* @inheritdoc
34+
*/
35+
inline fun<T> use(operation: (Span) -> T): T {
36+
return startSpan().useWithoutActiveScope { span ->
37+
(span as Span).makeCurrent().use {
38+
operation(span)
39+
}
40+
}
41+
}
42+
43+
protected var parent: Context? = null
44+
override fun setParent(context: Context): Builder {
45+
parent = context
46+
delegate.setParent(context)
47+
return this as Builder
48+
}
49+
50+
override fun setNoParent(): Builder {
51+
parent = null
52+
delegate.setNoParent()
53+
return this as Builder
54+
}
55+
56+
override fun addLink(spanContext: SpanContext): Builder {
57+
delegate.addLink(spanContext)
58+
return this as Builder
59+
}
60+
61+
override fun addLink(
62+
spanContext: SpanContext,
63+
attributes: Attributes,
64+
): Builder {
65+
delegate.addLink(spanContext, attributes)
66+
return this as Builder
67+
}
68+
69+
override fun setAttribute(key: String, value: String): Builder {
70+
delegate.setAttribute(key, value)
71+
return this as Builder
72+
}
73+
74+
override fun setAttribute(key: String, value: Long): Builder {
75+
delegate.setAttribute(key, value)
76+
return this as Builder
77+
}
78+
79+
override fun setAttribute(key: String, value: Double): Builder {
80+
delegate.setAttribute(key, value)
81+
return this as Builder
82+
}
83+
84+
override fun setAttribute(key: String, value: Boolean): Builder {
85+
delegate.setAttribute(key, value)
86+
return this as Builder
87+
}
88+
89+
override fun <V : Any?> setAttribute(
90+
key: AttributeKey<V?>,
91+
value: V & Any,
92+
): Builder {
93+
delegate.setAttribute(key, value)
94+
return this as Builder
95+
}
96+
97+
override fun setAllAttributes(attributes: Attributes): Builder {
98+
delegate.setAllAttributes(attributes)
99+
return this as Builder
100+
}
101+
102+
override fun setSpanKind(spanKind: SpanKind): Builder {
103+
delegate.setSpanKind(spanKind)
104+
return this as Builder
105+
}
106+
107+
override fun setStartTimestamp(startTimestamp: Long, unit: TimeUnit): Builder {
108+
delegate.setStartTimestamp(startTimestamp, unit)
109+
return this as Builder
110+
}
111+
112+
override fun setStartTimestamp(startTimestamp: Instant): Builder {
113+
delegate.setStartTimestamp(startTimestamp)
114+
return this as Builder
115+
}
116+
117+
protected abstract fun doStartSpan(): Span
118+
119+
override fun startSpan(): Span {
120+
var parent = parent
121+
if (parent == null) {
122+
parent = Context.current()
123+
}
124+
requireNotNull(parent)
125+
126+
val contextValue = parent.get(AWS_PRODUCT_CONTEXT_KEY)
127+
if (contextValue == null) {
128+
val s = Span.fromContextOrNull(parent)
129+
if (s is AbstractBaseSpan) {
130+
setParent(s.context.with(Span.fromContext(parent)))
131+
} else {
132+
setParent(parent.with(AWS_PRODUCT_CONTEXT_KEY, resolvePluginName()))
133+
}
134+
}
135+
136+
setAttribute(
137+
PLUGIN_ATTRIBUTE_KEY,
138+
(parent.get(AWS_PRODUCT_CONTEXT_KEY) ?: resolvePluginName()).name
139+
)
140+
141+
return doStartSpan()
142+
}
143+
144+
private fun resolvePluginName() = PluginResolver.Companion.fromStackTrace(Thread.currentThread().stackTrace).product
145+
}
146+
147+
abstract class AbstractBaseSpan(internal val context: Context, private val delegate: Span) : Span by delegate {
148+
fun metadata(key: String, value: String) = setAttribute(key, value)
149+
150+
override fun makeCurrent(): Scope =
151+
context.with(this).makeCurrent()
152+
}
153+
154+
/**
155+
* Placeholder; will be generated
156+
*/
157+
class BaseSpan(context: Context, delegate: Span) : AbstractBaseSpan(context, delegate) {
158+
fun reason(reason: String) = metadata("reason", reason)
159+
}

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ import com.intellij.openapi.util.ThrowableComputable
1414
import com.intellij.util.ExceptionUtil
1515
import com.intellij.util.concurrency.AppExecutorUtil
1616
import com.intellij.util.concurrency.Semaphore
17+
import io.opentelemetry.context.Context
1718
import software.aws.toolkits.jetbrains.services.telemetry.PluginResolver
1819
import java.time.Duration
1920
import java.util.concurrent.Future
@@ -81,8 +82,9 @@ fun <T> pluginAwareExecuteOnPooledThread(action: () -> T): Future<T> {
8182
* worker thread will not contain original call stack. Necessary for telemetry.
8283
*/
8384
val pluginResolver = PluginResolver.fromCurrentThread()
85+
val context = Context.current()
8486
return ApplicationManager.getApplication().executeOnPooledThread<T> {
8587
PluginResolver.setThreadLocal(pluginResolver)
86-
action()
88+
context.wrap(action).call()
8789
}
8890
}

0 commit comments

Comments
 (0)