Create CollectionStages for SERVICE_ENDPOINT

dave-fn · dave-fn · commit 07b0746f19a6 · 2023-03-10T10:29:02.000-07:00
diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/AmazonAsyncHttpClient.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/AmazonAsyncHttpClient.java
@@ -36,6 +36,7 @@
 import software.amazon.awssdk.core.internal.http.pipeline.stages.AsyncBeforeTransmissionExecutionInterceptorsStage;
 import software.amazon.awssdk.core.internal.http.pipeline.stages.AsyncExecutionFailureExceptionReportingStage;
 import software.amazon.awssdk.core.internal.http.pipeline.stages.AsyncRetryableStage;
+import software.amazon.awssdk.core.internal.http.pipeline.stages.AsyncServiceEndpointMetricCollectionStage;
 import software.amazon.awssdk.core.internal.http.pipeline.stages.AsyncSigningStage;
 import software.amazon.awssdk.core.internal.http.pipeline.stages.MakeAsyncHttpRequestStage;
 import software.amazon.awssdk.core.internal.http.pipeline.stages.MakeRequestImmutableStage;
@@ -180,7 +181,8 @@ public <OutputT> CompletableFuture<OutputT> execute(
                                         .then(async(() -> new AfterExecutionInterceptorsStage<>()))
                                         .wrappedWith(AsyncExecutionFailureExceptionReportingStage::new)
                                         .wrappedWith(AsyncApiCallTimeoutTrackingStage::new)
-                                        .wrappedWith(AsyncApiCallMetricCollectionStage::new)::build)::build)
+                                        .wrappedWith(AsyncApiCallMetricCollectionStage::new)
+                                        .wrappedWith(AsyncServiceEndpointMetricCollectionStage::new)::build)::build)
                         .build(httpClientDependencies)
                         .execute(request, createRequestExecutionDependencies());
             } catch (RuntimeException e) {
diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/AmazonSyncHttpClient.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/AmazonSyncHttpClient.java
@@ -43,6 +43,8 @@
 import software.amazon.awssdk.core.internal.http.pipeline.stages.MergeCustomHeadersStage;
 import software.amazon.awssdk.core.internal.http.pipeline.stages.MergeCustomQueryParamsStage;
 import software.amazon.awssdk.core.internal.http.pipeline.stages.RetryableStage;
+import software.amazon.awssdk.core.internal.http.pipeline.stages.ServiceEndpointAttemptMetricCollectionStage;
+import software.amazon.awssdk.core.internal.http.pipeline.stages.ServiceEndpointMetricCollectionStage;
 import software.amazon.awssdk.core.internal.http.pipeline.stages.SigningStage;
 import software.amazon.awssdk.core.internal.http.pipeline.stages.TimeoutExceptionHandlingStage;
 import software.amazon.awssdk.core.internal.http.pipeline.stages.UnwrapResponseContainer;
@@ -182,10 +184,13 @@ public <OutputT> OutputT execute(HttpResponseHandler<Response<OutputT>> response
                                          .wrappedWith(ApiCallAttemptTimeoutTrackingStage::new)
                                          .wrappedWith(TimeoutExceptionHandlingStage::new)
                                          .wrappedWith((deps, wrapped) -> new ApiCallAttemptMetricCollectionStage<>(wrapped))
+                                         .wrappedWith((deps, wrapped) ->
+                                                          new ServiceEndpointAttemptMetricCollectionStage<>(wrapped))
                                          .wrappedWith(RetryableStage::new)::build)
                                .wrappedWith(StreamManagingStage::new)
                                .wrappedWith(ApiCallTimeoutTrackingStage::new)::build)
                                .wrappedWith((deps, wrapped) -> new ApiCallMetricCollectionStage<>(wrapped))
+                               .wrappedWith((deps, wrapped) -> new ServiceEndpointMetricCollectionStage<>(wrapped))
                     .then(() -> new UnwrapResponseContainer<>())
                     .then(() -> new AfterExecutionInterceptorsStage<>())
                     .wrappedWith(ExecutionFailureExceptionReportingStage::new)
diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/ApiCallAttemptMetricCollectionStage.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/ApiCallAttemptMetricCollectionStage.java
@@ -16,7 +16,6 @@
 package software.amazon.awssdk.core.internal.http.pipeline.stages;
 
 import static software.amazon.awssdk.core.internal.util.MetricUtils.collectHttpMetrics;
-import static software.amazon.awssdk.core.internal.util.MetricUtils.collectHttpRequestMetrics;
 import static software.amazon.awssdk.core.internal.util.MetricUtils.createAttemptMetricsCollector;
 
 import java.time.Duration;
@@ -48,8 +47,6 @@ public Response<OutputT> execute(SdkHttpFullRequest input, RequestExecutionConte
         context.attemptMetricCollector(apiCallAttemptMetrics);
         reportBackoffDelay(context);
 
-        collectHttpRequestMetrics(apiCallAttemptMetrics, input);
-
         Response<OutputT> response = wrapped.execute(input, context);
 
         collectHttpMetrics(apiCallAttemptMetrics, response.httpResponse());
diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/AsyncApiCallAttemptMetricCollectionStage.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/AsyncApiCallAttemptMetricCollectionStage.java
@@ -16,7 +16,6 @@
 package software.amazon.awssdk.core.internal.http.pipeline.stages;
 
 import static software.amazon.awssdk.core.internal.util.MetricUtils.collectHttpMetrics;
-import static software.amazon.awssdk.core.internal.util.MetricUtils.collectHttpRequestMetrics;
 import static software.amazon.awssdk.core.internal.util.MetricUtils.createAttemptMetricsCollector;
 
 import java.time.Duration;
@@ -53,8 +52,6 @@ public CompletableFuture<Response<OutputT>> execute(SdkHttpFullRequest input,
         context.attemptMetricCollector(apiCallAttemptMetrics);
         reportBackoffDelay(context);
 
-        collectHttpRequestMetrics(apiCallAttemptMetrics, input);
-
         CompletableFuture<Response<OutputT>> executeFuture = wrapped.execute(input, context);
         CompletableFuture<Response<OutputT>> metricsCollectedFuture = executeFuture.whenComplete((r, t) -> {
             if (t == null) {
diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/AsyncServiceEndpointMetricCollectionStage.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/AsyncServiceEndpointMetricCollectionStage.java
@@ -0,0 +1,59 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.core.internal.http.pipeline.stages;
+
+import java.util.concurrent.CompletableFuture;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.core.internal.http.RequestExecutionContext;
+import software.amazon.awssdk.core.internal.http.pipeline.RequestPipeline;
+import software.amazon.awssdk.core.metrics.CoreMetric;
+import software.amazon.awssdk.http.SdkHttpFullRequest;
+import software.amazon.awssdk.metrics.MetricCollector;
+import software.amazon.awssdk.utils.CompletableFutureUtils;
+
+/**
+ * Wrapper pipeline that tracks the {@link CoreMetric#SERVICE_ENDPOINT} metric.
+ */
+@SdkInternalApi
+public final class AsyncServiceEndpointMetricCollectionStage<OutputT> implements RequestPipeline<SdkHttpFullRequest,
+    CompletableFuture<OutputT>> {
+    private final RequestPipeline<SdkHttpFullRequest, CompletableFuture<OutputT>> wrapped;
+
+    public AsyncServiceEndpointMetricCollectionStage(RequestPipeline<SdkHttpFullRequest, CompletableFuture<OutputT>> wrapped) {
+        this.wrapped = wrapped;
+    }
+
+    @Override
+    public CompletableFuture<OutputT> execute(SdkHttpFullRequest input, RequestExecutionContext context) throws Exception {
+        MetricCollector metricCollector = context.executionContext().metricCollector();
+
+        CompletableFuture<OutputT> future = new CompletableFuture<>();
+
+        CompletableFuture<OutputT> executeFuture = wrapped.execute(input, context);
+
+        executeFuture.whenComplete((r, t) -> {
+            metricCollector.reportMetric(CoreMetric.SERVICE_ENDPOINT, input.getUri());
+
+            if (t != null) {
+                future.completeExceptionally(t);
+            } else {
+                future.complete(r);
+            }
+        });
+
+        return CompletableFutureUtils.forwardExceptionTo(future, executeFuture);
+    }
+}
diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/ServiceEndpointAttemptMetricCollectionStage.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/ServiceEndpointAttemptMetricCollectionStage.java
@@ -0,0 +1,53 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.core.internal.http.pipeline.stages;
+
+import static software.amazon.awssdk.core.internal.util.MetricUtils.collectServiceEndpointMetrics;
+import static software.amazon.awssdk.core.internal.util.MetricUtils.createAttemptMetricsCollector;
+
+import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.core.Response;
+import software.amazon.awssdk.core.internal.http.RequestExecutionContext;
+import software.amazon.awssdk.core.internal.http.pipeline.RequestPipeline;
+import software.amazon.awssdk.core.internal.http.pipeline.RequestToResponsePipeline;
+import software.amazon.awssdk.http.SdkHttpFullRequest;
+import software.amazon.awssdk.metrics.MetricCollector;
+
+/**
+ * Wrapper pipeline that initializes and tracks the API call attempt metric collection. This wrapper and any wrapped
+ * stages will track API call attempt metrics.
+ */
+@SdkInternalApi
+public final class ServiceEndpointAttemptMetricCollectionStage<OutputT> implements RequestToResponsePipeline<OutputT>  {
+    private final RequestPipeline<SdkHttpFullRequest, Response<OutputT>> wrapped;
+
+    public ServiceEndpointAttemptMetricCollectionStage(RequestPipeline<SdkHttpFullRequest, Response<OutputT>> wrapped) {
+        this.wrapped = wrapped;
+    }
+
+    @Override
+    public Response<OutputT> execute(SdkHttpFullRequest input, RequestExecutionContext context) throws Exception {
+        MetricCollector apiCallAttemptMetrics = createAttemptMetricsCollector(context);
+        context.attemptMetricCollector(apiCallAttemptMetrics);
+
+        Response<OutputT> response = wrapped.execute(input, context);
+
+        collectServiceEndpointMetrics(apiCallAttemptMetrics, input);
+
+        return response;
+    }
+
+}
diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/ServiceEndpointMetricCollectionStage.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/ServiceEndpointMetricCollectionStage.java
@@ -0,0 +1,45 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.core.internal.http.pipeline.stages;
+
+import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.core.Response;
+import software.amazon.awssdk.core.internal.http.RequestExecutionContext;
+import software.amazon.awssdk.core.internal.http.pipeline.RequestPipeline;
+import software.amazon.awssdk.core.internal.http.pipeline.RequestToResponsePipeline;
+import software.amazon.awssdk.core.metrics.CoreMetric;
+import software.amazon.awssdk.http.SdkHttpFullRequest;
+import software.amazon.awssdk.metrics.MetricCollector;
+
+/**
+ * Wrapper pipeline that tracks the {@link CoreMetric#SERVICE_ENDPOINT} metric.
+ */
+@SdkInternalApi
+public class ServiceEndpointMetricCollectionStage<OutputT> implements RequestToResponsePipeline<OutputT> {
+    private final RequestPipeline<SdkHttpFullRequest, Response<OutputT>>  wrapped;
+
+    public ServiceEndpointMetricCollectionStage(RequestPipeline<SdkHttpFullRequest, Response<OutputT>> wrapped) {
+        this.wrapped = wrapped;
+    }
+
+    @Override
+    public Response<OutputT> execute(SdkHttpFullRequest input, RequestExecutionContext context) throws Exception {
+        MetricCollector metricCollector = context.executionContext().metricCollector();
+        metricCollector.reportMetric(CoreMetric.SERVICE_ENDPOINT, input.getUri());
+
+        return wrapped.execute(input, context);
+    }
+}
diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/util/MetricUtils.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/util/MetricUtils.java
@@ -66,7 +66,7 @@ public static <T> Pair<T, Duration> measureDurationUnsafe(Callable<T> c) throws
         return Pair.of(result, d);
     }
 
-    public static void collectHttpRequestMetrics(MetricCollector metricCollector, SdkHttpFullRequest httpRequest) {
+    public static void collectServiceEndpointMetrics(MetricCollector metricCollector, SdkHttpFullRequest httpRequest) {
         if (metricCollector != null && !(metricCollector instanceof NoOpMetricCollector) && httpRequest != null) {
             metricCollector.reportMetric(CoreMetric.SERVICE_ENDPOINT, httpRequest.getUri());
         }
diff --git a/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/metrics/CoreMetricsTest.java b/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/metrics/CoreMetricsTest.java
@@ -24,6 +24,7 @@
 
 import java.io.ByteArrayInputStream;
 import java.io.IOException;
+import java.net.URI;
 import java.time.Duration;
 import java.util.List;
 import java.util.stream.Collectors;
diff --git a/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/metrics/async/BaseAsyncCoreMetricsTest.java b/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/metrics/async/BaseAsyncCoreMetricsTest.java
@@ -28,7 +28,6 @@
 import java.time.Duration;
 import java.util.concurrent.CompletableFuture;
 import java.util.function.Supplier;
-import org.junit.Ignore;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
@@ -201,8 +200,8 @@ private void verifyApiCallCollection(MetricCollection capturedCollection) {
             .isGreaterThanOrEqualTo(Duration.ZERO);
         assertThat(capturedCollection.metricValues(CoreMetric.API_CALL_DURATION).get(0))
             .isGreaterThan(FIXED_DELAY);
-        assertThat(capturedCollection.metricValues(CoreMetric.SERVICE_ENDPOINT).get(0))
-            .isEqualTo(URI.create("http://localhost"));
+        assertThat(capturedCollection.metricValues(CoreMetric.SERVICE_ENDPOINT).get(0)).toString()
+            .startsWith("http://localhost");
     }
 
     void stubSuccessfulResponse() {

Original file line number	Diff line number	Diff line change
`@@ -66,7 +66,7 @@ public static <T> Pair<T, Duration> measureDurationUnsafe(Callable<T> c) throws`
`66`	`66`	`return Pair.of(result, d);`
`67`	`67`	`}`
`68`	`68`
`69`		`- public static void collectHttpRequestMetrics(MetricCollector metricCollector, SdkHttpFullRequest httpRequest) {`
	`69`	`+ public static void collectServiceEndpointMetrics(MetricCollector metricCollector, SdkHttpFullRequest httpRequest) {`
`70`	`70`	`if (metricCollector != null && !(metricCollector instanceof NoOpMetricCollector) && httpRequest != null) {`
`71`	`71`	`metricCollector.reportMetric(CoreMetric.SERVICE_ENDPOINT, httpRequest.getUri());`
`72`	`72`	`}`