Skip to content

Commit a8a28ac

Browse files
marychattemarychatte
committed
Add extensions for server v2
1 parent 080f857 commit a8a28ac

File tree

4 files changed

+130
-16
lines changed

4 files changed

+130
-16
lines changed

‎instrumentation/ktor/ktor-2.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/ktor/v2_0/ServerInstrumentation.java

+3-3
Original file line numberDiff line numberDiff line change
@@ -50,9 +50,9 @@ public static class SetupFunction
5050
public Unit invoke(KtorServerTracing.Configuration configuration) {
5151
OpenTelemetry openTelemetry = GlobalOpenTelemetry.get();
5252
configuration.setOpenTelemetry(openTelemetry);
53-
configuration.setCapturedRequestHeaders(CommonConfig.get().getServerRequestHeaders());
54-
configuration.setCapturedResponseHeaders(CommonConfig.get().getServerResponseHeaders());
55-
configuration.setKnownMethods(CommonConfig.get().getKnownHttpRequestMethods());
53+
configuration.capturedRequestHeaders(CommonConfig.get().getServerRequestHeaders());
54+
configuration.capturedResponseHeaders(CommonConfig.get().getServerResponseHeaders());
55+
configuration.knownMethods(CommonConfig.get().getKnownHttpRequestMethods());
5656

5757
return kotlin.Unit.INSTANCE;
5858
}

‎instrumentation/ktor/ktor-2.0/library/src/main/kotlin/io/opentelemetry/instrumentation/ktor/v2_0/server/KtorServerTracing.kt

+120-3
Original file line numberDiff line numberDiff line change
@@ -5,18 +5,22 @@
55

66
package io.opentelemetry.instrumentation.ktor.v2_0.server
77

8+
import io.ktor.http.*
89
import io.ktor.server.application.*
910
import io.ktor.server.request.*
1011
import io.ktor.server.response.*
1112
import io.ktor.server.routing.*
1213
import io.ktor.util.*
1314
import io.ktor.util.pipeline.*
1415
import io.opentelemetry.api.OpenTelemetry
16+
import io.opentelemetry.api.common.AttributesBuilder
17+
import io.opentelemetry.api.trace.SpanKind
1518
import io.opentelemetry.context.Context
1619
import io.opentelemetry.extension.kotlin.asContextElement
1720
import io.opentelemetry.instrumentation.api.instrumenter.AttributesExtractor
1821
import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter
1922
import io.opentelemetry.instrumentation.api.instrumenter.SpanKindExtractor
23+
import io.opentelemetry.instrumentation.api.instrumenter.SpanStatusBuilder
2024
import io.opentelemetry.instrumentation.api.instrumenter.SpanStatusExtractor
2125
import io.opentelemetry.instrumentation.api.internal.InstrumenterUtil
2226
import io.opentelemetry.instrumentation.api.semconv.http.HttpServerAttributesExtractor
@@ -59,28 +63,143 @@ class KtorServerTracing private constructor(
5963
this.statusExtractor = extractor
6064
}
6165

66+
fun spanStatusExtractor(extract: SpanStatusData.() -> Unit) {
67+
setStatusExtractor {
68+
SpanStatusExtractor<ApplicationRequest, ApplicationResponse> { spanStatusBuilder: SpanStatusBuilder,
69+
request: ApplicationRequest,
70+
response: ApplicationResponse?,
71+
throwable: Throwable? ->
72+
extract(SpanStatusData(spanStatusBuilder, request, response, throwable))
73+
}
74+
}
75+
}
76+
77+
data class SpanStatusData(
78+
val spanStatusBuilder: SpanStatusBuilder,
79+
val request: ApplicationRequest,
80+
val response: ApplicationResponse?,
81+
val error: Throwable?
82+
)
83+
6284
fun setSpanKindExtractor(extractor: (SpanKindExtractor<ApplicationRequest>) -> SpanKindExtractor<ApplicationRequest>) {
6385
this.spanKindExtractor = extractor
6486
}
6587

88+
fun spanKindExtractor(extract: ApplicationRequest.() -> SpanKind) {
89+
setSpanKindExtractor {
90+
SpanKindExtractor<ApplicationRequest> { request: ApplicationRequest ->
91+
extract(request)
92+
}
93+
}
94+
}
95+
6696
fun addAttributeExtractor(extractor: AttributesExtractor<in ApplicationRequest, in ApplicationResponse>) {
6797
additionalExtractors.add(extractor)
6898
}
6999

100+
fun attributeExtractor(
101+
extractorBuilder: ExtractorBuilder.() -> Unit = {}
102+
) {
103+
val builder = ExtractorBuilder().apply(extractorBuilder).build()
104+
addAttributeExtractor(
105+
object : AttributesExtractor<ApplicationRequest, ApplicationResponse> {
106+
override fun onStart(attributes: AttributesBuilder, parentContext: Context, request: ApplicationRequest) {
107+
builder.onStart(OnStartData(attributes, parentContext, request))
108+
}
109+
110+
override fun onEnd(
111+
attributes: AttributesBuilder,
112+
context: Context,
113+
request: ApplicationRequest,
114+
response: ApplicationResponse?,
115+
error: Throwable?
116+
) {
117+
builder.onEnd(OnEndData(attributes, context, request, response, error))
118+
}
119+
}
120+
)
121+
}
122+
123+
class ExtractorBuilder {
124+
private var onStart: OnStartData.() -> Unit = {}
125+
private var onEnd: OnEndData.() -> Unit = {}
126+
127+
fun onStart(block: OnStartData.() -> Unit) {
128+
onStart = block
129+
}
130+
131+
fun onEnd(block: OnEndData.() -> Unit) {
132+
onEnd = block
133+
}
134+
135+
internal fun build(): Extractor {
136+
return Extractor(onStart, onEnd)
137+
}
138+
}
139+
140+
internal class Extractor(val onStart: OnStartData.() -> Unit, val onEnd: OnEndData.() -> Unit)
141+
142+
data class OnStartData(
143+
val attributes: AttributesBuilder,
144+
val parentContext: Context,
145+
val request: ApplicationRequest
146+
)
147+
148+
data class OnEndData(
149+
val attributes: AttributesBuilder,
150+
val parentContext: Context,
151+
val request: ApplicationRequest,
152+
val response: ApplicationResponse?,
153+
val error: Throwable?
154+
)
155+
70156
fun setCapturedRequestHeaders(requestHeaders: List<String>) {
71157
httpAttributesExtractorBuilder.setCapturedRequestHeaders(requestHeaders)
72158
}
73159

160+
fun capturedRequestHeaders(headers: Iterable<String>) {
161+
setCapturedRequestHeaders(headers.toList())
162+
}
163+
164+
fun capturedRequestHeaders(vararg headers: String) {
165+
capturedRequestHeaders(headers.asIterable())
166+
}
167+
74168
fun setCapturedResponseHeaders(responseHeaders: List<String>) {
75169
httpAttributesExtractorBuilder.setCapturedResponseHeaders(responseHeaders)
76170
}
77171

172+
fun capturedResponseHeaders(headers: Iterable<String>) {
173+
setCapturedResponseHeaders(headers.toList())
174+
}
175+
176+
fun capturedResponseHeaders(vararg headers: String) {
177+
capturedResponseHeaders(headers.asIterable())
178+
}
179+
78180
fun setKnownMethods(knownMethods: Set<String>) {
79181
httpAttributesExtractorBuilder.setKnownMethods(knownMethods)
80182
httpSpanNameExtractorBuilder.setKnownMethods(knownMethods)
81183
httpServerRouteBuilder.setKnownMethods(knownMethods)
82184
}
83185

186+
fun knownMethods(vararg methods: String) {
187+
setKnownMethods(methods.toSet())
188+
}
189+
190+
fun knownMethods(methods: Iterable<String>) {
191+
setKnownMethods(methods.toSet())
192+
}
193+
194+
@JvmName("knownMethodsJvm")
195+
fun knownMethods(methods: Iterable<HttpMethod>) {
196+
knownMethods(methods.map { it.value })
197+
}
198+
199+
fun knownMethods(vararg methods: HttpMethod) {
200+
knownMethods(methods.map { it.value })
201+
}
202+
84203
internal fun isOpenTelemetryInitialized(): Boolean = this::openTelemetry.isInitialized
85204
}
86205

@@ -107,9 +226,7 @@ class KtorServerTracing private constructor(
107226
override fun install(pipeline: Application, configure: Configuration.() -> Unit): KtorServerTracing {
108227
val configuration = Configuration().apply(configure)
109228

110-
if (!configuration.isOpenTelemetryInitialized()) {
111-
throw IllegalArgumentException("OpenTelemetry must be set")
112-
}
229+
require(configuration.isOpenTelemetryInitialized()) { "OpenTelemetry must be set" }
113230

114231
val httpAttributesGetter = KtorHttpServerAttributesGetter.INSTANCE
115232

‎instrumentation/ktor/ktor-2.0/library/src/test/kotlin/io/opentelemetry/instrumentation/ktor/v2_0/server/KtorServerSpanKindExtractorTest.kt

+5-8
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@ import io.ktor.server.request.*
1313
import io.ktor.server.response.*
1414
import io.ktor.server.routing.*
1515
import io.opentelemetry.api.trace.SpanKind
16-
import io.opentelemetry.instrumentation.api.instrumenter.SpanKindExtractor
1716
import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension
1817
import io.opentelemetry.instrumentation.testing.junit.http.AbstractHttpServerUsingTest
1918
import io.opentelemetry.instrumentation.testing.junit.http.HttpServerInstrumentationExtension
@@ -60,13 +59,11 @@ class KtorServerSpanKindExtractorTest : AbstractHttpServerUsingTest<ApplicationE
6059
return embeddedServer(Netty, port = port) {
6160
install(KtorServerTracing) {
6261
setOpenTelemetry(testing.openTelemetry)
63-
setSpanKindExtractor {
64-
SpanKindExtractor { req ->
65-
if (req.uri.startsWith("/from-pubsub/")) {
66-
SpanKind.CONSUMER
67-
} else {
68-
SpanKind.SERVER
69-
}
62+
spanKindExtractor {
63+
if (uri.startsWith("/from-pubsub/")) {
64+
SpanKind.CONSUMER
65+
} else {
66+
SpanKind.SERVER
7067
}
7168
}
7269
}

‎instrumentation/ktor/ktor-2.0/library/src/test/kotlin/io/opentelemetry/instrumentation/ktor/v2_0/server/KtorTestUtil.kt

+2-2
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,8 @@ class KtorTestUtil {
1414
fun installOpenTelemetry(application: Application, openTelemetry: OpenTelemetry) {
1515
application.install(KtorServerTracing) {
1616
setOpenTelemetry(openTelemetry)
17-
setCapturedRequestHeaders(listOf(AbstractHttpServerTest.TEST_REQUEST_HEADER))
18-
setCapturedResponseHeaders(listOf(AbstractHttpServerTest.TEST_RESPONSE_HEADER))
17+
capturedRequestHeaders(AbstractHttpServerTest.TEST_REQUEST_HEADER)
18+
capturedResponseHeaders(AbstractHttpServerTest.TEST_RESPONSE_HEADER)
1919
}
2020
}
2121
}

0 commit comments

Comments
 (0)