Skip to content

Commit 838b641

Browse files
Merge branch 'main' into samgst/emitSessionStart
2 parents f6f6914 + 06181dd commit 838b641

File tree

9 files changed

+790
-14
lines changed

9 files changed

+790
-14
lines changed

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileContextProvider.kt

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -144,9 +144,10 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi
144144
}
145145

146146
return supplementalContext?.let {
147+
val latency = System.currentTimeMillis() - startFetchingTimestamp
147148
if (it.contents.isNotEmpty()) {
148149
val logStr = buildString {
149-
append("Successfully fetched supplemental context with strategy ${it.strategy} with ${it.latency} ms")
150+
append("Successfully fetched supplemental context with strategy ${it.strategy} with $latency ms")
150151
it.contents.forEachIndexed { index, chunk ->
151152
append(
152153
"""
@@ -166,8 +167,7 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi
166167
LOG.warn { "Failed to fetch supplemental context, empty list." }
167168
}
168169

169-
// TODO: fix this latency
170-
it.copy(latency = System.currentTimeMillis() - startFetchingTimestamp)
170+
it.copy(latency = latency)
171171
}
172172
} catch (e: TimeoutCancellationException) {
173173
LOG.debug {

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()
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/core/credentials/sso/DiskCache.kt

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ import software.aws.toolkits.core.utils.touch
3232
import software.aws.toolkits.core.utils.tryDirOp
3333
import software.aws.toolkits.core.utils.tryFileOp
3434
import software.aws.toolkits.core.utils.tryOrNull
35+
import software.aws.toolkits.core.utils.warn
3536
import software.aws.toolkits.telemetry.AuthTelemetry
3637
import software.aws.toolkits.telemetry.Result
3738
import java.io.InputStream
@@ -105,7 +106,7 @@ class DiskCache(
105106
val inputStream = clientRegistrationCache(cacheKey).tryInputStreamIfExists()
106107
if (inputStream == null) {
107108
val stage = LoadCredentialStage.ACCESS_FILE
108-
LOG.warn("Failed to load Client Registration: cache file does not exist")
109+
LOG.warn { "Failed to load Client Registration: cache file does not exist" }
109110
AuthTelemetry.modifyConnection(
110111
action = "Load cache file",
111112
source = "loadClientRegistration",
@@ -224,7 +225,7 @@ class DiskCache(
224225
if (clientRegistration.expiresAt.isNotExpired()) {
225226
return clientRegistration
226227
} else {
227-
LOG.warn("Client Registration is expired")
228+
LOG.warn { "Client Registration is expired" }
228229
AuthTelemetry.modifyConnection(
229230
action = "Validate Credentials",
230231
source = "loadClientRegistration",
@@ -235,7 +236,7 @@ class DiskCache(
235236
return null
236237
}
237238
} catch (e: Exception) {
238-
LOG.warn("Client Registration could not be read")
239+
LOG.warn { "Client Registration could not be read" }
239240
AuthTelemetry.modifyConnection(
240241
action = "Validate Credentials",
241242
source = "loadClientRegistration",

plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/SsoAccessTokenProvider.kt

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -176,7 +176,7 @@ class SsoAccessTokenProvider(
176176
result = Result.Succeeded
177177
)
178178
} catch (e: Exception) {
179-
getLogger<SsoAccessTokenProvider>().warn("Failed to save access token ${e.message}")
179+
LOG.warn { "Failed to save access token ${e.message}" }
180180
AuthTelemetry.modifyConnection(
181181
action = "Write file",
182182
source = "accessToken",
@@ -218,7 +218,7 @@ class SsoAccessTokenProvider(
218218
result = Result.Succeeded
219219
)
220220
} catch (e: Exception) {
221-
getLogger<SsoAccessTokenProvider>().warn("Failed to save client registration ${e.message}")
221+
LOG.warn { "Failed to save client registration ${e.message}" }
222222
AuthTelemetry.modifyConnection(
223223
action = "Write file",
224224
source = "registerDAGClient",
@@ -238,7 +238,7 @@ class SsoAccessTokenProvider(
238238
}
239239

240240
if (!ssoUrl.contains("identitycenter")) {
241-
getLogger<SsoAccessTokenProvider>().warn { "$ssoUrl does not appear to be a valid issuer URL" }
241+
LOG.warn { "$ssoUrl does not appear to be a valid issuer URL" }
242242
}
243243

244244
val registerResponse = client.registerClient {
@@ -270,7 +270,7 @@ class SsoAccessTokenProvider(
270270
result = Result.Succeeded
271271
)
272272
} catch (e: Exception) {
273-
getLogger<SsoAccessTokenProvider>().warn("Failed to save client registration${e.message}")
273+
LOG.warn { "Failed to save client registration${e.message}" }
274274
AuthTelemetry.modifyConnection(
275275
action = "Write file",
276276
source = "registerPkceClient",
@@ -498,7 +498,7 @@ class SsoAccessTokenProvider(
498498
requestId = requestId,
499499
result = Result.Failed
500500
)
501-
getLogger<SsoAccessTokenProvider>().warn("RefreshAccessTokenFailed: ${e.message}")
501+
LOG.warn { "RefreshAccessTokenFailed: ${e.message}" }
502502
throw e
503503
}
504504
}
@@ -577,5 +577,6 @@ class SsoAccessTokenProvider(
577577
// Default number of seconds to poll for token, https://tools.ietf.org/html/draft-ietf-oauth-device-flow-15#section-3.5
578578
const val DEFAULT_INTERVAL_SECS = 5L
579579
const val SLOW_DOWN_DELAY_SECS = 5L
580+
private val LOG = getLogger<SsoAccessTokenProvider>()
580581
}
581582
}
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)