Bug fixes & refactoring of new business metrics utils

Author: 0marperez

runtime/auth/identity-api/common/src/aws/smithy/kotlin/runtime/identity/IdentityProviderChain.kt

@@ -6,6 +6,8 @@
 package aws.smithy.kotlin.runtime.identity
 
 import aws.smithy.kotlin.runtime.InternalApi
+import aws.smithy.kotlin.runtime.businessmetrics.BusinessMetrics
+import aws.smithy.kotlin.runtime.businessmetrics.copyBusinessMetrics
 import aws.smithy.kotlin.runtime.collections.*
 import aws.smithy.kotlin.runtime.io.Closeable
 import aws.smithy.kotlin.runtime.telemetry.logging.logger
@@ -38,15 +40,14 @@ public abstract class IdentityProviderChain<P : IdentityProvider, I : Identity>(
         for (provider in providers) {
             logger.trace { "Attempting to resolve identity from $provider" }
 
-            val attributesAreMutable = attributes is MutableAttributes
-            val attributesSnapshot = if (attributesAreMutable) attributes.copy() else null
+            val businessMetricsSnapshot = attributes.copyBusinessMetrics()
 
             try {
                 @Suppress("UNCHECKED_CAST")
                 return@withSpan provider.resolve(attributes) as I
             } catch (ex: Exception) {
                 logger.debug { "unable to resolve identity from $provider: ${ex.message}" }
-                if (attributesAreMutable) (attributes as MutableAttributes).resetTo(attributesSnapshot!!)
+                if (attributes is MutableAttributes) attributes[BusinessMetrics] = businessMetricsSnapshot
                 chainException.value.addSuppressed(ex)
             }
         }

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 copyBusinessMetrics (Laws/smithy/kotlin/runtime/collections/Attributes;)Ljava/util/Set;
 	public static final fun emitBusinessMetric (Laws/smithy/kotlin/runtime/collections/Attributes;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;
@@ -134,7 +135,6 @@ public final class aws/smithy/kotlin/runtime/collections/AttributesBuilder {
 
 public final class aws/smithy/kotlin/runtime/collections/AttributesKt {
 	public static final fun attributesOf (Lkotlin/jvm/functions/Function1;)Laws/smithy/kotlin/runtime/collections/Attributes;
-	public static final fun copy (Laws/smithy/kotlin/runtime/collections/Attributes;)Laws/smithy/kotlin/runtime/collections/Attributes;
 	public static final fun emptyAttributes ()Laws/smithy/kotlin/runtime/collections/Attributes;
 	public static final fun get (Laws/smithy/kotlin/runtime/collections/Attributes;Laws/smithy/kotlin/runtime/collections/AttributeKey;)Ljava/lang/Object;
 	public static final fun isNotEmpty (Laws/smithy/kotlin/runtime/collections/Attributes;)Z
@@ -143,7 +143,6 @@ public final class aws/smithy/kotlin/runtime/collections/AttributesKt {
 	public static final fun mutableAttributesOf (Lkotlin/jvm/functions/Function1;)Laws/smithy/kotlin/runtime/collections/MutableAttributes;
 	public static final fun putIfAbsent (Laws/smithy/kotlin/runtime/collections/MutableAttributes;Laws/smithy/kotlin/runtime/collections/AttributeKey;Ljava/lang/Object;)V
 	public static final fun putIfAbsentNotNull (Laws/smithy/kotlin/runtime/collections/MutableAttributes;Laws/smithy/kotlin/runtime/collections/AttributeKey;Ljava/lang/Object;)V
-	public static final fun resetTo (Laws/smithy/kotlin/runtime/collections/MutableAttributes;Laws/smithy/kotlin/runtime/collections/Attributes;)V
 	public static final fun setIfValueNotNull (Laws/smithy/kotlin/runtime/collections/MutableAttributes;Laws/smithy/kotlin/runtime/collections/AttributeKey;Ljava/lang/Object;)V
 	public static final fun take (Laws/smithy/kotlin/runtime/collections/MutableAttributes;Laws/smithy/kotlin/runtime/collections/AttributeKey;)Ljava/lang/Object;
 	public static final fun takeOrNull (Laws/smithy/kotlin/runtime/collections/MutableAttributes;Laws/smithy/kotlin/runtime/collections/AttributeKey;)Ljava/lang/Object;

runtime/runtime-core/common/src/aws/smithy/kotlin/runtime/businessmetrics/BusinessMetricsUtils.kt

@@ -84,3 +84,17 @@ public enum class SmithyBusinessMetric(public override val identifier: String) :
 public fun Attributes.emitBusinessMetric(metric: BusinessMetric) {
     if (this is ExecutionContext) this.emitBusinessMetric(metric)
 }
+
+/**
+ * Creates a copy of an [Attributes] business metrics.
+ * Returns an empty set if no business metrics are available.
+ */
+public fun Attributes.copyBusinessMetrics(): MutableSet<String> {
+    val copy = mutableSetOf<String>()
+    if (!this.contains(BusinessMetrics)) return copy
+
+    this[BusinessMetrics].forEach { metric ->
+        copy.add(metric)
+    }
+    return copy
+}

runtime/runtime-core/common/src/aws/smithy/kotlin/runtime/collections/Attributes.kt

@@ -235,27 +235,3 @@ public inline fun attributesOf(block: AttributesBuilder.() -> Unit): Attributes
  * Returns a new [MutableAttributes] instance with elements from this set of attributes.
  */
 public fun Attributes.toMutableAttributes(): MutableAttributes = AttributesImpl(this)
-
-/**
- * Creates a copy of this [Attributes] instance.
- *
- * This function generates a new [Attributes] instance containing all the key-value pairs
- * from the current set of attributes. The new instance is independent of the original,
- * meaning changes to the copied attributes will not affect the original set.
- */
-public fun Attributes.copy(): Attributes {
-    val copy = mutableAttributes()
-    (this.keys as Set<AttributeKey<Any>>).forEach { key ->
-        copy[key] = this[key]
-    }
-    return copy
-}
-
-/**
- * Resets this [MutableAttributes] instance to the state of the specified [Attributes].
- */
-public fun MutableAttributes.resetTo(attributes: Attributes) {
-    val obsoleteKeys = this.keys.subtract(attributes.keys)
-    obsoleteKeys.forEach { key -> this.remove(key as AttributeKey<Any>) }
-    this.merge(attributes)
-}