misc: credentials business metrics (#1164)

smithy-lang · Nov 4, 2024 · 1614926 · 1614926
1 parent e416285
commit 1614926
Show file tree

Hide file tree

Showing 4 changed files with 36 additions and 10 deletions.
diff --git a/.../http-client/common/src/aws/smithy/kotlin/runtime/http/operation/SdkOperationExecution.kt b/.../http-client/common/src/aws/smithy/kotlin/runtime/http/operation/SdkOperationExecution.kt
@@ -6,6 +6,7 @@
 package aws.smithy.kotlin.runtime.http.operation
 
 import aws.smithy.kotlin.runtime.InternalApi
+import aws.smithy.kotlin.runtime.businessmetrics.BusinessMetrics
 import aws.smithy.kotlin.runtime.businessmetrics.SmithyBusinessMetric
 import aws.smithy.kotlin.runtime.businessmetrics.emitBusinessMetric
 import aws.smithy.kotlin.runtime.client.LogMode
@@ -23,9 +24,11 @@ import aws.smithy.kotlin.runtime.http.request.dumpRequest
 import aws.smithy.kotlin.runtime.http.request.immutableView
 import aws.smithy.kotlin.runtime.http.request.toBuilder
 import aws.smithy.kotlin.runtime.http.response.dumpResponse
+import aws.smithy.kotlin.runtime.identity.Identity
 import aws.smithy.kotlin.runtime.io.Handler
 import aws.smithy.kotlin.runtime.io.middleware.Middleware
 import aws.smithy.kotlin.runtime.io.middleware.Phase
+import aws.smithy.kotlin.runtime.operation.ExecutionContext
 import aws.smithy.kotlin.runtime.retries.RetryStrategy
 import aws.smithy.kotlin.runtime.retries.StandardRetryStrategy
 import aws.smithy.kotlin.runtime.retries.policy.RetryPolicy
@@ -285,6 +288,9 @@ internal class AuthHandler<Input, Output>(
             identityProvider.resolve(request.context)
         }
 
+        // emit identity business metrics
+        emitIdentityBusinessMetrics(identity, request.context)
+
         val resolveEndpointReq = ResolveEndpointRequest(request.context, request.subject.immutableView(), identity)
 
         if (endpointResolver != null) {
@@ -408,3 +414,9 @@ private class InterceptorTransmitMiddleware<I, O>(
         return call
     }
 }
+
+/**
+ * Emits an [Identity]'s attributes' [BusinessMetrics] into the [ExecutionContext]
+ */
+private fun emitIdentityBusinessMetrics(identity: Identity, context: ExecutionContext) =
+    identity.attributes.getOrNull(BusinessMetrics)?.toList()?.reversed()?.forEach(context::emitBusinessMetric)
diff --git a/runtime/runtime-core/api/runtime-core.api b/runtime/runtime-core/api/runtime-core.api
@@ -83,6 +83,7 @@ public abstract interface class aws/smithy/kotlin/runtime/businessmetrics/Busine
 
 public final class aws/smithy/kotlin/runtime/businessmetrics/BusinessMetricsUtilsKt {
 	public static final fun containsBusinessMetric (Laws/smithy/kotlin/runtime/operation/ExecutionContext;Laws/smithy/kotlin/runtime/businessmetrics/BusinessMetric;)Z
+	public static final fun emitBusinessMetric (Laws/smithy/kotlin/runtime/collections/MutableAttributes;Laws/smithy/kotlin/runtime/businessmetrics/BusinessMetric;)V
 	public static final fun emitBusinessMetric (Laws/smithy/kotlin/runtime/operation/ExecutionContext;Laws/smithy/kotlin/runtime/businessmetrics/BusinessMetric;)V
 	public static final fun getAccountIdBasedEndpointAccountId ()Laws/smithy/kotlin/runtime/collections/AttributeKey;
 	public static final fun getBusinessMetrics ()Laws/smithy/kotlin/runtime/collections/AttributeKey;

diff --git a/...runtime-core/common/src/aws/smithy/kotlin/runtime/businessmetrics/BusinessMetricsUtils.kt b/...runtime-core/common/src/aws/smithy/kotlin/runtime/businessmetrics/BusinessMetricsUtils.kt
@@ -6,14 +6,15 @@ package aws.smithy.kotlin.runtime.businessmetrics
 
 import aws.smithy.kotlin.runtime.InternalApi
 import aws.smithy.kotlin.runtime.collections.AttributeKey
+import aws.smithy.kotlin.runtime.collections.MutableAttributes
 import aws.smithy.kotlin.runtime.collections.get
 import aws.smithy.kotlin.runtime.operation.ExecutionContext
 
 /**
  * Keeps track of all business metrics along an operations execution
  */
 @InternalApi
-public val BusinessMetrics: AttributeKey<MutableSet<String>> = AttributeKey("aws.sdk.kotlin#BusinessMetrics")
+public val BusinessMetrics: AttributeKey<MutableSet<BusinessMetric>> = AttributeKey("aws.smithy.kotlin#BusinessMetrics")
 
 /**
  * The account ID in an account ID based endpoint
@@ -33,9 +34,21 @@ public val ServiceEndpointOverride: AttributeKey<Boolean> = AttributeKey("aws.sm
 @InternalApi
 public fun ExecutionContext.emitBusinessMetric(metric: BusinessMetric) {
     if (this.attributes.contains(BusinessMetrics)) {
-        this.attributes[BusinessMetrics].add(metric.identifier)
+        this.attributes[BusinessMetrics].add(metric)
     } else {
-        this.attributes[BusinessMetrics] = mutableSetOf(metric.identifier)
+        this.attributes[BusinessMetrics] = mutableSetOf(metric)
+    }
+}
+
+/**
+ * Emit a business metric to the mutable attributes
+ */
+@InternalApi
+public fun MutableAttributes.emitBusinessMetric(metric: BusinessMetric) {
+    if (this.contains(BusinessMetrics)) {
+        this[BusinessMetrics].add(metric)
+    } else {
+        this[BusinessMetrics] = mutableSetOf(metric)
     }
 }
 
@@ -45,7 +58,7 @@ public fun ExecutionContext.emitBusinessMetric(metric: BusinessMetric) {
 @InternalApi
 public fun ExecutionContext.removeBusinessMetric(metric: BusinessMetric) {
     if (this.attributes.contains(BusinessMetrics)) {
-        this.attributes[BusinessMetrics].remove(metric.identifier)
+        this.attributes[BusinessMetrics].remove(metric)
     }
 }
 
@@ -54,7 +67,7 @@ public fun ExecutionContext.removeBusinessMetric(metric: BusinessMetric) {
  */
 @InternalApi
 public fun ExecutionContext.containsBusinessMetric(metric: BusinessMetric): Boolean =
-    (this.attributes.contains(BusinessMetrics)) && this.attributes[BusinessMetrics].contains(metric.identifier)
+    (this.attributes.contains(BusinessMetrics)) && this.attributes[BusinessMetrics].contains(metric)
 
 /**
  * Valid business metrics

diff --git a/...me-core/common/test/aws/smithy/kotlin/runtime/businessmetrics/BusinessMetricsUtilsTest.kt b/...me-core/common/test/aws/smithy/kotlin/runtime/businessmetrics/BusinessMetricsUtilsTest.kt
@@ -17,7 +17,7 @@ class BusinessMetricsUtilsTest {
         executionContext.emitBusinessMetric(SmithyBusinessMetric.GZIP_REQUEST_COMPRESSION)
 
         assertTrue(executionContext.attributes.contains(BusinessMetrics))
-        assertTrue(executionContext.attributes[BusinessMetrics].contains(SmithyBusinessMetric.GZIP_REQUEST_COMPRESSION.identifier))
+        assertTrue(executionContext.attributes[BusinessMetrics].contains(SmithyBusinessMetric.GZIP_REQUEST_COMPRESSION))
     }
 
     @Test
@@ -27,8 +27,8 @@ class BusinessMetricsUtilsTest {
         executionContext.emitBusinessMetric(SmithyBusinessMetric.SIGV4A_SIGNING)
 
         assertTrue(executionContext.attributes.contains(BusinessMetrics))
-        assertTrue(executionContext.attributes[BusinessMetrics].contains(SmithyBusinessMetric.GZIP_REQUEST_COMPRESSION.identifier))
-        assertTrue(executionContext.attributes[BusinessMetrics].contains(SmithyBusinessMetric.SIGV4A_SIGNING.identifier))
+        assertTrue(executionContext.attributes[BusinessMetrics].contains(SmithyBusinessMetric.GZIP_REQUEST_COMPRESSION))
+        assertTrue(executionContext.attributes[BusinessMetrics].contains(SmithyBusinessMetric.SIGV4A_SIGNING))
     }
 
     @Test
@@ -37,12 +37,12 @@ class BusinessMetricsUtilsTest {
         executionContext.emitBusinessMetric(SmithyBusinessMetric.GZIP_REQUEST_COMPRESSION)
 
         assertTrue(executionContext.attributes.contains(BusinessMetrics))
-        assertTrue(executionContext.attributes[BusinessMetrics].contains(SmithyBusinessMetric.GZIP_REQUEST_COMPRESSION.identifier))
+        assertTrue(executionContext.attributes[BusinessMetrics].contains(SmithyBusinessMetric.GZIP_REQUEST_COMPRESSION))
 
         executionContext.removeBusinessMetric(SmithyBusinessMetric.GZIP_REQUEST_COMPRESSION)
 
         assertTrue(executionContext.attributes.contains(BusinessMetrics))
-        assertFalse(executionContext.attributes[BusinessMetrics].contains(SmithyBusinessMetric.GZIP_REQUEST_COMPRESSION.identifier))
+        assertFalse(executionContext.attributes[BusinessMetrics].contains(SmithyBusinessMetric.GZIP_REQUEST_COMPRESSION))
     }
 
     @Test