address review comment

Author: laurit

instrumentation/ktor/ktor-2-common/library/src/main/kotlin/io/opentelemetry/instrumentation/ktor/internal/KtorBuilderUtil.kt

@@ -12,7 +12,7 @@ import io.ktor.server.response.*
 import io.opentelemetry.instrumentation.api.incubator.builder.internal.DefaultHttpClientInstrumenterBuilder
 import io.opentelemetry.instrumentation.api.incubator.builder.internal.DefaultHttpServerInstrumenterBuilder
 import io.opentelemetry.instrumentation.ktor.client.AbstractKtorClientTracingBuilder
-import io.opentelemetry.instrumentation.ktor.server.AbstractKtorServerTracing
+import io.opentelemetry.instrumentation.ktor.server.AbstractKtorServerTracingBuilder
 
 /**
  * This class is internal and is hence not for public use. Its APIs are unstable and can change at
@@ -21,6 +21,6 @@ import io.opentelemetry.instrumentation.ktor.server.AbstractKtorServerTracing
 object KtorBuilderUtil {
   lateinit var clientBuilderExtractor: (AbstractKtorClientTracingBuilder) -> DefaultHttpClientInstrumenterBuilder<HttpRequestData, HttpResponse>
   lateinit var serverBuilderExtractor: (
-    AbstractKtorServerTracing.Configuration
+    AbstractKtorServerTracingBuilder
   ) -> DefaultHttpServerInstrumenterBuilder<ApplicationRequest, ApplicationResponse>
 }

instrumentation/ktor/ktor-2-common/library/src/main/kotlin/io/opentelemetry/instrumentation/ktor/internal/KtorServerTracer.kt

@@ -0,0 +1,29 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.ktor.internal
+
+import io.ktor.server.application.*
+import io.ktor.server.request.*
+import io.ktor.server.response.*
+import io.opentelemetry.context.Context
+import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter
+
+class KtorServerTracer(
+  private val instrumenter: Instrumenter<ApplicationRequest, ApplicationResponse>,
+) {
+  fun start(call: ApplicationCall): Context? {
+    val parentContext = Context.current()
+    if (!instrumenter.shouldStart(parentContext, call.request)) {
+      return null
+    }
+
+    return instrumenter.start(parentContext, call.request)
+  }
+
+  fun end(context: Context, call: ApplicationCall, error: Throwable?) {
+    instrumenter.end(context, call.request, call.response, error)
+  }
+}

instrumentation/ktor/ktor-2-common/library/src/main/kotlin/io/opentelemetry/instrumentation/ktor/internal/KtorServerTracingUtil.kt

@@ -5,25 +5,79 @@
 
 package io.opentelemetry.instrumentation.ktor.internal
 
+import io.ktor.server.application.*
 import io.ktor.server.request.*
 import io.ktor.server.response.*
+import io.ktor.util.*
+import io.ktor.util.pipeline.*
+import io.opentelemetry.context.Context
+import io.opentelemetry.extension.kotlin.asContextElement
 import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter
 import io.opentelemetry.instrumentation.api.instrumenter.SpanKindExtractor
 import io.opentelemetry.instrumentation.api.internal.InstrumenterUtil
-import io.opentelemetry.instrumentation.ktor.server.AbstractKtorServerTracing
+import io.opentelemetry.instrumentation.ktor.server.AbstractKtorServerTracingBuilder
 import io.opentelemetry.instrumentation.ktor.server.ApplicationRequestGetter
+import kotlinx.coroutines.withContext
 
 /**
  * This class is internal and is hence not for public use. Its APIs are unstable and can change at
  * any time.
  */
 object KtorServerTracingUtil {
 
-  fun instrumenter(configuration: AbstractKtorServerTracing.Configuration): Instrumenter<ApplicationRequest, ApplicationResponse> {
+  fun configureTracing(builder: AbstractKtorServerTracingBuilder, application: Application) {
+    val contextKey = AttributeKey<Context>("OpenTelemetry")
+    val errorKey = AttributeKey<Throwable>("OpenTelemetryException")
+
+    val instrumenter = instrumenter(builder)
+    val tracer = KtorServerTracer(instrumenter)
+    val startPhase = PipelinePhase("OpenTelemetry")
+
+    application.insertPhaseBefore(ApplicationCallPipeline.Monitoring, startPhase)
+    application.intercept(startPhase) {
+      val context = tracer.start(call)
+
+      if (context != null) {
+        call.attributes.put(contextKey, context)
+        withContext(context.asContextElement()) {
+          try {
+            proceed()
+          } catch (err: Throwable) {
+            // Stash error for reporting later since need ktor to finish setting up the response
+            call.attributes.put(errorKey, err)
+            throw err
+          }
+        }
+      } else {
+        proceed()
+      }
+    }
+
+    val postSendPhase = PipelinePhase("OpenTelemetryPostSend")
+    application.sendPipeline.insertPhaseAfter(ApplicationSendPipeline.After, postSendPhase)
+    application.sendPipeline.intercept(postSendPhase) {
+      val context = call.attributes.getOrNull(contextKey)
+      if (context != null) {
+        var error: Throwable? = call.attributes.getOrNull(errorKey)
+        try {
+          proceed()
+        } catch (t: Throwable) {
+          error = t
+          throw t
+        } finally {
+          tracer.end(context, call, error)
+        }
+      } else {
+        proceed()
+      }
+    }
+  }
+
+  private fun instrumenter(builder: AbstractKtorServerTracingBuilder): Instrumenter<ApplicationRequest, ApplicationResponse> {
     return InstrumenterUtil.buildUpstreamInstrumenter(
-      configuration.serverBuilder.instrumenterBuilder(),
+      builder.serverBuilder.instrumenterBuilder(),
       ApplicationRequestGetter,
-      configuration.spanKindExtractor(SpanKindExtractor.alwaysServer())
+      builder.spanKindExtractor(SpanKindExtractor.alwaysServer())
     )
   }
 }