Skip to content

Commit ca24bab

Browse files
rlirli
authored
jetbrains: Generate OpenTelemetry-style metric builders (#899)
## Problem At the moment, metrics in the JetBrains extensions are collected as singular datapoints independent of each other. Modern observability frameworks, as well as our long-term goals are to be able to correlate these data flows over a single user flow. This requires the 'trace ID' concept which is high-friction in the current model. ## Solution Generate 'fluent builders' that wrap the OpenTelemetry SDK to allow us to leverage the SDK to record trace/span IDs automatically, but still expose a good amount of type safety based on the metric schema. This PR references base classes proposed in aws/aws-toolkit-jetbrains#4982. ## License By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
1 parent 9f6a8fb commit ca24bab

File tree

31 files changed

+995
-138
lines changed

31 files changed

+995
-138
lines changed

‎telemetry/jetbrains/.editorconfig

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,4 +9,8 @@ tab_width = 4
99
ij_continuation_indent_size = 4
1010

1111
[*.{kt,kts}]
12-
ktlint_code_style = ktlint_official
12+
ij_kotlin_allow_trailing_comma = true
13+
ij_kotlin_code_style_defaults = KOTLIN_OFFICIAL
14+
ij_kotlin_name_count_to_use_star_import = 2147483647
15+
ij_kotlin_name_count_to_use_star_import_for_members = 2147483647
16+
ij_kotlin_packages_to_use_import_on_demand = unset

‎telemetry/jetbrains/build.gradle.kts

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -39,15 +39,18 @@ dependencies {
3939
implementation(libs.kotlin.poet)
4040
implementation(libs.jackson.module.kotlin)
4141
implementation(libs.json.schema)
42-
testImplementation(libs.junit4)
42+
43+
testImplementation(platform(libs.junit5.bom))
44+
testImplementation(libs.junit5.jupiter)
4345
testImplementation(libs.assertj)
46+
47+
testRuntimeOnly(libs.junit5.launcher)
4448
}
4549

4650
tasks {
4751
withType<KotlinCompile> {
4852
compilerOptions {
4953
jvmTarget = JvmTarget.JVM_17
50-
freeCompilerArgs.add("-Xcontext-receivers")
5154
}
5255
}
5356

@@ -108,6 +111,8 @@ tasks.withType<GenerateModuleMetadata> {
108111
}
109112

110113
tasks.withType<Test> {
114+
useJUnitPlatform()
115+
111116
testLogging {
112117
exceptionFormat = TestExceptionFormat.FULL
113118
}

‎telemetry/jetbrains/gradle/libs.versions.toml

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,17 +4,19 @@ jackson = "2.17.2"
44
jlleitschuh-ktlint = "12.1.1"
55
# deprecated; should move to json-skema
66
jsonSchema = "1.14.4"
7-
junit4 = "4.13.2"
7+
junit5 = "5.11.3"
88
kotlin = "2.0.20"
9-
kotlin-poet = "1.18.1"
9+
kotlin-poet = "2.0.0"
1010
nexus = "2.0.0"
1111

1212
[libraries]
1313
assertj = { module = "org.assertj:assertj-core", version.ref = "assertJ" }
1414
kotlin-poet = { module = "com.squareup:kotlinpoet", version.ref = "kotlin-poet" }
1515
jackson-module-kotlin = { module = "com.fasterxml.jackson.module:jackson-module-kotlin", version.ref = "jackson" }
1616
json-schema = { module = "com.github.erosb:everit-json-schema", version.ref = "jsonSchema" }
17-
junit4 = { module = "junit:junit", version.ref = "junit4" }
17+
junit5-bom = { module = "org.junit:junit-bom", version.ref = "junit5" }
18+
junit5-jupiter = { module = "org.junit.jupiter:junit-jupiter" }
19+
junit5-launcher = { module = "org.junit.platform:junit-platform-launcher" }
1820

1921
[plugins]
2022
jlleitschuh-ktlint = { id = "org.jlleitschuh.gradle.ktlint", version.ref = "jlleitschuh-ktlint" }

‎telemetry/jetbrains/src/main/kotlin/software/aws/toolkits/telemetry/generator/OTelTelemetryGenerator.kt

Lines changed: 332 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,332 @@
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.telemetry.generator
5+
6+
import com.squareup.kotlinpoet.BOOLEAN
7+
import com.squareup.kotlinpoet.ClassName
8+
import com.squareup.kotlinpoet.CodeBlock
9+
import com.squareup.kotlinpoet.FileSpec
10+
import com.squareup.kotlinpoet.FunSpec
11+
import com.squareup.kotlinpoet.KModifier
12+
import com.squareup.kotlinpoet.ParameterizedTypeName.Companion.parameterizedBy
13+
import com.squareup.kotlinpoet.PropertySpec
14+
import com.squareup.kotlinpoet.STRING
15+
import com.squareup.kotlinpoet.TypeName
16+
import com.squareup.kotlinpoet.TypeSpec
17+
import com.squareup.kotlinpoet.TypeVariableName
18+
import java.io.File
19+
20+
object OTelTelemetryGenerator {
21+
private const val PACKAGE_NAME_IMPL = "software.aws.toolkits.telemetry.impl"
22+
23+
private val OTEL_CONTEXT = ClassName("io.opentelemetry.context", "Context")
24+
private val OTEL_TRACER = ClassName("io.opentelemetry.api.trace", "Tracer")
25+
private val OTEL_SPAN_BUILDER = ClassName("io.opentelemetry.api.trace", "SpanBuilder")
26+
private val OTEL_SPAN = ClassName("io.opentelemetry.api.trace", "Span")
27+
private val OTEL_RW_SPAN = ClassName("io.opentelemetry.sdk.trace", "ReadWriteSpan")
28+
29+
private const val TOOLKIT_OTEL_PACKAGE = "software.aws.toolkits.jetbrains.services.telemetry.otel"
30+
private val TOOLKIT_ABSTRACT_BASE_SPAN = ClassName(TOOLKIT_OTEL_PACKAGE, "AbstractBaseSpan")
31+
private val TOOLKIT_ABSTRACT_SPAN_BUILDER = ClassName(TOOLKIT_OTEL_PACKAGE, "AbstractSpanBuilder")
32+
private val TOOLKIT_DEFAULT_SPAN_BUILDER = ClassName(TOOLKIT_OTEL_PACKAGE, "DefaultSpanBuilder")
33+
private val TOOLKIT_OTEL_SERVICE = ClassName(TOOLKIT_OTEL_PACKAGE, "OTelService")
34+
35+
private val SPAN_TYPE_TYPEVAR = TypeVariableName("SpanType")
36+
37+
private val GENERATED_BASE_SPAN = ClassName(PACKAGE_NAME_IMPL, "BaseSpan")
38+
39+
private val indent = " ".repeat(4)
40+
private val commonMetadataTypes =
41+
setOf(
42+
"duration",
43+
"httpStatusCode",
44+
"reason",
45+
"reasonDesc",
46+
"requestId",
47+
"requestServiceType",
48+
"result",
49+
// handled by OpenTelemetry emitter
50+
// "traceId",
51+
// "metricId",
52+
// "parentId",
53+
// handled as special cases in base
54+
// "passive",
55+
// "value",
56+
// "unit",
57+
)
58+
59+
fun generateTelemetryFromFiles(
60+
inputFiles: List<File>,
61+
defaultDefinitions: List<String> = ResourceLoader.DEFINITIONS_FILES,
62+
outputFolder: File,
63+
) {
64+
val telemetryDefinitions = TelemetryParser.parseFiles(defaultDefinitions, inputFiles)
65+
66+
FileSpec.builder(GENERATED_BASE_SPAN)
67+
.indent(indent)
68+
.generateHeader()
69+
.addType(baseSpan(telemetryDefinitions))
70+
.build()
71+
.writeTo(outputFolder)
72+
73+
val telemetryKt =
74+
FileSpec.builder(PACKAGE_NAME, "Telemetry")
75+
.indent(indent)
76+
.generateHeader()
77+
78+
val telemetryRootBuilder = TypeSpec.objectBuilder("Telemetry")
79+
80+
telemetryDefinitions.metrics.groupBy { it.namespace() }
81+
.toSortedMap()
82+
.forEach { (namespace, metrics) ->
83+
generateMetrics(telemetryRootBuilder, outputFolder, namespace, metrics)
84+
}
85+
86+
telemetryKt
87+
.addType(telemetryRootBuilder.build())
88+
.build()
89+
.writeTo(outputFolder)
90+
}
91+
92+
// public open class BaseSpan<SpanType : BaseSpan<SpanType>>(
93+
// context: Context?,
94+
// `delegate`: Span,
95+
// ) : AbstractBaseSpan<SpanType>(context, delegate as ReadWriteSpan) {
96+
private fun baseSpan(telemetryDefinitions: TelemetrySchema) =
97+
TypeSpec.classBuilder(GENERATED_BASE_SPAN)
98+
.addModifiers(KModifier.OPEN)
99+
.primaryConstructor(
100+
FunSpec.constructorBuilder()
101+
.addParameter("context", OTEL_CONTEXT.copy(nullable = true))
102+
.addParameter("delegate", OTEL_SPAN)
103+
.build(),
104+
)
105+
.addTypeVariable(SPAN_TYPE_TYPEVAR.copy(bounds = listOf(GENERATED_BASE_SPAN.parameterizedBy(SPAN_TYPE_TYPEVAR))))
106+
.superclass(
107+
TOOLKIT_ABSTRACT_BASE_SPAN
108+
.parameterizedBy(SPAN_TYPE_TYPEVAR),
109+
)
110+
.addSuperclassConstructorParameter("context, delegate as %T", OTEL_RW_SPAN)
111+
.apply {
112+
commonMetadataTypes.forEach { t ->
113+
val type = telemetryDefinitions.types.firstOrNull { it.name == t } ?: return@forEach
114+
115+
addFunctions(MetadataSchema(type, false).overloadedFunSpec(SPAN_TYPE_TYPEVAR))
116+
}
117+
118+
// special case
119+
addFunction(
120+
FunSpec.builder("success")
121+
.addParameter("success", BOOLEAN)
122+
.returns(SPAN_TYPE_TYPEVAR)
123+
.addStatement("result(if(success) MetricResult.Succeeded else MetricResult.Failed)")
124+
.addStatement("return this as %T", SPAN_TYPE_TYPEVAR)
125+
.build(),
126+
)
127+
}
128+
.build()
129+
130+
private fun generateMetrics(
131+
rootBuilder: TypeSpec.Builder,
132+
outputFolder: File,
133+
namespace: String,
134+
metrics: List<MetricSchema>,
135+
) {
136+
val tracerName = ClassName(PACKAGE_NAME_IMPL, "${namespace.capitalize()}Tracer")
137+
138+
val tracerKt =
139+
FileSpec.builder(tracerName)
140+
.indent(indent)
141+
.generateHeader()
142+
143+
// public class AmazonqTracer internal constructor(
144+
// private val `delegate`: Tracer,
145+
// ) : Tracer {
146+
// /**
147+
// * When user opens CWSPR chat panel
148+
// */
149+
// public val openChat: AmazonqopenChatSpanBuilder
150+
// get() = AmazonqopenChatSpanBuilder(delegate.spanBuilder("amazonq_openChat"))
151+
val tracer =
152+
TypeSpec.classBuilder(tracerName)
153+
.addSuperinterface(OTEL_TRACER)
154+
.primaryConstructor(
155+
FunSpec.constructorBuilder()
156+
.addModifiers(KModifier.INTERNAL)
157+
.addParameter("delegate", OTEL_TRACER)
158+
.build(),
159+
)
160+
.addProperty(
161+
PropertySpec.builder("delegate", OTEL_TRACER)
162+
.initializer("delegate")
163+
.addModifiers(KModifier.PRIVATE)
164+
.build(),
165+
)
166+
.addFunction(
167+
FunSpec.builder("spanBuilder")
168+
.addModifiers(KModifier.OVERRIDE)
169+
.addParameter("spanName", String::class)
170+
.returns(TOOLKIT_DEFAULT_SPAN_BUILDER)
171+
.addStatement("return %T(delegate.spanBuilder(spanName))", TOOLKIT_DEFAULT_SPAN_BUILDER)
172+
.build(),
173+
)
174+
.apply {
175+
metrics.forEach { metricSchema ->
176+
val metricName = metricSchema.name.split("_", limit = 2)[1]
177+
val metricSpanName = ClassName(PACKAGE_NAME_IMPL, "${namespace.capitalize()}${metricName}Span")
178+
val metricSpanBuilderName = ClassName(PACKAGE_NAME_IMPL, "${namespace.capitalize()}${metricName}SpanBuilder")
179+
180+
tracerKt.generateMetricSpan(metricSchema, metricSpanName)
181+
tracerKt.generateMetricSpanBuilder(metricSpanName, metricSpanBuilderName)
182+
183+
// /**
184+
// * When user opens CWSPR chat panel
185+
// */
186+
// public val openChat: AmazonqopenChatSpanBuilder
187+
// get() = AmazonqopenChatSpanBuilder(delegate.spanBuilder("amazonq_openChat"))
188+
addProperty(
189+
PropertySpec.builder(metricName, metricSpanBuilderName)
190+
.getter(
191+
FunSpec.builder("get()")
192+
.addStatement("""return %T(delegate.spanBuilder(%S))""", metricSpanBuilderName, metricSchema.name)
193+
.build(),
194+
)
195+
.addKdoc(metricSchema.description)
196+
.build(),
197+
)
198+
}
199+
}
200+
.build()
201+
202+
tracerKt
203+
.addType(tracer)
204+
.build()
205+
.writeTo(outputFolder)
206+
207+
rootBuilder.addProperty(
208+
PropertySpec.builder(namespace, tracerName)
209+
.getter(
210+
FunSpec.builder("get()")
211+
.addStatement("return %T(%T.getSdk().getTracer(%S))", tracerName, TOOLKIT_OTEL_SERVICE, namespace)
212+
.build(),
213+
)
214+
.build(),
215+
)
216+
}
217+
218+
private fun FileSpec.Builder.generateMetricSpanBuilder(
219+
metricSpanName: ClassName,
220+
metricSpanBuilderName: ClassName,
221+
) {
222+
// public class AmazonqopenChatSpanBuilder internal constructor(
223+
// `delegate`: SpanBuilder,
224+
// ) : AbstractSpanBuilder<AmazonqopenChatSpanBuilder, AmazonqopenChatSpan>(delegate) {
225+
// override fun doStartSpan(): AmazonqopenChatSpan = AmazonqopenChatSpan(parent, delegate.startSpan())
226+
// }
227+
val metricSpanBuilder =
228+
TypeSpec.classBuilder(metricSpanBuilderName)
229+
.primaryConstructor(
230+
FunSpec.constructorBuilder()
231+
.addModifiers(KModifier.INTERNAL)
232+
.addParameter("delegate", OTEL_SPAN_BUILDER)
233+
.build(),
234+
)
235+
.superclass(
236+
TOOLKIT_ABSTRACT_SPAN_BUILDER.parameterizedBy(metricSpanBuilderName, metricSpanName),
237+
)
238+
.addSuperclassConstructorParameter("delegate")
239+
.addFunction(
240+
FunSpec.builder("doStartSpan")
241+
.returns(metricSpanName)
242+
.addModifiers(KModifier.OVERRIDE)
243+
.addStatement("return %T(parent, delegate.startSpan())", metricSpanName)
244+
.build(),
245+
)
246+
.build()
247+
addType(metricSpanBuilder)
248+
}
249+
250+
private fun FileSpec.Builder.generateMetricSpan(
251+
metricSchema: MetricSchema,
252+
metricSpanName: ClassName,
253+
) {
254+
// public class AmazonqopenChatSpan internal constructor(
255+
// context: Context?,
256+
// span: Span,
257+
// ) : BaseSpan<AmazonqopenChatSpan>(context, span) {
258+
// init {
259+
// passive(false)
260+
// }
261+
//
262+
// override val requiredFields: Collection<String> = setOf()
263+
// }
264+
val metricSpan =
265+
TypeSpec.classBuilder(metricSpanName)
266+
.primaryConstructor(
267+
FunSpec.constructorBuilder()
268+
.addModifiers(KModifier.INTERNAL)
269+
.addParameter("context", OTEL_CONTEXT.copy(nullable = true))
270+
.addParameter("span", OTEL_SPAN)
271+
.build(),
272+
)
273+
.addKdoc(metricSchema.description)
274+
.superclass(GENERATED_BASE_SPAN.parameterizedBy(metricSpanName))
275+
.addSuperclassConstructorParameter("context, span")
276+
.apply {
277+
if (!metricSchema.passive) {
278+
addInitializerBlock(CodeBlock.builder().addStatement("passive(false)").build())
279+
}
280+
281+
metricSchema.metadata.filterNot { it.type.name in commonMetadataTypes }.forEach { metadata ->
282+
addFunctions(metadata.overloadedFunSpec(metricSpanName))
283+
}
284+
285+
val requiredAttributes = metricSchema.metadata.filter { it.required != false }
286+
addProperty(
287+
PropertySpec.builder("requiredFields", Collection::class.parameterizedBy(String::class), KModifier.OVERRIDE)
288+
.initializer(
289+
"""setOf(${ "%S,".repeat(requiredAttributes.size) })""",
290+
*requiredAttributes.map { it.type.name }.toTypedArray(),
291+
)
292+
.build(),
293+
)
294+
}
295+
.build()
296+
addType(metricSpan)
297+
}
298+
299+
private fun MetadataSchema.overloadedFunSpec(returnType: TypeName): List<FunSpec> {
300+
val types =
301+
if (type.allowedValues?.isNotEmpty() == true) {
302+
listOf(ClassName(PACKAGE_NAME, type.typeName))
303+
} else {
304+
type.type.kotlinTypes()
305+
}
306+
307+
return types.map { t ->
308+
val needsToString = (t != STRING || type.allowedValues?.isNotEmpty() == true)
309+
val nullable = required == false
310+
311+
FunSpec.builder(type.name)
312+
.addParameter(type.name, t.copy(nullable = nullable))
313+
.returns(returnType)
314+
.apply {
315+
val valueParam =
316+
if (needsToString) {
317+
if (nullable) {
318+
"%N?.let { it.toString() }"
319+
} else {
320+
"%N.toString()"
321+
}
322+
} else {
323+
"%N"
324+
}
325+
326+
addStatement("return metadata(%S, $valueParam)", type.name, type.name)
327+
}
328+
.addKdoc(type.description)
329+
.build()
330+
}
331+
}
332+
}

0 commit comments

Comments
 (0)