fix: ensure coroutines are propagated for accurate spans

nais · Nov 19, 2024 · 99fce34 · 99fce34
1 parent 0ae959f
commit 99fce34
Show file tree

Hide file tree

Showing 6 changed files with 43 additions and 38 deletions.
diff --git a/build.gradle.kts b/build.gradle.kts
@@ -100,6 +100,7 @@ subprojects {
         implementation("io.opentelemetry:opentelemetry-sdk:${opentelemetryVersion}")
         implementation("io.opentelemetry:opentelemetry-sdk-extension-autoconfigure:$opentelemetryVersion");
         implementation("io.opentelemetry:opentelemetry-exporter-otlp:${opentelemetryVersion}")
+        implementation("io.opentelemetry:opentelemetry-extension-kotlin:${opentelemetryVersion}")
         implementation("net.logstash.logback:logstash-logback-encoder:${logstashVersion}")
         runtimeOnly("ch.qos.logback:logback-classic:${logbackVersion}")
     }

diff --git a/wonderwalled-azure/local.env b/wonderwalled-azure/local.env
@@ -6,7 +6,7 @@ AZURE_ENABLED=true
 # expected audience for user token from wonderwall
 AZURE_APP_CLIENT_ID=default
 # this key is useless in production; disregard security notices
-AZURE_APP_CLIENT_JWK='{"p":"_LNnIjBshCrFuxtjUC2KKzg_NTVv26UZh5j12_9r5mYTxb8yW047jOYFEGvIdMkTRLGOBig6fLWzgd62lnLainzV35J6K6zr4jQfTldLondlkldMR6nQrp1KfnNUuRbKvzpNKkhl12-f1l91l0tCx3s4blztvWgdzN2xBfvWV68","kty":"RSA","q":"9MIWsbIA3WjiR_Ful5FM8NCgb6JdS2D6ySHVepoNI-iAPilcltF_J2orjfLqAxeztTskPi45wtF_-eV4GIYSzvMo-gFiXLMrvEa7WaWizMi_7Bu9tEk3m_f3IDLN9lwULYoebkDbiXx6GOiuj0VkuKz8ckYFNKLCMP9QRLFff-0","d":"J6UX848X8tNz-09PFvcFDUVqak32GXzoPjnuDjBsxNUvG7LxenLmM_i8tvYl0EW9Ztn4AiCqJUoHw5cX3jz_mSqGl7ciaDedpKm_AetcZwHiEuT1EpSKRPMmOMQSqcJqXrdbbWB8gdUrnTKZIlJCfj7yqgT16ypC43TnwjA0UwxhG5pHaYjKI3pPdoHg2BzA-iubHjVn15Sz7-pnjBmeGDbEFa7ADY-1yPHCmqqvPKTNhoCNW6RpG34Id9hXslPa3X-7pAhJrDBd0_NPlktSA2rUkifYiZURhHR5ijhe0v3uw6kYP8f_foVm_C8O1ExkxXh9Dg8KDZ89dbsSOtBc0Q","e":"AQAB","use":"sig","kid":"l7C_WJgbZ_6e59vPrFETAehX7Dsp7fIyvSV4XhotsGs","qi":"cQFN5q5WhYkzgd1RS0rGqvpX1AkmZMrLv2MW04gSfu0dDwpbsSAu8EUCQW9oA4pr6V7R9CBSu9kdN2iY5SR-hZvEad5nDKPV1F3TMQYv5KpRiS_0XhfV5PcolUJVO_4p3h8d-mo2hh1Sw2fairAKOzvnwJCQ6DFkiY7H1cqwA54","dp":"YTql9AGtvyy158gh7jeXcgmySEbHQzvDFulDr-IXIg8kjHGEbp0rTIs0Z50RA95aC5RFkRjpaBKBfvaySjDm5WIi6GLzntpp6B8l7H6qG1jVO_la4Df2kzjx8LVvY8fhOrKz_hDdHodUeKdCF3RdvWMr00ruLnJhBPJHqoW7cwE","alg":"RS256","dq":"IZA4AngRbEtEtG7kJn6zWVaSmZxfRMXwvgIYvy4-3Qy2AVA0tS3XTPVfMaD8_B2U9CY_CxPVseR-sysHc_12uNBZbycfcOzU84WTjXCMSZ7BysPnGMDtkkLHra-p1L29upz1HVNhh5H9QEswHM98R2LZX2ZAsn4bORLZ1AGqweU","n":"8ZqUp5Cs90XpNn8tJBdUUxdGH4bjqKjFj8lyB3x50RpTuECuwzX1NpVqyFENDiEtMja5fdmJl6SErjnhj6kbhcmfmFibANuG-0WlV5yMysdSbocd75C1JQbiPdpHdXrijmVFMfDnoZTQ-ErNsqqngTNkn5SXBcPenli6Cf9MTSchZuh_qFj_B7Fp3CWKehTiyBcLlNOIjYsXX8WQjZkWKGpQ23AWjZulngWRektLcRWuEKTWaRBtbAr3XAfSmcqTICrebaD3IMWKHDtvzHAt_pt4wnZ06clgeO2Wbc980usnpsF7g8k9p81RcbS4JEZmuuA9NCmOmbyADXwgA9_-Aw"}'
+AZURE_APP_JWK='{"p":"_LNnIjBshCrFuxtjUC2KKzg_NTVv26UZh5j12_9r5mYTxb8yW047jOYFEGvIdMkTRLGOBig6fLWzgd62lnLainzV35J6K6zr4jQfTldLondlkldMR6nQrp1KfnNUuRbKvzpNKkhl12-f1l91l0tCx3s4blztvWgdzN2xBfvWV68","kty":"RSA","q":"9MIWsbIA3WjiR_Ful5FM8NCgb6JdS2D6ySHVepoNI-iAPilcltF_J2orjfLqAxeztTskPi45wtF_-eV4GIYSzvMo-gFiXLMrvEa7WaWizMi_7Bu9tEk3m_f3IDLN9lwULYoebkDbiXx6GOiuj0VkuKz8ckYFNKLCMP9QRLFff-0","d":"J6UX848X8tNz-09PFvcFDUVqak32GXzoPjnuDjBsxNUvG7LxenLmM_i8tvYl0EW9Ztn4AiCqJUoHw5cX3jz_mSqGl7ciaDedpKm_AetcZwHiEuT1EpSKRPMmOMQSqcJqXrdbbWB8gdUrnTKZIlJCfj7yqgT16ypC43TnwjA0UwxhG5pHaYjKI3pPdoHg2BzA-iubHjVn15Sz7-pnjBmeGDbEFa7ADY-1yPHCmqqvPKTNhoCNW6RpG34Id9hXslPa3X-7pAhJrDBd0_NPlktSA2rUkifYiZURhHR5ijhe0v3uw6kYP8f_foVm_C8O1ExkxXh9Dg8KDZ89dbsSOtBc0Q","e":"AQAB","use":"sig","kid":"l7C_WJgbZ_6e59vPrFETAehX7Dsp7fIyvSV4XhotsGs","qi":"cQFN5q5WhYkzgd1RS0rGqvpX1AkmZMrLv2MW04gSfu0dDwpbsSAu8EUCQW9oA4pr6V7R9CBSu9kdN2iY5SR-hZvEad5nDKPV1F3TMQYv5KpRiS_0XhfV5PcolUJVO_4p3h8d-mo2hh1Sw2fairAKOzvnwJCQ6DFkiY7H1cqwA54","dp":"YTql9AGtvyy158gh7jeXcgmySEbHQzvDFulDr-IXIg8kjHGEbp0rTIs0Z50RA95aC5RFkRjpaBKBfvaySjDm5WIi6GLzntpp6B8l7H6qG1jVO_la4Df2kzjx8LVvY8fhOrKz_hDdHodUeKdCF3RdvWMr00ruLnJhBPJHqoW7cwE","alg":"RS256","dq":"IZA4AngRbEtEtG7kJn6zWVaSmZxfRMXwvgIYvy4-3Qy2AVA0tS3XTPVfMaD8_B2U9CY_CxPVseR-sysHc_12uNBZbycfcOzU84WTjXCMSZ7BysPnGMDtkkLHra-p1L29upz1HVNhh5H9QEswHM98R2LZX2ZAsn4bORLZ1AGqweU","n":"8ZqUp5Cs90XpNn8tJBdUUxdGH4bjqKjFj8lyB3x50RpTuECuwzX1NpVqyFENDiEtMja5fdmJl6SErjnhj6kbhcmfmFibANuG-0WlV5yMysdSbocd75C1JQbiPdpHdXrijmVFMfDnoZTQ-ErNsqqngTNkn5SXBcPenli6Cf9MTSchZuh_qFj_B7Fp3CWKehTiyBcLlNOIjYsXX8WQjZkWKGpQ23AWjZulngWRektLcRWuEKTWaRBtbAr3XAfSmcqTICrebaD3IMWKHDtvzHAt_pt4wnZ06clgeO2Wbc980usnpsF7g8k9p81RcbS4JEZmuuA9NCmOmbyADXwgA9_-Aw"}'
 AZURE_OPENID_CONFIG_ISSUER=http://localhost:7070/azure
 AZURE_OPENID_CONFIG_TOKEN_ENDPOINT=http://localhost:7070/azure/token
 AZURE_OPENID_CONFIG_JWKS_URI=http://localhost:7070/azure/jwks
diff --git a/wonderwalled-common/src/main/kotlin/io/nais/common/Auth.kt b/wonderwalled-common/src/main/kotlin/io/nais/common/Auth.kt
@@ -22,6 +22,7 @@ import io.ktor.server.request.uri
 import io.ktor.server.response.respondRedirect
 import io.opentelemetry.api.OpenTelemetry
 import io.opentelemetry.api.common.AttributeKey
+import io.opentelemetry.api.common.Attributes
 import io.opentelemetry.api.trace.Tracer
 import org.slf4j.Logger
 import org.slf4j.LoggerFactory
@@ -81,43 +82,47 @@ data class TokenIntrospectionResponse(
 class AuthClient(
     private val config: AuthClientConfig,
     private val provider: IdentityProvider,
-    private val openTelemetry: OpenTelemetry? = null,
+    private val openTelemetry: OpenTelemetry = OpenTelemetry.noop(),
     private val httpClient: HttpClient = defaultHttpClient(openTelemetry),
 ) {
-    private val tracer: Tracer? = openTelemetry?.getTracer("io.nais.common.AuthClient")
+    private val tracer: Tracer = openTelemetry.getTracer("io.nais.common.AuthClient")
 
     suspend fun token(target: String) =
-        tracer.withSpan("auth/token", traceAttributes(target)) {
+        tracer.withSpan("auth/token", parameters = {
+            setAllAttributes(traceAttributes(target))
+        }) {
             httpClient.post(config.tokenEndpoint) {
                 contentType(ContentType.Application.Json)
                 setBody(TokenRequest(target, provider))
             }.body<TokenResponse>()
         }
 
     suspend fun exchange(target: String, userToken: String) =
-        tracer.withSpan("auth/exchange", traceAttributes(target)) {
+        tracer.withSpan("auth/exchange", parameters = {
+            setAllAttributes(traceAttributes(target))
+        }) {
             httpClient.post(config.tokenExchangeEndpoint) {
                 contentType(ContentType.Application.Json)
                 setBody(TokenExchangeRequest(target, provider, userToken))
             }.body<TokenResponse>()
         }
 
     suspend fun introspect(accessToken: String) =
-        tracer.withSpan("auth/introspect", traceAttributes()) {
+        tracer.withSpan("auth/introspect", parameters = {
+            setAllAttributes(traceAttributes())
+        }) {
             httpClient.post(config.tokenIntrospectionEndpoint) {
                 contentType(ContentType.Application.Json)
                 setBody(TokenIntrospectionRequest(accessToken, provider))
             }.body<TokenIntrospectionResponse>()
         }
 
-    private fun traceAttributes(target: String? = null) =
-        mutableMapOf(
-            attributeKeyIdentityProvider to provider.alias,
-        ).apply {
-            if (target != null) {
-                this[attributeKeyTarget] = target
-            }
+    private fun traceAttributes(target: String? = null) = Attributes.builder().apply {
+        put(attributeKeyIdentityProvider, provider.alias)
+        if (target != null) {
+            put(attributeKeyTarget, target)
         }
+    }.build()
 
     companion object {
         private val attributeKeyTarget: AttributeKey<String> = AttributeKey.stringKey("target")
@@ -219,7 +224,8 @@ data class OpenIdConfiguration(
 private fun invalidOpenIdConfigurationException(
     expected: List<String>,
     got: String,
-): RuntimeException = RuntimeException("authority does not match the issuer returned by provider: got $got, expected one of $expected")
+): RuntimeException =
+    RuntimeException("authority does not match the issuer returned by provider: got $got, expected one of $expected")
 
 @JsonIgnoreProperties(ignoreUnknown = true)
 data class AccessToken(

diff --git a/wonderwalled-common/src/main/kotlin/io/nais/common/Client.kt b/wonderwalled-common/src/main/kotlin/io/nais/common/Client.kt
@@ -26,14 +26,6 @@ fun defaultHttpClient(
     openTelemetry?.let {
         install(KtorClientTracing) {
             setOpenTelemetry(it)
-            attributeExtractor {
-                onStart {
-                    attributes.put("start-time", System.currentTimeMillis())
-                }
-                onEnd {
-                    attributes.put("end-time", System.currentTimeMillis())
-                }
-            }
         }
     }
 }

diff --git a/wonderwalled-common/src/main/kotlin/io/nais/common/Tracing.kt b/wonderwalled-common/src/main/kotlin/io/nais/common/Tracing.kt
@@ -1,10 +1,15 @@
 package io.nais.common
 
 import io.opentelemetry.api.OpenTelemetry
-import io.opentelemetry.api.common.AttributeKey
+import io.opentelemetry.api.trace.Span
+import io.opentelemetry.api.trace.SpanBuilder
+import io.opentelemetry.api.trace.StatusCode
 import io.opentelemetry.api.trace.Tracer
+import io.opentelemetry.extension.kotlin.asContextElement
 import io.opentelemetry.sdk.autoconfigure.AutoConfiguredOpenTelemetrySdk
 import io.opentelemetry.semconv.ServiceAttributes
+import kotlinx.coroutines.withContext
+import kotlin.coroutines.coroutineContext
 
 fun openTelemetry(serviceName: String): OpenTelemetry {
     return AutoConfiguredOpenTelemetrySdk.builder().addResourceCustomizer { oldResource, _ ->
@@ -15,24 +20,25 @@ fun openTelemetry(serviceName: String): OpenTelemetry {
     }.build().openTelemetrySdk
 }
 
-suspend fun <T> Tracer?.withSpan(
+suspend fun <T> Tracer.withSpan(
     spanName: String,
-    attributes: Map<AttributeKey<String>, String> = emptyMap(),
-    block: suspend () -> T
+    parameters: (SpanBuilder.() -> Unit)? = null,
+    block: suspend (span: Span) -> T
 ): T {
-    val span = this?.spanBuilder(spanName)?.startSpan()?.apply {
-        attributes.forEach { (key, value) -> this.setAttribute(key, value) }
+    val span: Span = this.spanBuilder(spanName).run {
+        if (parameters != null) parameters()
+        startSpan()
     }
 
-    if (this == null || span == null) {
-        return block()
-    }
-
-    return try {
-        span.makeCurrent().use { _ ->
-            block()
+    return withContext(coroutineContext + span.asContextElement()) {
+        try {
+            block(span)
+        } catch (throwable: Throwable) {
+            span.setStatus(StatusCode.ERROR)
+            span.recordException(throwable)
+            throw throwable
+        } finally {
+            span.end()
         }
-    } finally {
-        span.end()
     }
 }
diff --git a/wonderwalled-maskinporten/local.env b/wonderwalled-maskinporten/local.env
@@ -14,7 +14,7 @@ AZURE_ENABLED=true
 # expected audience for user token from wonderwall
 AZURE_APP_CLIENT_ID=default
 # this key is useless in production; disregard security notices
-AZURE_APP_CLIENT_JWK='{"p":"_LNnIjBshCrFuxtjUC2KKzg_NTVv26UZh5j12_9r5mYTxb8yW047jOYFEGvIdMkTRLGOBig6fLWzgd62lnLainzV35J6K6zr4jQfTldLondlkldMR6nQrp1KfnNUuRbKvzpNKkhl12-f1l91l0tCx3s4blztvWgdzN2xBfvWV68","kty":"RSA","q":"9MIWsbIA3WjiR_Ful5FM8NCgb6JdS2D6ySHVepoNI-iAPilcltF_J2orjfLqAxeztTskPi45wtF_-eV4GIYSzvMo-gFiXLMrvEa7WaWizMi_7Bu9tEk3m_f3IDLN9lwULYoebkDbiXx6GOiuj0VkuKz8ckYFNKLCMP9QRLFff-0","d":"J6UX848X8tNz-09PFvcFDUVqak32GXzoPjnuDjBsxNUvG7LxenLmM_i8tvYl0EW9Ztn4AiCqJUoHw5cX3jz_mSqGl7ciaDedpKm_AetcZwHiEuT1EpSKRPMmOMQSqcJqXrdbbWB8gdUrnTKZIlJCfj7yqgT16ypC43TnwjA0UwxhG5pHaYjKI3pPdoHg2BzA-iubHjVn15Sz7-pnjBmeGDbEFa7ADY-1yPHCmqqvPKTNhoCNW6RpG34Id9hXslPa3X-7pAhJrDBd0_NPlktSA2rUkifYiZURhHR5ijhe0v3uw6kYP8f_foVm_C8O1ExkxXh9Dg8KDZ89dbsSOtBc0Q","e":"AQAB","use":"sig","kid":"l7C_WJgbZ_6e59vPrFETAehX7Dsp7fIyvSV4XhotsGs","qi":"cQFN5q5WhYkzgd1RS0rGqvpX1AkmZMrLv2MW04gSfu0dDwpbsSAu8EUCQW9oA4pr6V7R9CBSu9kdN2iY5SR-hZvEad5nDKPV1F3TMQYv5KpRiS_0XhfV5PcolUJVO_4p3h8d-mo2hh1Sw2fairAKOzvnwJCQ6DFkiY7H1cqwA54","dp":"YTql9AGtvyy158gh7jeXcgmySEbHQzvDFulDr-IXIg8kjHGEbp0rTIs0Z50RA95aC5RFkRjpaBKBfvaySjDm5WIi6GLzntpp6B8l7H6qG1jVO_la4Df2kzjx8LVvY8fhOrKz_hDdHodUeKdCF3RdvWMr00ruLnJhBPJHqoW7cwE","alg":"RS256","dq":"IZA4AngRbEtEtG7kJn6zWVaSmZxfRMXwvgIYvy4-3Qy2AVA0tS3XTPVfMaD8_B2U9CY_CxPVseR-sysHc_12uNBZbycfcOzU84WTjXCMSZ7BysPnGMDtkkLHra-p1L29upz1HVNhh5H9QEswHM98R2LZX2ZAsn4bORLZ1AGqweU","n":"8ZqUp5Cs90XpNn8tJBdUUxdGH4bjqKjFj8lyB3x50RpTuECuwzX1NpVqyFENDiEtMja5fdmJl6SErjnhj6kbhcmfmFibANuG-0WlV5yMysdSbocd75C1JQbiPdpHdXrijmVFMfDnoZTQ-ErNsqqngTNkn5SXBcPenli6Cf9MTSchZuh_qFj_B7Fp3CWKehTiyBcLlNOIjYsXX8WQjZkWKGpQ23AWjZulngWRektLcRWuEKTWaRBtbAr3XAfSmcqTICrebaD3IMWKHDtvzHAt_pt4wnZ06clgeO2Wbc980usnpsF7g8k9p81RcbS4JEZmuuA9NCmOmbyADXwgA9_-Aw"}'
+AZURE_APP_JWK='{"p":"_LNnIjBshCrFuxtjUC2KKzg_NTVv26UZh5j12_9r5mYTxb8yW047jOYFEGvIdMkTRLGOBig6fLWzgd62lnLainzV35J6K6zr4jQfTldLondlkldMR6nQrp1KfnNUuRbKvzpNKkhl12-f1l91l0tCx3s4blztvWgdzN2xBfvWV68","kty":"RSA","q":"9MIWsbIA3WjiR_Ful5FM8NCgb6JdS2D6ySHVepoNI-iAPilcltF_J2orjfLqAxeztTskPi45wtF_-eV4GIYSzvMo-gFiXLMrvEa7WaWizMi_7Bu9tEk3m_f3IDLN9lwULYoebkDbiXx6GOiuj0VkuKz8ckYFNKLCMP9QRLFff-0","d":"J6UX848X8tNz-09PFvcFDUVqak32GXzoPjnuDjBsxNUvG7LxenLmM_i8tvYl0EW9Ztn4AiCqJUoHw5cX3jz_mSqGl7ciaDedpKm_AetcZwHiEuT1EpSKRPMmOMQSqcJqXrdbbWB8gdUrnTKZIlJCfj7yqgT16ypC43TnwjA0UwxhG5pHaYjKI3pPdoHg2BzA-iubHjVn15Sz7-pnjBmeGDbEFa7ADY-1yPHCmqqvPKTNhoCNW6RpG34Id9hXslPa3X-7pAhJrDBd0_NPlktSA2rUkifYiZURhHR5ijhe0v3uw6kYP8f_foVm_C8O1ExkxXh9Dg8KDZ89dbsSOtBc0Q","e":"AQAB","use":"sig","kid":"l7C_WJgbZ_6e59vPrFETAehX7Dsp7fIyvSV4XhotsGs","qi":"cQFN5q5WhYkzgd1RS0rGqvpX1AkmZMrLv2MW04gSfu0dDwpbsSAu8EUCQW9oA4pr6V7R9CBSu9kdN2iY5SR-hZvEad5nDKPV1F3TMQYv5KpRiS_0XhfV5PcolUJVO_4p3h8d-mo2hh1Sw2fairAKOzvnwJCQ6DFkiY7H1cqwA54","dp":"YTql9AGtvyy158gh7jeXcgmySEbHQzvDFulDr-IXIg8kjHGEbp0rTIs0Z50RA95aC5RFkRjpaBKBfvaySjDm5WIi6GLzntpp6B8l7H6qG1jVO_la4Df2kzjx8LVvY8fhOrKz_hDdHodUeKdCF3RdvWMr00ruLnJhBPJHqoW7cwE","alg":"RS256","dq":"IZA4AngRbEtEtG7kJn6zWVaSmZxfRMXwvgIYvy4-3Qy2AVA0tS3XTPVfMaD8_B2U9CY_CxPVseR-sysHc_12uNBZbycfcOzU84WTjXCMSZ7BysPnGMDtkkLHra-p1L29upz1HVNhh5H9QEswHM98R2LZX2ZAsn4bORLZ1AGqweU","n":"8ZqUp5Cs90XpNn8tJBdUUxdGH4bjqKjFj8lyB3x50RpTuECuwzX1NpVqyFENDiEtMja5fdmJl6SErjnhj6kbhcmfmFibANuG-0WlV5yMysdSbocd75C1JQbiPdpHdXrijmVFMfDnoZTQ-ErNsqqngTNkn5SXBcPenli6Cf9MTSchZuh_qFj_B7Fp3CWKehTiyBcLlNOIjYsXX8WQjZkWKGpQ23AWjZulngWRektLcRWuEKTWaRBtbAr3XAfSmcqTICrebaD3IMWKHDtvzHAt_pt4wnZ06clgeO2Wbc980usnpsF7g8k9p81RcbS4JEZmuuA9NCmOmbyADXwgA9_-Aw"}'
 AZURE_OPENID_CONFIG_ISSUER=http://localhost:7070/azure
 AZURE_OPENID_CONFIG_TOKEN_ENDPOINT=http://localhost:7070/azure/token
 AZURE_OPENID_CONFIG_JWKS_URI=http://localhost:7070/azure/jwks