wip: tearing apart Telemetry APIs

ianbotsf · ianbotsf · commit e2d70b088fc2 · 2024-08-08T15:03:07.000Z
diff --git a/runtime/observability/telemetry-api/api/telemetry-api.api b/runtime/observability/telemetry-api/api/telemetry-api.api
@@ -53,7 +53,7 @@ public abstract class aws/smithy/kotlin/runtime/telemetry/context/AbstractContex
 
 public abstract class aws/smithy/kotlin/runtime/telemetry/context/AbstractContextManager : aws/smithy/kotlin/runtime/telemetry/context/ContextManager {
 	public fun <init> ()V
-	public fun current ()Laws/smithy/kotlin/runtime/telemetry/context/Context;
+	public fun current (Lkotlin/coroutines/CoroutineContext;)Laws/smithy/kotlin/runtime/telemetry/context/Context;
 }
 
 public abstract class aws/smithy/kotlin/runtime/telemetry/context/AbstractScope : aws/smithy/kotlin/runtime/telemetry/context/Scope {
@@ -72,7 +72,7 @@ public final class aws/smithy/kotlin/runtime/telemetry/context/Context$Companion
 
 public abstract interface class aws/smithy/kotlin/runtime/telemetry/context/ContextManager {
 	public static final field Companion Laws/smithy/kotlin/runtime/telemetry/context/ContextManager$Companion;
-	public abstract fun current ()Laws/smithy/kotlin/runtime/telemetry/context/Context;
+	public abstract fun current (Lkotlin/coroutines/CoroutineContext;)Laws/smithy/kotlin/runtime/telemetry/context/Context;
 }
 
 public final class aws/smithy/kotlin/runtime/telemetry/context/ContextManager$Companion {
diff --git a/runtime/observability/telemetry-api/common/src/aws/smithy/kotlin/runtime/telemetry/context/AbstractContextManager.kt b/runtime/observability/telemetry-api/common/src/aws/smithy/kotlin/runtime/telemetry/context/AbstractContextManager.kt
@@ -4,10 +4,12 @@
  */
 package aws.smithy.kotlin.runtime.telemetry.context
 
+import kotlin.coroutines.CoroutineContext
+
 /**
  * An abstract implementation of a context manager. By default, this class uses no-op implementations for all members
  * unless overridden in a subclass.
  */
 public abstract class AbstractContextManager : ContextManager {
-    override fun current(): Context = Context.None
+    override fun current(ctx: CoroutineContext): Context = Context.None
 }
diff --git a/runtime/observability/telemetry-api/common/src/aws/smithy/kotlin/runtime/telemetry/context/ContextManager.kt b/runtime/observability/telemetry-api/common/src/aws/smithy/kotlin/runtime/telemetry/context/ContextManager.kt
@@ -5,6 +5,8 @@
 
 package aws.smithy.kotlin.runtime.telemetry.context
 
+import kotlin.coroutines.CoroutineContext
+
 /**
  * Responsible for managing the current context with callers current execution unit. For example, some implementations
  * use Thread Local storage for managing the current context.
@@ -20,5 +22,5 @@ public interface ContextManager {
     /**
      * Return the current [Context]
      */
-    public fun current(): Context
+    public fun current(ctx: CoroutineContext): Context
 }
diff --git a/runtime/observability/telemetry-api/common/src/aws/smithy/kotlin/runtime/telemetry/trace/CoroutineContextTraceExt.kt b/runtime/observability/telemetry-api/common/src/aws/smithy/kotlin/runtime/telemetry/trace/CoroutineContextTraceExt.kt
@@ -9,6 +9,7 @@ import aws.smithy.kotlin.runtime.ExperimentalApi
 import aws.smithy.kotlin.runtime.InternalApi
 import aws.smithy.kotlin.runtime.collections.Attributes
 import aws.smithy.kotlin.runtime.collections.emptyAttributes
+import aws.smithy.kotlin.runtime.operation.ExecutionContext
 import aws.smithy.kotlin.runtime.telemetry.TelemetryProviderContext
 import aws.smithy.kotlin.runtime.telemetry.context.TelemetryContextElement
 import aws.smithy.kotlin.runtime.telemetry.context.telemetryContext
@@ -69,7 +70,7 @@ public suspend inline fun <R> withSpan(
 ): R = try {
     // after a span is created the current telemetry context may be updated, ensure the new coroutine gets the latest
     // or else traces may be disconnected from their parent
-    val updatedCtx = coroutineContext[TelemetryProviderContext]?.provider?.contextManager?.current()
+    val updatedCtx = coroutineContext[TelemetryProviderContext]?.provider?.contextManager?.current(coroutineContext)
     val telemetryCtxElement = (updatedCtx?.let { TelemetryContextElement(it) } ?: coroutineContext[TelemetryContextElement]) ?: EmptyCoroutineContext
     withContext(context + TraceSpanContext(span) + telemetryCtxElement) {
         block(span)
diff --git a/runtime/observability/telemetry-provider-ism/api/telemetry-provider-ism.api b/runtime/observability/telemetry-provider-ism/api/telemetry-provider-ism.api
@@ -0,0 +1,50 @@
+public final class aws/smithy/kotlin/runtime/telemetry/ism/IsmContextManager : aws/smithy/kotlin/runtime/telemetry/context/ContextManager {
+	public static final field Companion Laws/smithy/kotlin/runtime/telemetry/ism/IsmContextManager$Companion;
+	public fun current (Lkotlin/coroutines/CoroutineContext;)Laws/smithy/kotlin/runtime/telemetry/context/Context;
+}
+
+public final class aws/smithy/kotlin/runtime/telemetry/ism/IsmContextManager$Companion {
+	public final fun createWithScopeListener ()Lkotlin/Pair;
+}
+
+public final class aws/smithy/kotlin/runtime/telemetry/ism/IsmMetricsProvider : aws/smithy/kotlin/runtime/telemetry/metrics/MeterProvider {
+	public fun <init> (Laws/smithy/kotlin/runtime/telemetry/ism/OperationMetricsCollector;)V
+	public fun getOrCreateMeter (Ljava/lang/String;)Laws/smithy/kotlin/runtime/telemetry/metrics/Meter;
+}
+
+public final class aws/smithy/kotlin/runtime/telemetry/ism/IsmMetricsProviderKt {
+	public static final fun MetricRecord (Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Object;Laws/smithy/kotlin/runtime/collections/Attributes;Laws/smithy/kotlin/runtime/telemetry/context/Context;Laws/smithy/kotlin/runtime/time/Instant;)Laws/smithy/kotlin/runtime/telemetry/ism/MetricRecord;
+	public static final fun OperationMetrics (Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/util/Map;Ljava/util/Map;)Laws/smithy/kotlin/runtime/telemetry/ism/OperationMetrics;
+	public static final fun ScopeMetrics (Ljava/util/Map;Ljava/util/Map;)Laws/smithy/kotlin/runtime/telemetry/ism/ScopeMetrics;
+}
+
+public abstract interface class aws/smithy/kotlin/runtime/telemetry/ism/MetricRecord {
+	public abstract fun getAttributes ()Laws/smithy/kotlin/runtime/collections/Attributes;
+	public abstract fun getContext ()Laws/smithy/kotlin/runtime/telemetry/context/Context;
+	public abstract fun getDescription ()Ljava/lang/String;
+	public abstract fun getName ()Ljava/lang/String;
+	public abstract fun getTimestamp ()Laws/smithy/kotlin/runtime/time/Instant;
+	public abstract fun getUnits ()Ljava/lang/String;
+	public abstract fun getValue ()Ljava/lang/Object;
+}
+
+public abstract interface class aws/smithy/kotlin/runtime/telemetry/ism/OperationMetrics : aws/smithy/kotlin/runtime/telemetry/ism/ScopeMetrics {
+	public abstract fun getOperation ()Ljava/lang/String;
+	public abstract fun getSdkInvocationId ()Ljava/lang/String;
+	public abstract fun getService ()Ljava/lang/String;
+}
+
+public abstract interface class aws/smithy/kotlin/runtime/telemetry/ism/OperationMetricsCollector {
+	public abstract fun onOperationMetrics (Laws/smithy/kotlin/runtime/telemetry/ism/OperationMetrics;)V
+}
+
+public abstract interface class aws/smithy/kotlin/runtime/telemetry/ism/ScopeMetrics {
+	public abstract fun getChildScopes ()Ljava/util/Map;
+	public abstract fun getRecords ()Ljava/util/Map;
+}
+
+public abstract interface class aws/smithy/kotlin/runtime/telemetry/ism/SpanListener {
+	public abstract fun onCloseSpan (Laws/smithy/kotlin/runtime/telemetry/context/Context;)V
+	public abstract fun onNewSpan (Laws/smithy/kotlin/runtime/telemetry/context/Context;Ljava/lang/String;Laws/smithy/kotlin/runtime/collections/Attributes;)Laws/smithy/kotlin/runtime/telemetry/context/Context;
+}
+
diff --git a/runtime/observability/telemetry-provider-ism/common/src/aws/smithy/kotlin/runtime/telemetry/ism/IsmContextManager.kt b/runtime/observability/telemetry-provider-ism/common/src/aws/smithy/kotlin/runtime/telemetry/ism/IsmContextManager.kt
@@ -2,29 +2,21 @@
  * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
-package aws.smithy.kotlin.runtime.telemetry.ism.context
+package aws.smithy.kotlin.runtime.telemetry.ism
 
 import aws.smithy.kotlin.runtime.collections.Attributes
 import aws.smithy.kotlin.runtime.http.operation.OperationAttributes
 import aws.smithy.kotlin.runtime.telemetry.context.Context
 import aws.smithy.kotlin.runtime.telemetry.context.ContextManager
 import aws.smithy.kotlin.runtime.telemetry.context.Scope
-import aws.smithy.kotlin.runtime.telemetry.ism.metrics.MetricRecord
-import aws.smithy.kotlin.runtime.telemetry.ism.metrics.ScopeMetrics
+import aws.smithy.kotlin.runtime.telemetry.context.telemetryContext
+import kotlin.coroutines.CoroutineContext
 
 public interface SpanListener {
     public fun onNewSpan(parentContext: Context?, name: String, attributes: Attributes): Context
     public fun onCloseSpan(context: Context)
 }
 
-internal interface ContextStorage {
-    fun get(): Context
-    fun getAndSet(value: Context): Context
-    fun requireAndSet(expect: Context, update: Context)
-}
-
-internal expect fun ContextStorage(initialContext: Context): ContextStorage
-
 public class IsmContextManager private constructor() : ContextManager {
     public companion object {
         public fun createWithScopeListener(): Pair<IsmContextManager, SpanListener> {
@@ -40,9 +32,8 @@ public class IsmContextManager private constructor() : ContextManager {
     }
 
     private val rootContext = object : HierarchicalContext(null) { }
-    private val storage = ContextStorage(rootContext)
 
-    override fun current(): Context = storage.get()
+    override fun current(ctx: CoroutineContext): Context = ctx.telemetryContext ?: rootContext
 
     private fun onNewSpan(parentContext: Context?, name: String, attributes: Attributes): Context =
         when (parentContext) {
@@ -62,13 +53,10 @@ public class IsmContextManager private constructor() : ContextManager {
             else -> OtherContext(parentContext)
         }
 
-    private fun onCloseSpan(context: Context) {
-        TODO()
-    }
+    private fun onCloseSpan(context: Context) = Unit
 
     private abstract inner class HierarchicalContext(val parent: Context?) : Context {
         override fun makeCurrent(): Scope {
-            storage.requireAndSet(parent ?: Context.None, this)
             return IsmScope(this)
         }
     }
@@ -90,8 +78,6 @@ public class IsmContextManager private constructor() : ContextManager {
     private inner class OtherContext(parent: Context?) : HierarchicalContext(parent)
 
     private inner class IsmScope(val context: Context) : Scope {
-        override fun close() {
-            storage.requireAndSet(context, Context.None)
-        }
+        override fun close() = Unit
     }
 }
diff --git a/runtime/observability/telemetry-provider-ism/common/src/aws/smithy/kotlin/runtime/telemetry/ism/IsmMetricsProvider.kt b/runtime/observability/telemetry-provider-ism/common/src/aws/smithy/kotlin/runtime/telemetry/ism/IsmMetricsProvider.kt
@@ -2,7 +2,7 @@
  * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
-package aws.smithy.kotlin.runtime.telemetry.ism.metrics
+package aws.smithy.kotlin.runtime.telemetry.ism
 
 import aws.smithy.kotlin.runtime.ExperimentalApi
 import aws.smithy.kotlin.runtime.collections.Attributes
@@ -57,8 +57,7 @@ private data class ScopeMetricsImpl(
 public fun ScopeMetrics(
     records: Map<String, List<MetricRecord<*>>>,
     childScopes: Map<String, ScopeMetrics>,
-): ScopeMetrics =
-    ScopeMetricsImpl(records, childScopes)
+): ScopeMetrics = ScopeMetricsImpl(records, childScopes)
 
 public interface OperationMetrics : ScopeMetrics {
     public val service: String
diff --git a/runtime/observability/telemetry-provider-ism/common/src/aws/smithy/kotlin/runtime/telemetry/ism/IsmTracerProvider.kt b/runtime/observability/telemetry-provider-ism/common/src/aws/smithy/kotlin/runtime/telemetry/ism/IsmTracerProvider.kt
@@ -0,0 +1,52 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package aws.smithy.kotlin.runtime.telemetry.ism
+
+import aws.smithy.kotlin.runtime.collections.AttributeKey
+import aws.smithy.kotlin.runtime.collections.Attributes
+import aws.smithy.kotlin.runtime.telemetry.context.Context
+import aws.smithy.kotlin.runtime.telemetry.trace.*
+
+public class IsmTracerProvider : TracerProvider {
+    override fun getOrCreateTracer(scope: String): Tracer {
+        TODO("Not yet implemented")
+    }
+}
+
+public class IsmTracer : Tracer {
+    override fun createSpan(
+        name: String,
+        initialAttributes: Attributes,
+        spanKind: SpanKind,
+        parentContext: Context?
+    ): TraceSpan {
+        TODO("Not yet implemented")
+    }
+}
+
+public class IsmTraceSpan : TraceSpan {
+    override val spanContext: SpanContext
+        get() = TODO("Not yet implemented")
+
+    override fun <T : Any> set(key: AttributeKey<T>, value: T) {
+        TODO("Not yet implemented")
+    }
+
+    override fun mergeAttributes(attributes: Attributes) {
+        TODO("Not yet implemented")
+    }
+
+    override fun emitEvent(name: String, attributes: Attributes) {
+        TODO("Not yet implemented")
+    }
+
+    override fun setStatus(status: SpanStatus) {
+        TODO("Not yet implemented")
+    }
+
+    override fun close() {
+        TODO("Not yet implemented")
+    }
+}
diff --git a/runtime/observability/telemetry-provider-ism/common/src/aws/smithy/kotlin/runtime/telemetry/ism/trace/IsmTracerProvider.kt b/runtime/observability/telemetry-provider-ism/common/src/aws/smithy/kotlin/runtime/telemetry/ism/trace/IsmTracerProvider.kt
diff --git a/runtime/observability/telemetry-provider-ism/jvm/src/aws/smithy/kotlin/runtime/telemetry/ism/context/ContextStorageJvm.kt b/runtime/observability/telemetry-provider-ism/jvm/src/aws/smithy/kotlin/runtime/telemetry/ism/context/ContextStorageJvm.kt
diff --git a/runtime/observability/telemetry-provider-otel/jvm/src/aws/smithy/kotlin/runtime/telemetry/otel/OtelContextManager.kt b/runtime/observability/telemetry-provider-otel/jvm/src/aws/smithy/kotlin/runtime/telemetry/otel/OtelContextManager.kt
@@ -8,11 +8,12 @@ package aws.smithy.kotlin.runtime.telemetry.otel
 import aws.smithy.kotlin.runtime.telemetry.context.Context
 import aws.smithy.kotlin.runtime.telemetry.context.ContextManager
 import aws.smithy.kotlin.runtime.telemetry.context.Scope
+import kotlin.coroutines.CoroutineContext
 import io.opentelemetry.context.Context as OpenTelemetryContext
 import io.opentelemetry.context.Scope as OtelScope
 
 internal object OtelContextManager : ContextManager {
-    override fun current(): Context =
+    override fun current(ctx: CoroutineContext): Context =
         OtelContext(OpenTelemetryContext.current())
 }
 
diff --git a/runtime/protocol/http-client/api/http-client.api b/runtime/protocol/http-client/api/http-client.api
@@ -542,6 +542,13 @@ public final class aws/smithy/kotlin/runtime/http/operation/MutateMiddleware$Def
 	public static fun install (Laws/smithy/kotlin/runtime/http/operation/MutateMiddleware;Laws/smithy/kotlin/runtime/http/operation/SdkHttpOperation;)V
 }
 
+public final class aws/smithy/kotlin/runtime/http/operation/OperationAttributes {
+	public static final field INSTANCE Laws/smithy/kotlin/runtime/http/operation/OperationAttributes;
+	public final fun getAwsInvocationId ()Laws/smithy/kotlin/runtime/collections/AttributeKey;
+	public final fun getRpcOperation ()Laws/smithy/kotlin/runtime/collections/AttributeKey;
+	public final fun getRpcService ()Laws/smithy/kotlin/runtime/collections/AttributeKey;
+}
+
 public final class aws/smithy/kotlin/runtime/http/operation/OperationAuthConfig {
 	public static final field Companion Laws/smithy/kotlin/runtime/http/operation/OperationAuthConfig$Companion;
 	public fun <init> (Laws/smithy/kotlin/runtime/http/operation/AuthSchemeResolver;Ljava/util/Map;Laws/smithy/kotlin/runtime/identity/IdentityProviderConfig;)V
diff --git a/runtime/protocol/http-client/common/src/aws/smithy/kotlin/runtime/http/interceptors/OperationTelemetryInterceptor.kt b/runtime/protocol/http-client/common/src/aws/smithy/kotlin/runtime/http/interceptors/OperationTelemetryInterceptor.kt
@@ -18,6 +18,7 @@ import aws.smithy.kotlin.runtime.http.operation.OperationAttributes
 import aws.smithy.kotlin.runtime.http.operation.OperationMetrics
 import aws.smithy.kotlin.runtime.http.request.HttpRequest
 import aws.smithy.kotlin.runtime.http.response.HttpResponse
+import aws.smithy.kotlin.runtime.telemetry.context.Context
 import aws.smithy.kotlin.runtime.telemetry.metrics.recordSeconds
 import kotlin.time.TimeMark
 import kotlin.time.TimeSource
@@ -49,15 +50,18 @@ internal class OperationTelemetryInterceptor(
         OperationAttributes.RpcOperation to operation
     }
 
+    private val RequestInterceptorContext<*>.metricsContext: Context
+        get() = metrics.provider.contextManager.current(executionContext.coroutineContext)
+
     override fun readBeforeExecution(context: RequestInterceptorContext<Any>) {
         callStart = timeSource.markNow()
     }
 
     override fun readAfterExecution(context: ResponseInterceptorContext<Any, Any, HttpRequest?, HttpResponse?>) {
-        val currentCtx = metrics.provider.contextManager.current()
+        val metricsCtx = context.metricsContext
 
         callStart?.elapsedNow()?.let { callDuration ->
-            metrics.rpcCallDuration.recordSeconds(callDuration, perRpcAttributes, currentCtx)
+            metrics.rpcCallDuration.recordSeconds(callDuration, perRpcAttributes, metricsCtx)
         }
 
         context.response.exceptionOrNull()?.let { ex ->
@@ -69,7 +73,7 @@ internal class OperationTelemetryInterceptor(
             } else {
                 perRpcAttributes
             }
-            metrics.rpcErrors.add(1L, errorAttributes, currentCtx)
+            metrics.rpcErrors.add(1L, errorAttributes, metricsCtx)
         }
     }
 
@@ -79,15 +83,17 @@ internal class OperationTelemetryInterceptor(
 
     override fun readAfterSerialization(context: ProtocolRequestInterceptorContext<Any, HttpRequest>) {
         val serializeDuration = serializeStart?.elapsedNow() ?: return
-        metrics.serializationDuration.recordSeconds(serializeDuration, perRpcAttributes, metrics.provider.contextManager.current())
+        metrics.serializationDuration.recordSeconds(serializeDuration, perRpcAttributes, context.metricsContext)
     }
 
     override fun readAfterAttempt(context: ResponseInterceptorContext<Any, Any, HttpRequest, HttpResponse?>) {
-        metrics.rpcAttempts.add(1L, perRpcAttributes, metrics.provider.contextManager.current())
+        val metricsCtx = context.metricsContext
+
+        metrics.rpcAttempts.add(1L, perRpcAttributes, metricsCtx)
         attempts++
 
         val attemptDuration = attemptStart?.elapsedNow() ?: return
-        metrics.rpcAttemptDuration.recordSeconds(attemptDuration, perRpcAttributes, metrics.provider.contextManager.current())
+        metrics.rpcAttemptDuration.recordSeconds(attemptDuration, perRpcAttributes, metricsCtx)
 
         context.executionContext.takeOrNull(EngineAttributes.TimeToFirstByte)?.let { ttfb ->
             metrics.rpcAttemptOverheadDuration.recordSeconds(attemptDuration - ttfb, perRpcAttributes)
@@ -100,7 +106,7 @@ internal class OperationTelemetryInterceptor(
 
     override fun readAfterDeserialization(context: ResponseInterceptorContext<Any, Any, HttpRequest, HttpResponse>) {
         val deserializeDuration = deserializeStart?.elapsedNow() ?: return
-        metrics.deserializationDuration.recordSeconds(deserializeDuration, perRpcAttributes, metrics.provider.contextManager.current())
+        metrics.deserializationDuration.recordSeconds(deserializeDuration, perRpcAttributes, context.metricsContext)
     }
 
     override fun readBeforeAttempt(context: ProtocolRequestInterceptorContext<Any, HttpRequest>) {
diff --git a/runtime/protocol/http-client/common/src/aws/smithy/kotlin/runtime/http/operation/OperationTelemetry.kt b/runtime/protocol/http-client/common/src/aws/smithy/kotlin/runtime/http/operation/OperationTelemetry.kt

Original file line number	Diff line number	Diff line change
`@@ -53,7 +53,7 @@ public abstract class aws/smithy/kotlin/runtime/telemetry/context/AbstractContex`
`53`	`53`
`54`	`54`	`public abstract class aws/smithy/kotlin/runtime/telemetry/context/AbstractContextManager : aws/smithy/kotlin/runtime/telemetry/context/ContextManager {`
`55`	`55`	`public fun <init> ()V`
`56`		`- public fun current ()Laws/smithy/kotlin/runtime/telemetry/context/Context;`
	`56`	`+ public fun current (Lkotlin/coroutines/CoroutineContext;)Laws/smithy/kotlin/runtime/telemetry/context/Context;`
`57`	`57`	`}`
`58`	`58`
`59`	`59`	`public abstract class aws/smithy/kotlin/runtime/telemetry/context/AbstractScope : aws/smithy/kotlin/runtime/telemetry/context/Scope {`
`@@ -72,7 +72,7 @@ public final class aws/smithy/kotlin/runtime/telemetry/context/Context$Companion`
`72`	`72`
`73`	`73`	`public abstract interface class aws/smithy/kotlin/runtime/telemetry/context/ContextManager {`
`74`	`74`	`public static final field Companion Laws/smithy/kotlin/runtime/telemetry/context/ContextManager$Companion;`
`75`		`- public abstract fun current ()Laws/smithy/kotlin/runtime/telemetry/context/Context;`
	`75`	`+ public abstract fun current (Lkotlin/coroutines/CoroutineContext;)Laws/smithy/kotlin/runtime/telemetry/context/Context;`
`76`	`76`	`}`
`77`	`77`
`78`	`78`	`public final class aws/smithy/kotlin/runtime/telemetry/context/ContextManager$Companion {`