-
Notifications
You must be signed in to change notification settings - Fork 26
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
7 changed files
with
268 additions
and
16 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
107 changes: 107 additions & 0 deletions
107
...-engine-okhttp/jvm/src/aws/smithy/kotlin/runtime/http/engine/okhttp/MetricsInterceptor.kt
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,107 @@ | ||
| /* | ||
| * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
| package aws.smithy.kotlin.runtime.http.engine.okhttp | ||
|
|
||
| import aws.smithy.kotlin.runtime.telemetry.metrics.MonotonicCounter | ||
| import aws.smithy.kotlin.runtime.util.Attributes | ||
| import aws.smithy.kotlin.runtime.util.attributesOf | ||
| import okhttp3.* | ||
| import okio.* | ||
|
|
||
| /** | ||
| * Instrument the HTTP throughput metrics (e.g. bytes rcvd/sent) | ||
| */ | ||
| internal object MetricsInterceptor : Interceptor { | ||
| override fun intercept(chain: Interceptor.Chain): Response { | ||
| val originalRequest = chain.request() | ||
| val metrics = originalRequest.tag<SdkRequestTag>()?.metrics ?: return chain.proceed(originalRequest) | ||
|
|
||
| val attrs = attributesOf { "server.address" to "${originalRequest.url.host}:${originalRequest.url.port}" } | ||
| val request = if (originalRequest.body != null) { | ||
| originalRequest.newBuilder() | ||
| .method(originalRequest.method, originalRequest.body?.instrument(metrics.bytesSent, attrs)) | ||
| .build() | ||
| } else { | ||
| originalRequest | ||
| } | ||
|
|
||
| val originalResponse = chain.proceed(request) | ||
| val response = if (originalResponse.body.contentLength() != 0L) { | ||
| originalResponse.newBuilder() | ||
| .body(originalResponse.body.instrument(metrics.bytesReceived, attrs)) | ||
| .build() | ||
| } else { | ||
| originalResponse | ||
| } | ||
|
|
||
| return response | ||
| } | ||
| } | ||
|
|
||
| internal class InstrumentedSink( | ||
| private val delegate: BufferedSink, | ||
| private val counter: MonotonicCounter, | ||
| private val attributes: Attributes, | ||
| ) : Sink by delegate { | ||
| override fun write(source: Buffer, byteCount: Long) { | ||
| delegate.write(source, byteCount) | ||
| counter.add(byteCount, attributes) | ||
| } | ||
| override fun close() { | ||
| delegate.emit() | ||
| delegate.close() | ||
| } | ||
| } | ||
|
|
||
| internal class InstrumentedRequestBody( | ||
| private val delegate: RequestBody, | ||
| private val counter: MonotonicCounter, | ||
| private val attributes: Attributes, | ||
| ) : RequestBody() { | ||
| override fun contentType(): MediaType? = delegate.contentType() | ||
| override fun isOneShot(): Boolean = delegate.isOneShot() | ||
| override fun isDuplex(): Boolean = delegate.isDuplex() | ||
| override fun contentLength(): Long = delegate.contentLength() | ||
| override fun writeTo(sink: BufferedSink) { | ||
| val metricsSink = InstrumentedSink(sink, counter, attributes).buffer() | ||
| delegate.writeTo(metricsSink) | ||
| metricsSink.close() | ||
| } | ||
| } | ||
|
|
||
| internal fun RequestBody.instrument(counter: MonotonicCounter, attributes: Attributes): RequestBody = | ||
| InstrumentedRequestBody(this, counter, attributes) | ||
|
|
||
| internal class InstrumentedSource( | ||
| private val delegate: Source, | ||
| private val counter: MonotonicCounter, | ||
| private val attributes: Attributes, | ||
| ) : Source by delegate { | ||
| override fun timeout(): Timeout = delegate.timeout() | ||
| override fun read(sink: Buffer, byteCount: Long): Long { | ||
| val rc = delegate.read(sink, byteCount) | ||
| if (rc > 0L) { | ||
| counter.add(rc, attributes) | ||
| } | ||
| return rc | ||
| } | ||
| override fun close() { | ||
| delegate.close() | ||
| } | ||
| } | ||
|
|
||
| internal class InstrumentedResponseBody( | ||
| private val delegate: ResponseBody, | ||
| private val counter: MonotonicCounter, | ||
| private val attributes: Attributes, | ||
| ) : ResponseBody() { | ||
| override fun contentType(): MediaType? = delegate.contentType() | ||
| override fun contentLength(): Long = delegate.contentLength() | ||
| override fun source(): BufferedSource = | ||
| InstrumentedSource(delegate.source(), counter, attributes).buffer() | ||
| } | ||
|
|
||
| internal fun ResponseBody.instrument(counter: MonotonicCounter, attributes: Attributes): ResponseBody = | ||
| InstrumentedResponseBody(this, counter, attributes) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
129 changes: 129 additions & 0 deletions
129
...ne-okhttp/jvm/test/aws/smithy/kotlin/runtime/http/engine/okhttp/MetricsInterceptorTest.kt
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,129 @@ | ||
| /* | ||
| * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
| package aws.smithy.kotlin.runtime.http.engine.okhttp | ||
|
|
||
| import aws.smithy.kotlin.runtime.ExperimentalApi | ||
| import aws.smithy.kotlin.runtime.http.engine.internal.HttpClientMetrics | ||
| import aws.smithy.kotlin.runtime.operation.ExecutionContext | ||
| import aws.smithy.kotlin.runtime.telemetry.otel.OpenTelemetryProvider | ||
| import aws.smithy.kotlin.runtime.util.emptyAttributes | ||
| import io.opentelemetry.api.common.AttributeKey | ||
| import io.opentelemetry.sdk.metrics.data.MetricData | ||
| import io.opentelemetry.sdk.testing.junit5.OpenTelemetryExtension | ||
| import okhttp3.* | ||
| import okhttp3.MediaType.Companion.toMediaType | ||
| import okhttp3.RequestBody.Companion.toRequestBody | ||
| import okhttp3.ResponseBody.Companion.toResponseBody | ||
| import okio.blackholeSink | ||
| import okio.buffer | ||
| import org.junit.jupiter.api.extension.RegisterExtension | ||
| import kotlin.coroutines.EmptyCoroutineContext | ||
| import kotlin.test.Test | ||
| import kotlin.test.assertEquals | ||
| import kotlin.test.fail | ||
|
|
||
| @OptIn(ExperimentalApi::class) | ||
| class MetricsInterceptorTest { | ||
| companion object { | ||
| @JvmField | ||
| @RegisterExtension | ||
| val otelTesting = OpenTelemetryExtension.create() | ||
| } | ||
|
|
||
| private val provider = OpenTelemetryProvider(otelTesting.openTelemetry) | ||
| private val meter = provider.meterProvider.getOrCreateMeter("test") | ||
|
|
||
| @Test | ||
| fun testInstrumentedSource() { | ||
| val source = okio.Buffer() | ||
| val data = "a".repeat(15 * 1024) | ||
| source.writeUtf8(data) | ||
|
|
||
| val sink = okio.Buffer() | ||
| val counter = meter.createMonotonicCounter("TestCounter", "By") | ||
| val instrumented = InstrumentedSource(source, counter, emptyAttributes()) | ||
| do { | ||
| val rc = instrumented.read(sink, 399) | ||
| } while (rc >= 0L) | ||
|
|
||
| assertEquals(data.length.toLong(), sink.size) | ||
|
|
||
| val counted = otelTesting.metrics.first().longCounterSum() | ||
| assertEquals(data.length.toLong(), counted) | ||
| } | ||
|
|
||
| @Test | ||
| fun testInstrumentedSink() { | ||
| val source = okio.Buffer() | ||
| val data = "b".repeat(13 * 1024) | ||
| source.writeUtf8(data) | ||
|
|
||
| val sink = okio.Buffer() | ||
| val counter = meter.createMonotonicCounter("TestCounter", "By") | ||
| val instrumented = InstrumentedSink(sink, counter, emptyAttributes()) | ||
|
|
||
| val buffered = instrumented.buffer() | ||
| buffered.writeAll(source) | ||
| buffered.close() | ||
|
|
||
| assertEquals(data.length.toLong(), sink.size) | ||
|
|
||
| val counted = otelTesting.metrics.first().longCounterSum() | ||
| assertEquals(data.length.toLong(), counted) | ||
| } | ||
|
|
||
| @Test | ||
| fun testMetricsInterceptor() { | ||
| val reqData = "a".repeat(15 * 1024) | ||
| val reqBody = reqData.toRequestBody() | ||
| val metrics = HttpClientMetrics("test", provider) | ||
| val tag = SdkRequestTag(ExecutionContext(), EmptyCoroutineContext, metrics) | ||
| val request = Request.Builder() | ||
| .url("https://localhost:1/") | ||
| .method("PUT", reqBody) | ||
| .tag<SdkRequestTag>(tag) | ||
| .build() | ||
|
|
||
| val respData = "b".repeat(13 * 1024) | ||
| val respBody = respData.toResponseBody("text/plain; charset=utf-8".toMediaType()) | ||
| val mockResp = Response.Builder() | ||
| .request(request) | ||
| .protocol(Protocol.HTTP_1_1) | ||
| .code(200) | ||
| .message("Intercepted") | ||
| .body(respBody) | ||
| .build() | ||
|
|
||
| val client = OkHttpClient.Builder() | ||
| .addInterceptor(MetricsInterceptor) | ||
| .addInterceptor { chain -> | ||
| // consume the body and short circuit with a mock response | ||
| chain.request().body?.writeTo(blackholeSink().buffer()) | ||
| mockResp | ||
| } | ||
| .build() | ||
|
|
||
| val resp = client.newCall(request).execute() | ||
| val actualRespData = resp.body.source().readByteArray().decodeToString() | ||
| assertEquals(respData, actualRespData) | ||
|
|
||
| val actualBytesSent = otelTesting.metrics | ||
| .find { it.name == "smithy.client.http.bytes_sent" } ?: fail("expected bytes_sent") | ||
|
|
||
| val actualBytesReceived = otelTesting.metrics | ||
| .find { it.name == "smithy.client.http.bytes_received" } ?: fail("expected bytes_received") | ||
|
|
||
| assertEquals(reqData.length.toLong(), actualBytesSent.longCounterSum()) | ||
| assertEquals(respData.length.toLong(), actualBytesReceived.longCounterSum()) | ||
|
|
||
| val bytesSentAttr = actualBytesSent.longSumData.points.first().attributes.get(AttributeKey.stringKey("server.address")) | ||
| val bytesRecvAttr = actualBytesSent.longSumData.points.first().attributes.get(AttributeKey.stringKey("server.address")) | ||
| val expectedAttr = "localhost:1" | ||
| assertEquals(expectedAttr, bytesRecvAttr) | ||
| assertEquals(expectedAttr, bytesSentAttr) | ||
| } | ||
|
|
||
| private fun MetricData.longCounterSum(): Long = longSumData.points.sumOf { it.value } | ||
| } |
Oops, something went wrong.