Skip to content

Commit 06181dd

Browse files
rlirli
authored
Add OpenTelemetry base classes (#4982)
The OpenTelemetry SDK can automatically attach the same trace ID to spans recorded under the same 'context'. This is somewhat challenging because the JetBrains platform frequently requires switching threads and coroutine scope contexts, which loses the `ThreadLocal` state needed to connect parent->child when spans are started. This on its own is not too difficult, but becomes a significant ergonomic challenge when we also want to bring some type-safety to how span metadata is built. This PR allows spans generated in aws/aws-toolkit-common#899 to automatically perform context propagation. See internal documentation for details. A contributor-friendly version of the doc will be written after we validate this pattern with some initial metrics.
1 parent c765d61 commit 06181dd

File tree

6 files changed

+777
-3
lines changed

6 files changed

+777
-3
lines changed

‎plugins/core/jetbrains-community/build.gradle.kts‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,15 +18,22 @@ buildscript {
1818
}
1919
}
2020

21+
private val generatedSrcDir = project.layout.buildDirectory.dir("generated-src")
2122
sourceSets {
2223
main {
23-
java.srcDir(project.layout.buildDirectory.dir("generated-src"))
24+
java.srcDir(generatedSrcDir)
25+
}
26+
}
27+
28+
idea {
29+
module {
30+
generatedSourceDirs = generatedSourceDirs.toMutableSet() + generatedSrcDir.get().asFile
2431
}
2532
}
2633

2734
val generateTelemetry = tasks.register<GenerateTelemetry>("generateTelemetry") {
2835
inputFiles = listOf(file("${project.projectDir}/resources/telemetryOverride.json"))
29-
outputDirectory = project.layout.buildDirectory.dir("generated-src").get().asFile
36+
outputDirectory = generatedSrcDir.get().asFile
3037

3138
doFirst {
3239
outputDirectory.deleteRecursively()

‎plugins/core/jetbrains-community/src-233/com/intellij/platform/diagnostic/telemetry/helpers/UseWithoutActiveScope.kt‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
// Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
package com.intellij.platform.diagnostic.telemetry.helpers
5+
6+
import io.opentelemetry.api.common.AttributeKey
7+
import io.opentelemetry.api.common.Attributes
8+
import io.opentelemetry.api.trace.Span
9+
import io.opentelemetry.api.trace.StatusCode
10+
import kotlin.coroutines.cancellation.CancellationException
11+
12+
val EXCEPTION_ESCAPED = AttributeKey.booleanKey("exception.escaped")
13+
14+
inline fun <T> Span.useWithoutActiveScope(operation: (Span) -> T): T {
15+
try {
16+
return operation(this)
17+
} catch (e: CancellationException) {
18+
recordException(e, Attributes.of(EXCEPTION_ESCAPED, true))
19+
throw e
20+
} catch (e: Throwable) {
21+
recordException(e, Attributes.of(EXCEPTION_ESCAPED, true))
22+
setStatus(StatusCode.ERROR)
23+
throw e
24+
} finally {
25+
end()
26+
}
27+
}

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

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

0 commit comments

Comments
 (0)