Skip to content

Commit

Permalink
feat: business metrics
Browse files Browse the repository at this point in the history
  • Loading branch information
0marperez committed May 23, 2024
1 parent 2eaeea9 commit 165c837
Show file tree
Hide file tree
Showing 9 changed files with 247 additions and 3 deletions.
23 changes: 23 additions & 0 deletions aws-runtime/aws-http/api/aws-http.api
Original file line number Diff line number Diff line change
Expand Up @@ -162,6 +162,29 @@ public final class aws/sdk/kotlin/runtime/http/interceptors/AwsSpanInterceptor :
public fun readBeforeTransmit (Laws/smithy/kotlin/runtime/client/ProtocolRequestInterceptorContext;)V
}

public final class aws/sdk/kotlin/runtime/http/interceptors/BusinessMetricsInterceptor : aws/smithy/kotlin/runtime/client/Interceptor {
public fun <init> ()V
public fun modifyBeforeAttemptCompletion-gIAlu-s (Laws/smithy/kotlin/runtime/client/ResponseInterceptorContext;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
public fun modifyBeforeCompletion-gIAlu-s (Laws/smithy/kotlin/runtime/client/ResponseInterceptorContext;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
public fun modifyBeforeDeserialization (Laws/smithy/kotlin/runtime/client/ProtocolResponseInterceptorContext;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
public fun modifyBeforeRetryLoop (Laws/smithy/kotlin/runtime/client/ProtocolRequestInterceptorContext;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
public fun modifyBeforeSerialization (Laws/smithy/kotlin/runtime/client/RequestInterceptorContext;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
public fun modifyBeforeSigning (Laws/smithy/kotlin/runtime/client/ProtocolRequestInterceptorContext;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
public fun modifyBeforeTransmit (Laws/smithy/kotlin/runtime/client/ProtocolRequestInterceptorContext;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
public fun readAfterAttempt (Laws/smithy/kotlin/runtime/client/ResponseInterceptorContext;)V
public fun readAfterDeserialization (Laws/smithy/kotlin/runtime/client/ResponseInterceptorContext;)V
public fun readAfterExecution (Laws/smithy/kotlin/runtime/client/ResponseInterceptorContext;)V
public fun readAfterSerialization (Laws/smithy/kotlin/runtime/client/ProtocolRequestInterceptorContext;)V
public fun readAfterSigning (Laws/smithy/kotlin/runtime/client/ProtocolRequestInterceptorContext;)V
public fun readAfterTransmit (Laws/smithy/kotlin/runtime/client/ProtocolResponseInterceptorContext;)V
public fun readBeforeAttempt (Laws/smithy/kotlin/runtime/client/ProtocolRequestInterceptorContext;)V
public fun readBeforeDeserialization (Laws/smithy/kotlin/runtime/client/ProtocolResponseInterceptorContext;)V
public fun readBeforeExecution (Laws/smithy/kotlin/runtime/client/RequestInterceptorContext;)V
public fun readBeforeSerialization (Laws/smithy/kotlin/runtime/client/RequestInterceptorContext;)V
public fun readBeforeSigning (Laws/smithy/kotlin/runtime/client/ProtocolRequestInterceptorContext;)V
public fun readBeforeTransmit (Laws/smithy/kotlin/runtime/client/ProtocolRequestInterceptorContext;)V
}

public final class aws/sdk/kotlin/runtime/http/interceptors/UnsupportedSigningAlgorithmInterceptor : aws/smithy/kotlin/runtime/client/Interceptor {
public fun <init> ()V
public fun modifyBeforeAttemptCompletion-gIAlu-s (Laws/smithy/kotlin/runtime/client/ResponseInterceptorContext;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ import kotlin.jvm.JvmInline

internal const val AWS_EXECUTION_ENV = "AWS_EXECUTION_ENV"
public const val AWS_APP_ID_ENV: String = "AWS_SDK_UA_APP_ID"
private const val USER_AGENT_SPEC_VERSION = "2.1"

// non-standard environment variables/properties
public const val AWS_APP_ID_PROP: String = "aws.userAgentAppId"
Expand Down Expand Up @@ -65,7 +66,7 @@ public data class AwsUserAgentMetadata(
get() = buildList {
add(sdkMetadata)
customMetadata?.extras?.takeIf { it.containsKey("internal") }?.let { add("md/internal") }
add(uaPair("ua", "2.0")) // User agent specification version 2.0
add(uaPair("ua", USER_AGENT_SPEC_VERSION))
add(apiMetadata)
add(osMetadata)
add(languageMetadata)
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,60 @@
/*
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
* SPDX-License-Identifier: Apache-2.0
*/
package aws.sdk.kotlin.runtime.http.interceptors

import aws.sdk.kotlin.runtime.http.middleware.USER_AGENT
import aws.smithy.kotlin.runtime.businessmetrics.businessMetrics
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.http.request.toBuilder

/**
* Appends business metrics to the `User-Agent` header.
*/
public class BusinessMetricsInterceptor : HttpInterceptor {
override suspend fun modifyBeforeTransmit(context: ProtocolRequestInterceptorContext<Any, HttpRequest>): HttpRequest {
context.executionContext.getOrNull(businessMetrics)?.let { metrics ->
val metricsString = formatMetrics(metrics)
val currentUserAgentHeader = context.protocolRequest.headers[USER_AGENT]
val modifiedRequest = context.protocolRequest.toBuilder()

modifiedRequest.headers[USER_AGENT] = currentUserAgentHeader + metricsString

return modifiedRequest.build()
}
return context.protocolRequest
}
}

/**
* Makes sure the metrics do not exceed the maximum size and truncates them if so.
*/
private fun formatMetrics(metrics: MutableSet<String>): String {
val metricsString = metrics.joinToString(",", "m/")
val metricsByteArray = metricsString.toByteArray()
val maxSize = 1024

if (metricsByteArray.size <= maxSize) return metricsString

val commaByte = ','.code.toByte()
var lastCommaIndex: Int? = null

for (i in 0..1023) {
if (metricsByteArray[i] == commaByte) {
lastCommaIndex = i
}
}

lastCommaIndex?.let {
return metricsByteArray.decodeToString(
0,
lastCommaIndex,
true,
)
}

throw IllegalStateException("Business metrics are incorrectly formatted: $metricsString")
}
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@ class AwsUserAgentMetadataTest {
val expected = listOf(
"aws-sdk-kotlin/1.2.3",
"md/internal",
"ua/2.0",
"ua/2.1",
"api/test-service#1.2.3",
"os/linux#ubuntu-20.04",
"lang/kotlin#1.4.31",
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,102 @@
/*
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
* SPDX-License-Identifier: Apache-2.0
*/
package aws.sdk.kotlin.runtime.http.interceptors

import aws.sdk.kotlin.runtime.http.middleware.USER_AGENT
import aws.smithy.kotlin.runtime.businessmetrics.BusinessMetrics
import aws.smithy.kotlin.runtime.businessmetrics.emitBusinessMetric
import aws.smithy.kotlin.runtime.client.ProtocolRequestInterceptorContext
import aws.smithy.kotlin.runtime.collections.get
import aws.smithy.kotlin.runtime.http.*
import aws.smithy.kotlin.runtime.http.request.HttpRequest
import aws.smithy.kotlin.runtime.net.url.Url
import aws.smithy.kotlin.runtime.operation.ExecutionContext
import kotlinx.coroutines.test.runTest
import kotlin.test.Test
import kotlin.test.assertFalse
import kotlin.test.assertTrue

class BusinessMetricsInterceptorTest {
@Test
fun noBusinessMetrics() = runTest {
val executionContext = ExecutionContext()
val interceptor = BusinessMetricsInterceptor()
val request = interceptor.modifyBeforeTransmit(interceptorContext(executionContext))
val userAgentHeader = request.headers[USER_AGENT]!!

assertFalse(userAgentHeader.endsWith("m/"))
}

@Test
fun businessMetrics() = runTest {
val executionContext = ExecutionContext()
executionContext.emitBusinessMetric(BusinessMetrics.S3_EXPRESS_BUCKET)

val interceptor = BusinessMetricsInterceptor()
val request = interceptor.modifyBeforeTransmit(interceptorContext(executionContext))
val userAgentHeader = request.headers[USER_AGENT]!!

assertTrue(
userAgentHeader.endsWith(
"m/${BusinessMetrics.S3_EXPRESS_BUCKET.identifier}",
),
)
}

@Test
fun multipleBusinessMetrics() = runTest {
val executionContext = ExecutionContext()
executionContext.emitBusinessMetric(BusinessMetrics.S3_EXPRESS_BUCKET)
executionContext.emitBusinessMetric(BusinessMetrics.GZIP_REQUEST_COMPRESSION)

val interceptor = BusinessMetricsInterceptor()
val request = interceptor.modifyBeforeTransmit(interceptorContext(executionContext))
val userAgentHeader = request.headers[USER_AGENT]!!

assertTrue(
userAgentHeader.endsWith(
"m/${BusinessMetrics.S3_EXPRESS_BUCKET.identifier},${BusinessMetrics.GZIP_REQUEST_COMPRESSION.identifier}",
),
)
}

@Test
fun truncateBusinessMetrics() = runTest {
val executionContext = ExecutionContext()
executionContext.attributes[aws.smithy.kotlin.runtime.businessmetrics.businessMetrics] = mutableSetOf()

for (i in 0..1024) {
executionContext.attributes[aws.smithy.kotlin.runtime.businessmetrics.businessMetrics].add(i.toString())
}

val rawMetrics = executionContext[aws.smithy.kotlin.runtime.businessmetrics.businessMetrics]
val rawMetricsString = rawMetrics.joinToString(",", "m/")
val rawMetricsByteArray = rawMetricsString.toByteArray()
val maxSize = 1024

assertTrue(rawMetricsByteArray.size >= maxSize)

val interceptor = BusinessMetricsInterceptor()
val request = interceptor.modifyBeforeTransmit(interceptorContext(executionContext))
val userAgentHeader = request.headers[USER_AGENT]!!
val truncatedMetrics = "m/" + userAgentHeader.substringAfter("m/")

assertTrue(truncatedMetrics.toByteArray().size <= maxSize)
assertFalse(truncatedMetrics.endsWith(","))
}
}

private fun interceptorContext(executionContext: ExecutionContext): ProtocolRequestInterceptorContext<Any, HttpRequest> =
object : ProtocolRequestInterceptorContext<Any, HttpRequest> {
override val protocolRequest: HttpRequest = HttpRequest(
HttpMethod.GET,
Url.parse("https://test.aws.com?foo=bar"),
Headers {
append(USER_AGENT, "aws-sdk-kotlin/1.2.3 ua/2.1 api/test-service#1.2.3...")
},
)
override val executionContext: ExecutionContext = executionContext
override val request: Any = Unit
}
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,7 @@ class UserAgentTest {
assertTrue(request.headers.contains(X_AMZ_USER_AGENT))
assertEquals("aws-sdk-kotlin/1.2.3", request.headers[X_AMZ_USER_AGENT])
assertTrue(
request.headers[USER_AGENT]!!.startsWith("aws-sdk-kotlin/1.2.3 ua/2.0 api/test-service#1.2.3"),
request.headers[USER_AGENT]!!.startsWith("aws-sdk-kotlin/1.2.3 ua/2.1 api/test-service#1.2.3"),
"$USER_AGENT header didn't start with expected value. Found: ${request.headers[USER_AGENT]}",
)
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,7 @@ object AwsRuntimeTypes {
object Interceptors : RuntimeTypePackage(AwsKotlinDependency.AWS_HTTP, "interceptors") {
val AddUserAgentMetadataInterceptor = symbol("AddUserAgentMetadataInterceptor")
val UnsupportedSigningAlgorithmInterceptor = symbol("UnsupportedSigningAlgorithmInterceptor")
val BusinessMetricsInterceptor = symbol("BusinessMetricsInterceptor")
}

object Retries {
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,56 @@
/*
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
* SPDX-License-Identifier: Apache-2.0
*/
package aws.sdk.kotlin.codegen

import software.amazon.smithy.kotlin.codegen.core.KotlinWriter
import software.amazon.smithy.kotlin.codegen.core.RuntimeTypes
import software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration
import software.amazon.smithy.kotlin.codegen.integration.SectionWriter
import software.amazon.smithy.kotlin.codegen.integration.SectionWriterBinding
import software.amazon.smithy.kotlin.codegen.rendering.endpoints.EndpointBusinessMetrics
import software.amazon.smithy.kotlin.codegen.rendering.protocol.ProtocolGenerator
import software.amazon.smithy.kotlin.codegen.rendering.protocol.ProtocolMiddleware
import software.amazon.smithy.model.shapes.OperationShape

/**
* Renders the addition of the [BusinessMetricsInterceptor] and endpoint business metrics emitters
*/
class BusinessMetricsIntegration : KotlinIntegration {
override val sectionWriters: List<SectionWriterBinding>
get() = listOf(
SectionWriterBinding(EndpointBusinessMetrics, endpointBusinessMetricsSectionWriter),
)

private val endpointBusinessMetricsSectionWriter = SectionWriter { writer, _ ->
writer.write(
"if (endpoint.attributes.contains(#T)) request.context.#T(#T.SERVICE_ENDPOINT_OVERRIDE)",
RuntimeTypes.Core.BusinessMetrics.serviceEndpointOverride,
RuntimeTypes.Core.BusinessMetrics.EmitBusinessMetrics,
RuntimeTypes.Core.BusinessMetrics.BusinessMetrics,
)

writer.write(
"if (endpoint.attributes.contains(#T)) request.context.#T(#T.ACCOUNT_ID_BASED_ENDPOINT)",
RuntimeTypes.Core.BusinessMetrics.accountIdBasedEndPoint,
RuntimeTypes.Core.BusinessMetrics.EmitBusinessMetrics,
RuntimeTypes.Core.BusinessMetrics.BusinessMetrics,
)
}

override fun customizeMiddleware(
ctx: ProtocolGenerator.GenerationContext,
resolved: List<ProtocolMiddleware>,
): List<ProtocolMiddleware> = resolved + userAgentBusinessMetricsMiddleware

private val userAgentBusinessMetricsMiddleware = object : ProtocolMiddleware {
override val name: String = "UserAgentBusinessMetrics"
override fun render(ctx: ProtocolGenerator.GenerationContext, op: OperationShape, writer: KotlinWriter) {
writer.write(
"op.interceptors.add(#T())",
AwsRuntimeTypes.Http.Interceptors.BusinessMetricsInterceptor,
)
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -42,3 +42,4 @@ aws.sdk.kotlin.codegen.customization.cloudfrontkeyvaluestore.BackfillSigV4ACusto
aws.sdk.kotlin.codegen.customization.s3.express.SigV4S3ExpressAuthSchemeIntegration
aws.sdk.kotlin.codegen.customization.s3.express.S3ExpressIntegration
aws.sdk.kotlin.codegen.customization.s3.S3ExpiresIntegration
aws.sdk.kotlin.codegen.BusinessMetricsIntegration

0 comments on commit 165c837

Please sign in to comment.