Add SERVICE_ENDPOINT metric (#4307)

dagnir · dave-fn · web-flow · commit 7d54ff8231db · 2023-08-15T18:16:30.000-07:00
* Implement SERVICE_ENDPOINT core metric

* Create CollectionStages for SERVICE_ENDPOINT

* Remove erviceEndpointMetricCollectionStages since they're unnecessary

* Move collection of SERVICE_ENDPOINT metric

Simplify by collecting in the API Call metric collection stages where other
request metrics are already being collected.

* Remove unused import

* Update changelog entry

* Fix test

* Remove unused class

---------

Co-authored-by: David Negrete &lt;davidfn@amazon.com&gt;
diff --git a/.changes/next-release/bugfix-AWSSDKforJavav2-0a552ee.json b/.changes/next-release/bugfix-AWSSDKforJavav2-0a552ee.json
@@ -0,0 +1,6 @@
+{
+    "category": "AWS SDK for Java v2", 
+    "contributor": "", 
+    "type": "bugfix", 
+    "description": "Add support for the `SERVICE_ENDPOINT` metric. This metric represents the endpoint (scheme and authority) that the request was sent to."
+}
diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/ApiCallMetricCollectionStage.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/ApiCallMetricCollectionStage.java
@@ -21,6 +21,7 @@
 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.internal.util.MetricUtils;
 import software.amazon.awssdk.core.metrics.CoreMetric;
 import software.amazon.awssdk.http.SdkHttpFullRequest;
 import software.amazon.awssdk.metrics.MetricCollector;
@@ -39,6 +40,7 @@ public ApiCallMetricCollectionStage(RequestPipeline<SdkHttpFullRequest, Response
     @Override
     public Response<OutputT> execute(SdkHttpFullRequest input, RequestExecutionContext context) throws Exception {
         MetricCollector metricCollector = context.executionContext().metricCollector();
+        MetricUtils.collectServiceEndpointMetrics(metricCollector, input);
 
         // Note: at this point, any exception, even a service exception, will
         // be thrown from the wrapped pipeline so we can't use
diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/AsyncApiCallMetricCollectionStage.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/AsyncApiCallMetricCollectionStage.java
@@ -20,6 +20,7 @@
 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.internal.util.MetricUtils;
 import software.amazon.awssdk.core.metrics.CoreMetric;
 import software.amazon.awssdk.http.SdkHttpFullRequest;
 import software.amazon.awssdk.metrics.MetricCollector;
@@ -40,6 +41,7 @@ public AsyncApiCallMetricCollectionStage(RequestPipeline<SdkHttpFullRequest, Com
     @Override
     public CompletableFuture<OutputT> execute(SdkHttpFullRequest input, RequestExecutionContext context) throws Exception {
         MetricCollector metricCollector = context.executionContext().metricCollector();
+        MetricUtils.collectServiceEndpointMetrics(metricCollector, input);
 
         CompletableFuture<OutputT> future = new CompletableFuture<>();
 
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
@@ -18,13 +18,17 @@
 import static software.amazon.awssdk.core.http.HttpResponseHandler.X_AMZN_REQUEST_ID_HEADERS;
 import static software.amazon.awssdk.core.http.HttpResponseHandler.X_AMZ_ID_2_HEADER;
 
+import java.net.URI;
+import java.net.URISyntaxException;
 import java.time.Duration;
 import java.util.concurrent.Callable;
 import java.util.function.Supplier;
 import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.core.exception.SdkClientException;
 import software.amazon.awssdk.core.internal.http.RequestExecutionContext;
 import software.amazon.awssdk.core.metrics.CoreMetric;
 import software.amazon.awssdk.http.HttpMetric;
+import software.amazon.awssdk.http.SdkHttpFullRequest;
 import software.amazon.awssdk.http.SdkHttpFullResponse;
 import software.amazon.awssdk.metrics.MetricCollector;
 import software.amazon.awssdk.metrics.NoOpMetricCollector;
@@ -65,6 +69,23 @@ public static <T> Pair<T, Duration> measureDurationUnsafe(Callable<T> c) throws
         return Pair.of(result, d);
     }
 
+    /**
+     * Collect the SERVICE_ENDPOINT metric for this request.
+     */
+    public static void collectServiceEndpointMetrics(MetricCollector metricCollector, SdkHttpFullRequest httpRequest) {
+        if (metricCollector != null && !(metricCollector instanceof NoOpMetricCollector) && httpRequest != null) {
+            // Only interested in the service endpoint so don't include any path, query, or fragment component
+            URI requestUri = httpRequest.getUri();
+            try {
+                URI serviceEndpoint = new URI(requestUri.getScheme(), requestUri.getAuthority(), null, null, null);
+                metricCollector.reportMetric(CoreMetric.SERVICE_ENDPOINT, serviceEndpoint);
+            } catch (URISyntaxException e) {
+                // This should not happen since getUri() should return a valid URI
+                throw SdkClientException.create("Unable to collect SERVICE_ENDPOINT metric", e);
+            }
+        }
+    }
+
     public static void collectHttpMetrics(MetricCollector metricCollector, SdkHttpFullResponse httpResponse) {
         if (metricCollector != null && !(metricCollector instanceof NoOpMetricCollector) && httpResponse != null) {
             metricCollector.reportMetric(HttpMetric.HTTP_STATUS_CODE, httpResponse.statusCode());
diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/metrics/CoreMetric.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/metrics/CoreMetric.java
@@ -15,6 +15,7 @@
 
 package software.amazon.awssdk.core.metrics;
 
+import java.net.URI;
 import java.time.Duration;
 import software.amazon.awssdk.annotations.SdkPublicApi;
 import software.amazon.awssdk.core.retry.RetryPolicy;
@@ -50,6 +51,12 @@ public final class CoreMetric {
     public static final SdkMetric<Integer> RETRY_COUNT =
         metric("RetryCount", Integer.class, MetricLevel.ERROR);
 
+    /**
+     * The endpoint for the service.
+     */
+    public static final SdkMetric<URI> SERVICE_ENDPOINT =
+        metric("ServiceEndpoint", URI.class, MetricLevel.ERROR);
+
     /**
      * The duration of the API call. This includes all call attempts made.
      *
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,8 +24,10 @@
 
 import java.io.ByteArrayInputStream;
 import java.io.IOException;
+import java.net.URI;
 import java.time.Duration;
 import java.util.List;
+import java.util.concurrent.CompletableFuture;
 import java.util.stream.Collectors;
 import org.junit.After;
 import org.junit.Before;
@@ -41,6 +43,7 @@
 import software.amazon.awssdk.core.exception.SdkException;
 import software.amazon.awssdk.core.metrics.CoreMetric;
 import software.amazon.awssdk.core.internal.metrics.SdkErrorType;
+import software.amazon.awssdk.endpoints.Endpoint;
 import software.amazon.awssdk.http.AbortableInputStream;
 import software.amazon.awssdk.http.ExecutableHttpRequest;
 import software.amazon.awssdk.http.HttpExecuteRequest;
@@ -52,6 +55,8 @@
 import software.amazon.awssdk.metrics.MetricPublisher;
 import software.amazon.awssdk.regions.Region;
 import software.amazon.awssdk.services.protocolrestjson.ProtocolRestJsonClient;
+import software.amazon.awssdk.services.protocolrestjson.endpoints.ProtocolRestJsonEndpointParams;
+import software.amazon.awssdk.services.protocolrestjson.endpoints.ProtocolRestJsonEndpointProvider;
 import software.amazon.awssdk.services.protocolrestjson.model.EmptyModeledException;
 import software.amazon.awssdk.services.protocolrestjson.model.SimpleStruct;
 import software.amazon.awssdk.services.protocolrestjson.paginators.PaginatedOperationWithResultKeyIterable;
@@ -77,13 +82,17 @@ public class CoreMetricsTest {
     @Mock
     private MetricPublisher mockPublisher;
 
+    @Mock
+    private ProtocolRestJsonEndpointProvider mockEndpointProvider;
+
     @Before
     public void setup() throws IOException {
         client = ProtocolRestJsonClient.builder()
                 .httpClient(mockHttpClient)
                 .region(Region.US_WEST_2)
                 .credentialsProvider(mockCredentialsProvider)
                 .overrideConfiguration(c -> c.addMetricPublisher(mockPublisher).retryPolicy(b -> b.numRetries(MAX_RETRIES)))
+                .endpointProvider(mockEndpointProvider)
                 .build();
         AbortableInputStream content = contentStream("{}");
         SdkHttpFullResponse httpResponse = SdkHttpFullResponse.builder()
@@ -116,6 +125,11 @@ public void setup() throws IOException {
             }
             return AwsBasicCredentials.create("foo", "bar");
         });
+
+        when(mockEndpointProvider.resolveEndpoint(any(ProtocolRestJsonEndpointParams.class))).thenReturn(
+            CompletableFuture.completedFuture(Endpoint.builder()
+                                                      .url(URI.create("https://protocolrestjson.amazonaws.com"))
+                                                      .build()));
     }
 
     @After
@@ -183,6 +197,8 @@ public void testApiCall_operationSuccessful_addsMetrics() {
         assertThat(capturedCollection.metricValues(CoreMetric.MARSHALLING_DURATION).get(0))
             .isGreaterThanOrEqualTo(Duration.ZERO);
         assertThat(capturedCollection.metricValues(CoreMetric.RETRY_COUNT)).containsExactly(0);
+        assertThat(capturedCollection.metricValues(CoreMetric.SERVICE_ENDPOINT).get(0)).isEqualTo(URI.create(
+            "https://protocolrestjson.amazonaws.com"));
 
         assertThat(capturedCollection.children()).hasSize(1);
         MetricCollection attemptCollection = capturedCollection.children().get(0);
@@ -280,6 +296,24 @@ public void testApiCall_httpClientThrowsNetworkError_errorTypeIncludedInMetrics(
         }
     }
 
+    @Test
+    public void testApiCall_endpointProviderAddsPathQueryFragment_notReportedInServiceEndpointMetric() {
+        when(mockEndpointProvider.resolveEndpoint(any(ProtocolRestJsonEndpointParams.class)))
+            .thenReturn(CompletableFuture.completedFuture(Endpoint.builder()
+                                                                  .url(URI.create("https://protocolrestjson.amazonaws.com:8080/foo?bar#baz"))
+                                                                  .build()));
+
+        client.allTypes();
+
+        ArgumentCaptor<MetricCollection> collectionCaptor = ArgumentCaptor.forClass(MetricCollection.class);
+        verify(mockPublisher).publish(collectionCaptor.capture());
+
+        MetricCollection capturedCollection = collectionCaptor.getValue();
+
+        URI expectedServiceEndpoint = URI.create("https://protocolrestjson.amazonaws.com:8080");
+        assertThat(capturedCollection.metricValues(CoreMetric.SERVICE_ENDPOINT)).containsExactly(expectedServiceEndpoint);
+    }
+
 
     private static HttpExecuteResponse mockExecuteResponse(SdkHttpFullResponse httpResponse) {
         HttpExecuteResponse mockResponse = mock(HttpExecuteResponse.class);
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
@@ -203,6 +203,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)).toString()
+            .startsWith("http://localhost");
     }
 
     void stubSuccessfulResponse() {

Original file line number	Diff line number	Diff line change
`@@ -203,6 +203,8 @@ private void verifyApiCallCollection(MetricCollection capturedCollection) {`
`203`	`203`	`.isGreaterThanOrEqualTo(Duration.ZERO);`
`204`	`204`	`assertThat(capturedCollection.metricValues(CoreMetric.API_CALL_DURATION).get(0))`
`205`	`205`	`.isGreaterThan(FIXED_DELAY);`
	`206`	`+ assertThat(capturedCollection.metricValues(CoreMetric.SERVICE_ENDPOINT).get(0)).toString()`
	`207`	`+ .startsWith("http://localhost");`
`206`	`208`	`}`
`207`	`209`
`208`	`210`	`void stubSuccessfulResponse() {`