Add extension functions for Ktor plugins (#10963)

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

@@ -56,9 +56,9 @@ public class HttpClientInstrumentation implements TypeInstrumentation {
     public Unit invoke(KtorClientTracingBuilder builder) {
       OpenTelemetry openTelemetry = GlobalOpenTelemetry.get();
       builder.setOpenTelemetry(openTelemetry);
-      builder.setCapturedRequestHeaders(CommonConfig.get().getClientRequestHeaders());
-      builder.setCapturedResponseHeaders(CommonConfig.get().getClientResponseHeaders());
-      builder.setKnownMethods(CommonConfig.get().getKnownHttpRequestMethods());
+      builder.capturedRequestHeaders(CommonConfig.get().getClientRequestHeaders());
+      builder.capturedResponseHeaders(CommonConfig.get().getClientResponseHeaders());
+      builder.knownMethods(CommonConfig.get().getKnownHttpRequestMethods());
 
       return kotlin.Unit.INSTANCE;
     }

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

@@ -50,9 +50,9 @@ public class ServerInstrumentation implements TypeInstrumentation {
     public Unit invoke(KtorServerTracing.Configuration configuration) {
       OpenTelemetry openTelemetry = GlobalOpenTelemetry.get();
       configuration.setOpenTelemetry(openTelemetry);
-      configuration.setCapturedRequestHeaders(CommonConfig.get().getServerRequestHeaders());
-      configuration.setCapturedResponseHeaders(CommonConfig.get().getServerResponseHeaders());
-      configuration.setKnownMethods(CommonConfig.get().getKnownHttpRequestMethods());
+      configuration.capturedRequestHeaders(CommonConfig.get().getServerRequestHeaders());
+      configuration.capturedResponseHeaders(CommonConfig.get().getServerResponseHeaders());
+      configuration.knownMethods(CommonConfig.get().getKnownHttpRequestMethods());
 
       return kotlin.Unit.INSTANCE;
     }

instrumentation/ktor/ktor-2.0/library/src/main/kotlin/io/opentelemetry/instrumentation/ktor/v2_0/client/KtorClientTracingBuilder.kt

@@ -7,7 +7,10 @@ package io.opentelemetry.instrumentation.ktor.v2_0.client
 
 import io.ktor.client.request.*
 import io.ktor.client.statement.*
+import io.ktor.http.*
 import io.opentelemetry.api.OpenTelemetry
+import io.opentelemetry.api.common.AttributesBuilder
+import io.opentelemetry.context.Context
 import io.opentelemetry.instrumentation.api.incubator.semconv.http.HttpClientExperimentalMetrics
 import io.opentelemetry.instrumentation.api.incubator.semconv.http.HttpExperimentalAttributesExtractor
 import io.opentelemetry.instrumentation.api.instrumenter.AttributesExtractor
@@ -31,36 +34,137 @@ class KtorClientTracingBuilder {
     this.openTelemetry = openTelemetry
   }
 
-  fun setCapturedRequestHeaders(vararg headers: String) = setCapturedRequestHeaders(headers.asList())
+  @Deprecated(
+    "Please use method `capturedRequestHeaders`",
+    ReplaceWith("capturedRequestHeaders(headers.asIterable())")
+  )
+  fun setCapturedRequestHeaders(vararg headers: String) = capturedRequestHeaders(headers.asIterable())
 
-  fun setCapturedRequestHeaders(headers: List<String>) {
-    httpAttributesExtractorBuilder.setCapturedRequestHeaders(headers)
+  @Deprecated(
+    "Please use method `capturedRequestHeaders`",
+    ReplaceWith("capturedRequestHeaders(headers)")
+  )
+  fun setCapturedRequestHeaders(headers: List<String>) = capturedRequestHeaders(headers)
+
+  fun capturedRequestHeaders(vararg headers: String) = capturedRequestHeaders(headers.asIterable())
+
+  fun capturedRequestHeaders(headers: Iterable<String>) {
+    httpAttributesExtractorBuilder.setCapturedRequestHeaders(headers.toList())
   }
 
-  fun setCapturedResponseHeaders(vararg headers: String) = setCapturedResponseHeaders(headers.asList())
+  @Deprecated(
+    "Please use method `capturedResponseHeaders`",
+    ReplaceWith("capturedResponseHeaders(headers.asIterable())")
+  )
+  fun setCapturedResponseHeaders(vararg headers: String) = capturedResponseHeaders(headers.asIterable())
+
+  @Deprecated(
+    "Please use method `capturedResponseHeaders`",
+    ReplaceWith("capturedResponseHeaders(headers)")
+  )
+  fun setCapturedResponseHeaders(headers: List<String>) = capturedResponseHeaders(headers)
+
+  fun capturedResponseHeaders(vararg headers: String) = capturedResponseHeaders(headers.asIterable())
 
-  fun setCapturedResponseHeaders(headers: List<String>) {
-    httpAttributesExtractorBuilder.setCapturedResponseHeaders(headers)
+  fun capturedResponseHeaders(headers: Iterable<String>) {
+    httpAttributesExtractorBuilder.setCapturedResponseHeaders(headers.toList())
   }
 
-  fun setKnownMethods(knownMethods: Set<String>) {
-    httpAttributesExtractorBuilder.setKnownMethods(knownMethods)
-    httpSpanNameExtractorBuilder.setKnownMethods(knownMethods)
+  @Deprecated(
+    "Please use method `knownMethods`",
+    ReplaceWith("knownMethods(knownMethods)")
+  )
+  fun setKnownMethods(knownMethods: Set<String>) = knownMethods(knownMethods)
+
+  fun knownMethods(vararg methods: String) = knownMethods(methods.asIterable())
+
+  fun knownMethods(vararg methods: HttpMethod) = knownMethods(methods.asIterable())
+
+  @JvmName("knownMethodsJvm")
+  fun knownMethods(methods: Iterable<HttpMethod>) = knownMethods(methods.map { it.value })
+
+  fun knownMethods(methods: Iterable<String>) {
+    methods.toSet().apply {
+      httpAttributesExtractorBuilder.setKnownMethods(this)
+      httpSpanNameExtractorBuilder.setKnownMethods(this)
+    }
   }
 
+  @Deprecated("Please use method `attributeExtractor`")
   fun addAttributesExtractors(vararg extractors: AttributesExtractor<in HttpRequestData, in HttpResponse>) = addAttributesExtractors(extractors.asList())
 
+  @Deprecated("Please use method `attributeExtractor`")
   fun addAttributesExtractors(extractors: Iterable<AttributesExtractor<in HttpRequestData, in HttpResponse>>) {
-    additionalExtractors += extractors
+    extractors.forEach {
+      attributeExtractor {
+        onStart { it.onStart(attributes, parentContext, request) }
+        onEnd { it.onEnd(attributes, parentContext, request, response, error) }
+      }
+    }
+  }
+
+  fun attributeExtractor(extractorBuilder: ExtractorBuilder.() -> Unit = {}) {
+    val builder = ExtractorBuilder().apply(extractorBuilder).build()
+    additionalExtractors.add(
+      object : AttributesExtractor<HttpRequestData, HttpResponse> {
+        override fun onStart(attributes: AttributesBuilder, parentContext: Context, request: HttpRequestData) {
+          builder.onStart(OnStartData(attributes, parentContext, request))
+        }
+
+        override fun onEnd(attributes: AttributesBuilder, context: Context, request: HttpRequestData, response: HttpResponse?, error: Throwable?) {
+          builder.onEnd(OnEndData(attributes, context, request, response, error))
+        }
+      }
+    )
+  }
+
+  class ExtractorBuilder {
+    private var onStart: OnStartData.() -> Unit = {}
+    private var onEnd: OnEndData.() -> Unit = {}
+
+    fun onStart(block: OnStartData.() -> Unit) {
+      onStart = block
+    }
+
+    fun onEnd(block: OnEndData.() -> Unit) {
+      onEnd = block
+    }
+
+    internal fun build(): Extractor {
+      return Extractor(onStart, onEnd)
+    }
   }
 
+  internal class Extractor(val onStart: OnStartData.() -> Unit, val onEnd: OnEndData.() -> Unit)
+
+  data class OnStartData(
+    val attributes: AttributesBuilder,
+    val parentContext: Context,
+    val request: HttpRequestData
+  )
+
+  data class OnEndData(
+    val attributes: AttributesBuilder,
+    val parentContext: Context,
+    val request: HttpRequestData,
+    val response: HttpResponse?,
+    val error: Throwable?
+  )
+
   /**
    * Configures the instrumentation to emit experimental HTTP client metrics.
    *
    * @param emitExperimentalHttpClientMetrics `true` if the experimental HTTP client metrics are to be emitted.
    */
+  @Deprecated("Please use method `emitExperimentalHttpClientMetrics`")
   fun setEmitExperimentalHttpClientMetrics(emitExperimentalHttpClientMetrics: Boolean) {
-    this.emitExperimentalHttpClientMetrics = emitExperimentalHttpClientMetrics
+    if (emitExperimentalHttpClientMetrics) {
+      emitExperimentalHttpClientMetrics()
+    }
+  }
+
+  fun emitExperimentalHttpClientMetrics() {
+    emitExperimentalHttpClientMetrics = true
   }
 
   internal fun build(): KtorClientTracing {

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

@@ -5,6 +5,7 @@
 
 package io.opentelemetry.instrumentation.ktor.v2_0.server
 
+import io.ktor.http.*
 import io.ktor.server.application.*
 import io.ktor.server.request.*
 import io.ktor.server.response.*
@@ -12,11 +13,14 @@ import io.ktor.server.routing.*
 import io.ktor.util.*
 import io.ktor.util.pipeline.*
 import io.opentelemetry.api.OpenTelemetry
+import io.opentelemetry.api.common.AttributesBuilder
+import io.opentelemetry.api.trace.SpanKind
 import io.opentelemetry.context.Context
 import io.opentelemetry.extension.kotlin.asContextElement
 import io.opentelemetry.instrumentation.api.instrumenter.AttributesExtractor
 import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter
 import io.opentelemetry.instrumentation.api.instrumenter.SpanKindExtractor
+import io.opentelemetry.instrumentation.api.instrumenter.SpanStatusBuilder
 import io.opentelemetry.instrumentation.api.instrumenter.SpanStatusExtractor
 import io.opentelemetry.instrumentation.api.internal.InstrumenterUtil
 import io.opentelemetry.instrumentation.api.semconv.http.HttpServerAttributesExtractor
@@ -53,32 +57,151 @@ class KtorServerTracing private constructor(
       this.openTelemetry = openTelemetry
     }
 
+    @Deprecated("Please use method `spanStatusExtractor`")
     fun setStatusExtractor(
       extractor: (SpanStatusExtractor<ApplicationRequest, ApplicationResponse>) -> SpanStatusExtractor<in ApplicationRequest, in ApplicationResponse>
     ) {
-      this.statusExtractor = extractor
+      spanStatusExtractor { prevStatusExtractor ->
+        extractor(prevStatusExtractor).extract(spanStatusBuilder, request, response, error)
+      }
+    }
+
+    fun spanStatusExtractor(extract: SpanStatusData.(SpanStatusExtractor<ApplicationRequest, ApplicationResponse>) -> Unit) {
+      statusExtractor = { prevExtractor ->
+        SpanStatusExtractor<ApplicationRequest, ApplicationResponse> { spanStatusBuilder: SpanStatusBuilder,
+                                                                       request: ApplicationRequest,
+                                                                       response: ApplicationResponse?,
+                                                                       throwable: Throwable? ->
+          extract(SpanStatusData(spanStatusBuilder, request, response, throwable), prevExtractor)
+        }
+      }
     }
 
+    data class SpanStatusData(
+      val spanStatusBuilder: SpanStatusBuilder,
+      val request: ApplicationRequest,
+      val response: ApplicationResponse?,
+      val error: Throwable?
+    )
+
+    @Deprecated("Please use method `spanKindExtractor`")
     fun setSpanKindExtractor(extractor: (SpanKindExtractor<ApplicationRequest>) -> SpanKindExtractor<ApplicationRequest>) {
-      this.spanKindExtractor = extractor
+      spanKindExtractor { prevSpanKindExtractor ->
+        extractor(prevSpanKindExtractor).extract(this)
+      }
+    }
+
+    fun spanKindExtractor(extract: ApplicationRequest.(SpanKindExtractor<ApplicationRequest>) -> SpanKind) {
+      spanKindExtractor = { prevExtractor ->
+        SpanKindExtractor<ApplicationRequest> { request: ApplicationRequest ->
+          extract(request, prevExtractor)
+        }
+      }
     }
 
+    @Deprecated("Please use method `attributeExtractor`")
     fun addAttributeExtractor(extractor: AttributesExtractor<in ApplicationRequest, in ApplicationResponse>) {
-      additionalExtractors.add(extractor)
+      attributeExtractor {
+        onStart {
+          extractor.onStart(attributes, parentContext, request)
+        }
+        onEnd {
+          extractor.onEnd(attributes, parentContext, request, response, error)
+        }
+      }
+    }
+
+    fun attributeExtractor(extractorBuilder: ExtractorBuilder.() -> Unit = {}) {
+      val builder = ExtractorBuilder().apply(extractorBuilder).build()
+      additionalExtractors.add(
+        object : AttributesExtractor<ApplicationRequest, ApplicationResponse> {
+          override fun onStart(attributes: AttributesBuilder, parentContext: Context, request: ApplicationRequest) {
+            builder.onStart(OnStartData(attributes, parentContext, request))
+          }
+
+          override fun onEnd(attributes: AttributesBuilder, context: Context, request: ApplicationRequest, response: ApplicationResponse?, error: Throwable?) {
+            builder.onEnd(OnEndData(attributes, context, request, response, error))
+          }
+        }
+      )
     }
 
-    fun setCapturedRequestHeaders(requestHeaders: List<String>) {
-      httpAttributesExtractorBuilder.setCapturedRequestHeaders(requestHeaders)
+    class ExtractorBuilder {
+      private var onStart: OnStartData.() -> Unit = {}
+      private var onEnd: OnEndData.() -> Unit = {}
+
+      fun onStart(block: OnStartData.() -> Unit) {
+        onStart = block
+      }
+
+      fun onEnd(block: OnEndData.() -> Unit) {
+        onEnd = block
+      }
+
+      internal fun build(): Extractor {
+        return Extractor(onStart, onEnd)
+      }
     }
 
-    fun setCapturedResponseHeaders(responseHeaders: List<String>) {
-      httpAttributesExtractorBuilder.setCapturedResponseHeaders(responseHeaders)
+    internal class Extractor(val onStart: OnStartData.() -> Unit, val onEnd: OnEndData.() -> Unit)
+
+    data class OnStartData(
+      val attributes: AttributesBuilder,
+      val parentContext: Context,
+      val request: ApplicationRequest
+    )
+
+    data class OnEndData(
+      val attributes: AttributesBuilder,
+      val parentContext: Context,
+      val request: ApplicationRequest,
+      val response: ApplicationResponse?,
+      val error: Throwable?
+    )
+
+    @Deprecated(
+      "Please use method `capturedRequestHeaders`",
+      ReplaceWith("capturedRequestHeaders(headers)")
+    )
+    fun setCapturedRequestHeaders(headers: List<String>) = capturedRequestHeaders(headers)
+
+    fun capturedRequestHeaders(vararg headers: String) = capturedRequestHeaders(headers.asIterable())
+
+    fun capturedRequestHeaders(headers: Iterable<String>) {
+      httpAttributesExtractorBuilder.setCapturedRequestHeaders(headers.toList())
     }
 
-    fun setKnownMethods(knownMethods: Set<String>) {
-      httpAttributesExtractorBuilder.setKnownMethods(knownMethods)
-      httpSpanNameExtractorBuilder.setKnownMethods(knownMethods)
-      httpServerRouteBuilder.setKnownMethods(knownMethods)
+    @Deprecated(
+      "Please use method `capturedResponseHeaders`",
+      ReplaceWith("capturedResponseHeaders(headers)")
+    )
+    fun setCapturedResponseHeaders(headers: List<String>) = capturedResponseHeaders(headers)
+
+    fun capturedResponseHeaders(vararg headers: String) = capturedResponseHeaders(headers.asIterable())
+
+    fun capturedResponseHeaders(headers: Iterable<String>) {
+      httpAttributesExtractorBuilder.setCapturedResponseHeaders(headers.toList())
+    }
+
+    @Deprecated(
+      "Please use method `knownMethods`",
+      ReplaceWith("knownMethods(knownMethods)")
+    )
+    fun setKnownMethods(knownMethods: Set<String>) = knownMethods(knownMethods)
+
+    fun knownMethods(vararg methods: String) = knownMethods(methods.asIterable())
+
+    fun knownMethods(vararg methods: HttpMethod) = knownMethods(methods.asIterable())
+
+    @JvmName("knownMethodsJvm")
+    fun knownMethods(methods: Iterable<HttpMethod>) = knownMethods(methods.map { it.value })
+
+    fun knownMethods(methods: Iterable<String>) {
+      methods.toSet().apply {
+        httpAttributesExtractorBuilder.setKnownMethods(this)
+        httpSpanNameExtractorBuilder.setKnownMethods(this)
+        httpServerRouteBuilder.setKnownMethods(this)
+      }
     }
 
     internal fun isOpenTelemetryInitialized(): Boolean = this::openTelemetry.isInitialized
@@ -107,9 +230,7 @@ class KtorServerTracing private constructor(
     override fun install(pipeline: Application, configure: Configuration.() -> Unit): KtorServerTracing {
       val configuration = Configuration().apply(configure)
 
-      if (!configuration.isOpenTelemetryInitialized()) {
-        throw IllegalArgumentException("OpenTelemetry must be set")
-      }
+      require(configuration.isOpenTelemetryInitialized()) { "OpenTelemetry must be set" }
 
       val httpAttributesGetter = KtorHttpServerAttributesGetter.INSTANCE
 

instrumentation/ktor/ktor-2.0/library/src/test/kotlin/io/opentelemetry/instrumentation/ktor/v2_0/client/KtorHttpClientTest.kt

@@ -20,8 +20,8 @@ class KtorHttpClientTest : AbstractKtorHttpClientTest() {
   override fun HttpClientConfig<*>.installTracing() {
     install(KtorClientTracing) {
       setOpenTelemetry(TESTING.openTelemetry)
-      setCapturedRequestHeaders(listOf(TEST_REQUEST_HEADER))
-      setCapturedResponseHeaders(listOf(TEST_RESPONSE_HEADER))
+      capturedRequestHeaders(TEST_REQUEST_HEADER)
+      capturedResponseHeaders(TEST_RESPONSE_HEADER)
     }
   }
 }

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

@@ -13,7 +13,6 @@ import io.ktor.server.request.*
 import io.ktor.server.response.*
 import io.ktor.server.routing.*
 import io.opentelemetry.api.trace.SpanKind
-import io.opentelemetry.instrumentation.api.instrumenter.SpanKindExtractor
 import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension
 import io.opentelemetry.instrumentation.testing.junit.http.AbstractHttpServerUsingTest
 import io.opentelemetry.instrumentation.testing.junit.http.HttpServerInstrumentationExtension
@@ -60,13 +59,11 @@ class KtorServerSpanKindExtractorTest : AbstractHttpServerUsingTest<ApplicationE
     return embeddedServer(Netty, port = port) {
       install(KtorServerTracing) {
         setOpenTelemetry(testing.openTelemetry)
-        setSpanKindExtractor {
-          SpanKindExtractor { req ->
-            if (req.uri.startsWith("/from-pubsub/")) {
-              SpanKind.CONSUMER
-            } else {
-              SpanKind.SERVER
-            }
+        spanKindExtractor {
+          if (uri.startsWith("/from-pubsub/")) {
+            SpanKind.CONSUMER
+          } else {
+            SpanKind.SERVER
           }
         }
       }

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

@@ -14,8 +14,8 @@ class KtorTestUtil {
     fun installOpenTelemetry(application: Application, openTelemetry: OpenTelemetry) {
       application.install(KtorServerTracing) {
         setOpenTelemetry(openTelemetry)
-        setCapturedRequestHeaders(listOf(AbstractHttpServerTest.TEST_REQUEST_HEADER))
-        setCapturedResponseHeaders(listOf(AbstractHttpServerTest.TEST_RESPONSE_HEADER))
+        capturedRequestHeaders(AbstractHttpServerTest.TEST_REQUEST_HEADER)
+        capturedResponseHeaders(AbstractHttpServerTest.TEST_RESPONSE_HEADER)
       }
     }
   }