emit DDB_MAPPER business metric (#1426)

aws-runtime/aws-http/api/aws-http.api

@@ -141,6 +141,7 @@ public final class aws/sdk/kotlin/runtime/http/interceptors/AddUserAgentMetadata
 }
 
 public final class aws/sdk/kotlin/runtime/http/interceptors/AwsBusinessMetric : java/lang/Enum, aws/smithy/kotlin/runtime/businessmetrics/BusinessMetric {
+	public static final field DDB_MAPPER Laws/sdk/kotlin/runtime/http/interceptors/AwsBusinessMetric;
 	public static final field S3_EXPRESS_BUCKET Laws/sdk/kotlin/runtime/http/interceptors/AwsBusinessMetric;
 	public static fun getEntries ()Lkotlin/enums/EnumEntries;
 	public fun getIdentifier ()Ljava/lang/String;

aws-runtime/aws-http/common/src/aws/sdk/kotlin/runtime/http/interceptors/BusinessMetricsInterceptor.kt

@@ -64,4 +64,5 @@ private fun formatMetrics(metrics: MutableSet<String>): String {
 @InternalApi
 public enum class AwsBusinessMetric(public override val identifier: String) : BusinessMetric {
     S3_EXPRESS_BUCKET("J"),
+    DDB_MAPPER("d"),
 }

hll/dynamodb-mapper/dynamodb-mapper-codegen/src/main/kotlin/aws/sdk/kotlin/hll/dynamodbmapper/codegen/model/MapperPkg.kt

@@ -12,6 +12,7 @@ public object MapperPkg {
     public object Hl {
         public val Base: String = "aws.sdk.kotlin.hll.dynamodbmapper"
         public val Annotations: String = "$Base.annotations"
+        public val Internal: String = "$Base.internal"
         public val Items: String = "$Base.items"
         public val Model: String = "$Base.model"
         public val Ops: String = "$Base.operations"

hll/dynamodb-mapper/dynamodb-mapper-codegen/src/main/kotlin/aws/sdk/kotlin/hll/dynamodbmapper/codegen/model/MapperTypes.kt

@@ -22,6 +22,10 @@ public object MapperTypes {
         public val ManualPagination: TypeRef = TypeRef(MapperPkg.Hl.Annotations, "ManualPagination")
     }
 
+    public object Internal {
+        public val withWrappedClient: TypeRef = TypeRef(MapperPkg.Hl.Internal, "withWrappedClient")
+    }
+
     public object Items {
         public fun itemSchema(typeVar: String): TypeRef =
             TypeRef(MapperPkg.Hl.Items, "ItemSchema", genericArgs = listOf(TypeVar(typeVar)))

hll/dynamodb-mapper/dynamodb-mapper-ops-codegen/src/main/kotlin/aws/sdk/kotlin/hll/dynamodbmapper/codegen/operations/rendering/OperationRenderer.kt

@@ -70,7 +70,12 @@ internal class OperationRenderer(
                 }
                 write("schema) },")
 
-                write("lowLevelInvoke = spec.mapper.client::#L,", operation.methodName)
+                withBlock("lowLevelInvoke = { lowLevelReq ->", "},") {
+                    withBlock("spec.mapper.client.#T { client ->", "}", MapperTypes.Internal.withWrappedClient) {
+                        write("client.#L(lowLevelReq)", operation.methodName)
+                    }
+                }
+
                 write("deserialize = #L::convert,", operation.response.lowLevelName)
                 write("interceptors = spec.mapper.config.interceptors,")
             }

hll/dynamodb-mapper/dynamodb-mapper/build.gradle.kts

@@ -26,6 +26,7 @@ kotlin {
     sourceSets {
         commonMain {
             dependencies {
+                implementation(project(":aws-runtime:aws-http"))
                 api(project(":services:dynamodb"))
                 api(project(":hll:hll-mapping-core"))
                 api(libs.kotlinx.coroutines.core)

hll/dynamodb-mapper/dynamodb-mapper/common/src/aws/sdk/kotlin/hll/dynamodbmapper/internal/DynamoDbMapperImpl.kt

@@ -8,7 +8,12 @@ import aws.sdk.kotlin.hll.dynamodbmapper.DynamoDbMapper
 import aws.sdk.kotlin.hll.dynamodbmapper.items.ItemSchema
 import aws.sdk.kotlin.hll.dynamodbmapper.model.internal.tableImpl
 import aws.sdk.kotlin.hll.dynamodbmapper.pipeline.InterceptorAny
+import aws.sdk.kotlin.runtime.http.interceptors.AwsBusinessMetric
 import aws.sdk.kotlin.services.dynamodb.DynamoDbClient
+import aws.sdk.kotlin.services.dynamodb.withConfig
+import aws.smithy.kotlin.runtime.businessmetrics.emitBusinessMetric
+import aws.smithy.kotlin.runtime.client.RequestInterceptorContext
+import aws.smithy.kotlin.runtime.http.interceptors.HttpInterceptor
 
 internal data class DynamoDbMapperImpl(
     override val client: DynamoDbClient,
@@ -35,3 +40,16 @@ internal class MapperConfigBuilderImpl : DynamoDbMapper.Config.Builder {
 
     override fun build() = MapperConfigImpl(interceptors.toList())
 }
+
+/**
+ * An interceptor that emits the DynamoDB Mapper business metric
+ */
+private object BusinessMetricInterceptor : HttpInterceptor {
+    override suspend fun modifyBeforeSerialization(context: RequestInterceptorContext<Any>): Any {
+        context.executionContext.emitBusinessMetric(AwsBusinessMetric.DDB_MAPPER)
+        return context.request
+    }
+}
+
+internal inline fun <T> DynamoDbClient.withWrappedClient(block: (DynamoDbClient) -> T): T =
+    withConfig { interceptors += BusinessMetricInterceptor }.use(block)

hll/dynamodb-mapper/dynamodb-mapper/common/test/aws/sdk/kotlin/hll/dynamodbmapper/DynamoDbMapperTest.kt

@@ -0,0 +1,109 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package aws.sdk.kotlin.hll.dynamodbmapper
+
+import aws.sdk.kotlin.hll.dynamodbmapper.items.AttributeDescriptor
+import aws.sdk.kotlin.hll.dynamodbmapper.items.ItemSchema
+import aws.sdk.kotlin.hll.dynamodbmapper.items.KeySpec
+import aws.sdk.kotlin.hll.dynamodbmapper.items.SimpleItemConverter
+import aws.sdk.kotlin.hll.dynamodbmapper.operations.scanPaginated
+import aws.sdk.kotlin.hll.dynamodbmapper.testutils.DdbLocalTest
+import aws.sdk.kotlin.hll.dynamodbmapper.values.scalars.IntConverter
+import aws.sdk.kotlin.hll.dynamodbmapper.values.scalars.StringConverter
+import aws.sdk.kotlin.runtime.http.interceptors.AwsBusinessMetric
+import aws.sdk.kotlin.services.dynamodb.scan
+import aws.sdk.kotlin.services.dynamodb.withConfig
+import aws.smithy.kotlin.runtime.businessmetrics.BusinessMetric
+import aws.smithy.kotlin.runtime.businessmetrics.BusinessMetrics
+import aws.smithy.kotlin.runtime.client.ProtocolRequestInterceptorContext
+import aws.smithy.kotlin.runtime.collections.get
+import aws.smithy.kotlin.runtime.http.interceptors.HttpInterceptor
+import aws.smithy.kotlin.runtime.http.request.HttpRequest
+import kotlinx.coroutines.flow.collect
+import kotlinx.coroutines.test.runTest
+import kotlin.test.assertFalse
+import kotlin.test.assertTrue
+
+class DynamoDbMapperTest : DdbLocalTest() {
+    companion object {
+        private const val TABLE_NAME = "dummy"
+
+        private data class DummyData(var foo: String = "", var bar: Int = 0)
+
+        private val dummyConverter = SimpleItemConverter(
+            ::DummyData,
+            { this },
+            AttributeDescriptor("foo", DummyData::foo, DummyData::foo::set, StringConverter),
+            AttributeDescriptor("bar", DummyData::bar, DummyData::bar::set, IntConverter),
+        )
+
+        private val dummySchema = ItemSchema(dummyConverter, KeySpec.String("foo"), KeySpec.Number("bar"))
+    }
+
+    @BeforeAll
+    fun setUp() = runTest {
+        createTable(TABLE_NAME, dummySchema)
+    }
+
+    @Test
+    fun testBusinessMetricEmission() = runTest {
+        val interceptor = MetricCapturingInterceptor()
+
+        val ddb = lowLevelAccess { withConfig { interceptors += interceptor } }
+        interceptor.assertEmpty()
+
+        // No metric for low-level client
+        lowLevelAccess { scan { tableName = TABLE_NAME } }
+        interceptor.assertMetric(AwsBusinessMetric.DDB_MAPPER, exists = false)
+        interceptor.reset()
+
+        // Metric for high-level client
+        val mapper = mapper(ddb)
+        val table = mapper.getTable(TABLE_NAME, dummySchema)
+        table.scanPaginated { }.collect()
+        interceptor.assertMetric(AwsBusinessMetric.DDB_MAPPER)
+        interceptor.reset()
+
+        // Still no metric for low-level client (i.e., LL wasn't modified by HL)
+        lowLevelAccess { scan { tableName = TABLE_NAME } }
+        interceptor.assertMetric(AwsBusinessMetric.DDB_MAPPER, exists = false)
+        interceptor.reset()
+
+        // Original client can be closed, mapper is unaffected
+        lowLevelAccess { close() }
+        table.scanPaginated { }.collect()
+        interceptor.assertMetric(AwsBusinessMetric.DDB_MAPPER)
+    }
+}
+
+private class MetricCapturingInterceptor : HttpInterceptor {
+    private val capturedMetrics = mutableSetOf<String>()
+
+    override fun readBeforeTransmit(context: ProtocolRequestInterceptorContext<Any, HttpRequest>) {
+        capturedMetrics += context.executionContext[BusinessMetrics]
+    }
+
+    fun assertMetric(metric: BusinessMetric, exists: Boolean = true) {
+        if (exists) {
+            assertTrue(
+                metric.identifier in capturedMetrics,
+                "Expected metrics to contain $metric. Actual values: $capturedMetrics",
+            )
+        } else {
+            assertFalse(
+                metric.identifier in capturedMetrics,
+                "Expected metrics *not* to contain $metric. Actual values: $capturedMetrics",
+            )
+        }
+    }
+
+    fun assertEmpty() {
+        assertTrue(capturedMetrics.isEmpty(), "Expected metrics to be empty. Actual values: $capturedMetrics")
+    }
+
+    fun reset() {
+        capturedMetrics.clear()
+    }
+}

hll/dynamodb-mapper/dynamodb-mapper/common/test/aws/sdk/kotlin/hll/dynamodbmapper/operations/PutItemTest.kt

@@ -43,7 +43,7 @@ class PutItemTest : DdbLocalTest() {
 
         table.putItem { item = Item(id = "foo", value = 42) }
 
-        val resp = ddb.getItem(TABLE_NAME, "id" to "foo")
+        val resp = lowLevelAccess { getItem(TABLE_NAME, "id" to "foo") }
 
         val item = assertNotNull(resp.item)
         assertEquals("foo", item["id"]?.asSOrNull())

hll/dynamodb-mapper/dynamodb-mapper/common/test/aws/sdk/kotlin/hll/dynamodbmapper/testutils/DdbLocalTest.kt

@@ -8,14 +8,19 @@ import aws.sdk.kotlin.hll.dynamodbmapper.DynamoDbMapper
 import aws.sdk.kotlin.hll.dynamodbmapper.items.ItemSchema
 import aws.sdk.kotlin.hll.dynamodbmapper.model.Item
 import aws.sdk.kotlin.runtime.auth.credentials.StaticCredentialsProvider
+import aws.sdk.kotlin.runtime.http.interceptors.AwsBusinessMetric
 import aws.sdk.kotlin.services.dynamodb.DynamoDbClient
 import aws.sdk.kotlin.services.dynamodb.deleteTable
 import aws.sdk.kotlin.services.dynamodb.waiters.waitUntilTableNotExists
+import aws.smithy.kotlin.runtime.client.ProtocolRequestInterceptorContext
+import aws.smithy.kotlin.runtime.http.interceptors.HttpInterceptor
+import aws.smithy.kotlin.runtime.http.request.HttpRequest
 import aws.smithy.kotlin.runtime.net.Host
 import aws.smithy.kotlin.runtime.net.Scheme
 import aws.smithy.kotlin.runtime.net.url.Url
 import io.kotest.core.spec.style.AnnotationSpec
 import kotlinx.coroutines.runBlocking
+import kotlin.test.assertContains
 import kotlin.test.assertEquals
 import kotlin.test.assertNotNull
 
@@ -41,6 +46,9 @@ abstract class DdbLocalTest : AnnotationSpec() {
         }
     }
 
+    private val requests = mutableListOf<HttpRequest>()
+    private val requestInterceptor = RequestCapturingInterceptor(this@DdbLocalTest.requests)
+
     private val ddbHolder = lazy {
         DynamoDbClient {
             endpointUrl = Url {
@@ -55,15 +63,20 @@ abstract class DdbLocalTest : AnnotationSpec() {
                 accessKeyId = "DUMMY"
                 secretAccessKey = "DUMMY"
             }
+
+            interceptors += requestInterceptor
         }
     }
 
     /**
      * An instance of a low-level [DynamoDbClient] utilizing the DynamoDB Local instance which may be used for setting
      * up or verifying various mapper tests. If this is the first time accessing the value, the client will be
      * initialized.
+     *
+     * **Important**: This low-level client should only be accessed via [lowLevelAccess] to ensure that User-Agent
+     * header verification succeeds.
      */
-    val ddb by ddbHolder
+    private val ddb by ddbHolder
 
     private val tempTables = mutableListOf<String>()
 
@@ -95,16 +108,58 @@ abstract class DdbLocalTest : AnnotationSpec() {
         lsis: Map<String, ItemSchema<*>>,
         items: List<Item>,
     ) {
-        ddb.createTable(name, schema, gsis, lsis)
-        tempTables += name
-        ddb.putItems(name, items)
+        lowLevelAccess {
+            createTable(name, schema, gsis, lsis)
+            tempTables += name
+            putItems(name, items)
+        }
     }
 
     /**
      * Returns a [DynamoDbMapper] instance utilizing the DynamoDB Local instance
      * @param config A function to set the configuration of the mapper before it's built
      */
-    fun mapper(config: DynamoDbMapper.Config.Builder.() -> Unit = { }) = DynamoDbMapper(ddb, config)
+    fun mapper(
+        ddb: DynamoDbClient? = null,
+        config: DynamoDbMapper.Config.Builder.() -> Unit = { },
+    ) = DynamoDbMapper(ddb ?: this.ddb, config)
+
+    @BeforeEach
+    fun initializeTest() {
+        requestInterceptor.enabled = true
+    }
+
+    /**
+     * Executes requests on a low-level [DynamoDbClient] and _does not_ log any requests executed in [block]. (This
+     * skips verifying that low-level requests contain the [AwsBusinessMetric.DDB_MAPPER] metric.)
+     */
+    protected suspend fun <T> lowLevelAccess(block: suspend DynamoDbClient.() -> T): T {
+        requestInterceptor.enabled = false
+        return block(ddb).also { requestInterceptor.enabled = true }
+    }
+
+    @AfterEach
+    fun postVerify() {
+        requests.forEach { req ->
+            val uaString = requireNotNull(req.headers["User-Agent"]) {
+                "Missing User-Agent header for request $req"
+            }
+
+            val components = uaString.split(" ")
+
+            val metricsComponent = requireNotNull(components.find { it.startsWith("m/") }) {
+                """User-Agent header "$uaString" doesn't contain business metrics for request $req"""
+            }
+
+            val metrics = metricsComponent.removePrefix("m/").split(",")
+
+            assertContains(
+                metrics,
+                AwsBusinessMetric.DDB_MAPPER.identifier,
+                """Mapper business metric not present in User-Agent header "$uaString" for request $req""",
+            )
+        }
+    }
 
     @AfterAll
     fun cleanUp() {
@@ -120,3 +175,13 @@ abstract class DdbLocalTest : AnnotationSpec() {
         }
     }
 }
+
+private class RequestCapturingInterceptor(val requests: MutableList<HttpRequest>) : HttpInterceptor {
+    var enabled = true
+
+    override fun readBeforeTransmit(context: ProtocolRequestInterceptorContext<Any, HttpRequest>) {
+        if (enabled) {
+            requests += context.protocolRequest
+        }
+    }
+}