Replace Span.makeCurrent() with Span.asContextElement()

gradle/libs.versions.toml

@@ -57,6 +57,7 @@ okhttp4 = { module = "com.squareup.okhttp3:okhttp", version.ref = "okhttp4-versi
 okhttp-coroutines = { module = "com.squareup.okhttp3:okhttp-coroutines", version.ref = "okhttp-version" }
 opentelemetry-api = { module = "io.opentelemetry:opentelemetry-api", version.ref = "otel-version" }
 opentelemetry-sdk-testing = {module = "io.opentelemetry:opentelemetry-sdk-testing", version.ref = "otel-version" }
+opentelemetry-kotlin-extension = { module = "io.opentelemetry:opentelemetry-extension-kotlin", version.ref = "otel-version" }
 slf4j-api = { module = "org.slf4j:slf4j-api", version.ref = "slf4j-version" }
 slf4j-api-v1x = { module = "org.slf4j:slf4j-api", version.ref = "slf4j-v1x-version" }
 slf4j-simple = { module = "org.slf4j:slf4j-simple", version.ref = "slf4j-version" }

runtime/observability/telemetry-api/api/telemetry-api.api

@@ -364,6 +364,7 @@ public final class aws/smithy/kotlin/runtime/telemetry/metrics/UpDownCounter$Def
 
 public abstract class aws/smithy/kotlin/runtime/telemetry/trace/AbstractTraceSpan : aws/smithy/kotlin/runtime/telemetry/trace/TraceSpan {
 	public fun <init> ()V
+	public fun asContextElement ()Lkotlin/coroutines/CoroutineContext;
 	public fun close ()V
 	public fun emitEvent (Ljava/lang/String;Laws/smithy/kotlin/runtime/collections/Attributes;)V
 	public fun getSpanContext ()Laws/smithy/kotlin/runtime/telemetry/trace/SpanContext;
@@ -424,6 +425,7 @@ public final class aws/smithy/kotlin/runtime/telemetry/trace/SpanStatus : java/l
 
 public abstract interface class aws/smithy/kotlin/runtime/telemetry/trace/TraceSpan : aws/smithy/kotlin/runtime/telemetry/context/Scope {
 	public static final field Companion Laws/smithy/kotlin/runtime/telemetry/trace/TraceSpan$Companion;
+	public abstract fun asContextElement ()Lkotlin/coroutines/CoroutineContext;
 	public abstract fun close ()V
 	public abstract fun emitEvent (Ljava/lang/String;Laws/smithy/kotlin/runtime/collections/Attributes;)V
 	public abstract fun getSpanContext ()Laws/smithy/kotlin/runtime/telemetry/trace/SpanContext;

runtime/observability/telemetry-api/common/src/aws/smithy/kotlin/runtime/telemetry/trace/AbstractTraceSpan.kt

@@ -6,6 +6,8 @@ package aws.smithy.kotlin.runtime.telemetry.trace
 
 import aws.smithy.kotlin.runtime.collections.AttributeKey
 import aws.smithy.kotlin.runtime.collections.Attributes
+import kotlin.coroutines.CoroutineContext
+import kotlin.coroutines.EmptyCoroutineContext
 
 /**
  * An abstract implementation of a trace span. By default, this class uses no-op implementations for all members unless
@@ -18,4 +20,5 @@ public abstract class AbstractTraceSpan : TraceSpan {
     override operator fun <T : Any> set(key: AttributeKey<T>, value: T) { }
     override fun mergeAttributes(attributes: Attributes) { }
     override fun close() { }
+    override fun asContextElement(): CoroutineContext = EmptyCoroutineContext
 }

runtime/observability/telemetry-api/common/src/aws/smithy/kotlin/runtime/telemetry/trace/CoroutineContextTraceExt.kt

@@ -71,7 +71,7 @@ public suspend inline fun <R> withSpan(
     // or else traces may be disconnected from their parent
     val updatedCtx = coroutineContext[TelemetryProviderContext]?.provider?.contextManager?.current()
     val telemetryCtxElement = (updatedCtx?.let { TelemetryContextElement(it) } ?: coroutineContext[TelemetryContextElement]) ?: EmptyCoroutineContext
-    withContext(context + TraceSpanContext(span) + telemetryCtxElement) {
+    withContext(context + TraceSpanContext(span) + telemetryCtxElement + span.asContextElement()) {
         block(span)
     }
 } catch (ex: Exception) {

runtime/observability/telemetry-api/common/src/aws/smithy/kotlin/runtime/telemetry/trace/TraceSpan.kt

@@ -9,6 +9,7 @@ import aws.smithy.kotlin.runtime.collections.AttributeKey
 import aws.smithy.kotlin.runtime.collections.Attributes
 import aws.smithy.kotlin.runtime.collections.emptyAttributes
 import aws.smithy.kotlin.runtime.telemetry.context.Scope
+import kotlin.coroutines.CoroutineContext
 
 /**
  * Represents a single operation/task within a trace. Each trace contains a root span and
@@ -27,6 +28,8 @@ public interface TraceSpan : Scope {
      */
     public val spanContext: SpanContext
 
+    public fun asContextElement(): CoroutineContext
+
     /**
      * Set an attribute on the span
      * @param key the attribute key to use

runtime/observability/telemetry-provider-otel/build.gradle.kts

@@ -18,6 +18,7 @@ kotlin {
         jvmMain {
             dependencies {
                 api(libs.opentelemetry.api)
+                api(libs.opentelemetry.kotlin.extension)
             }
         }
         all {

runtime/observability/telemetry-provider-otel/jvm/src/aws/smithy/kotlin/runtime/telemetry/otel/OtelTracerProvider.kt

@@ -10,6 +10,8 @@ import aws.smithy.kotlin.runtime.collections.Attributes
 import aws.smithy.kotlin.runtime.telemetry.context.Context
 import aws.smithy.kotlin.runtime.telemetry.trace.*
 import io.opentelemetry.api.OpenTelemetry
+import io.opentelemetry.extension.kotlin.asContextElement
+import kotlin.coroutines.CoroutineContext
 import io.opentelemetry.api.trace.Span as OtelSpan
 import io.opentelemetry.api.trace.SpanContext as OtelSpanContext
 import io.opentelemetry.api.trace.SpanKind as OtelSpanKind
@@ -62,11 +64,11 @@ private class OtelSpanContextImpl(private val otelSpanContext: OtelSpanContext)
 internal class OtelTraceSpanImpl(
     private val otelSpan: OtelSpan,
 ) : TraceSpan {
-
-    private val spanScope = otelSpan.makeCurrent()
-
     override val spanContext: SpanContext
         get() = OtelSpanContextImpl(otelSpan.spanContext)
+
+    override fun asContextElement(): CoroutineContext = otelSpan.asContextElement()
+
     override fun <T : Any> set(key: AttributeKey<T>, value: T) {
         key.otelAttrKeyOrNull(value)?.let { otelKey ->
             otelSpan.setAttribute(otelKey, value)
@@ -90,7 +92,6 @@ internal class OtelTraceSpanImpl(
 
     override fun close() {
         otelSpan.end()
-        spanScope.close()
     }
 }