Skip to content

Commit a3a62b4

Browse files
zoewanggzoewangg
committed
Add metrics codegen changes for async clients
1 parent 0ee0be3 commit a3a62b4

File tree

14 files changed

+512
-74
lines changed

14 files changed

+512
-74
lines changed

‎codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/AsyncClientClass.java

Lines changed: 30 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
import java.net.URI;
3333
import java.nio.ByteBuffer;
3434
import java.util.List;
35+
import java.util.Optional;
3536
import java.util.concurrent.Executor;
3637
import java.util.stream.Collectors;
3738
import javax.lang.model.element.Modifier;
@@ -61,11 +62,17 @@
6162
import software.amazon.awssdk.core.client.handler.AsyncClientHandler;
6263
import software.amazon.awssdk.core.endpointdiscovery.EndpointDiscoveryRefreshCache;
6364
import software.amazon.awssdk.core.endpointdiscovery.EndpointDiscoveryRequest;
65+
import software.amazon.awssdk.core.internal.util.MetricUtils;
66+
import software.amazon.awssdk.core.metrics.CoreMetric;
67+
import software.amazon.awssdk.metrics.MetricCollector;
68+
import software.amazon.awssdk.metrics.MetricPublisher;
6469
import software.amazon.awssdk.protocols.json.AwsJsonProtocolFactory;
6570
import software.amazon.awssdk.utils.CompletableFutureUtils;
6671
import software.amazon.awssdk.utils.FunctionalUtils;
6772

6873
public final class AsyncClientClass extends AsyncClientInterface {
74+
private static final String PUBLISHER_NAME = "metricPublisher";
75+
private static final String METRIC_COLLECTOR_NAME = "apiCallMetricCollector";
6976
private final IntermediateModel model;
7077
private final PoetExtensions poetExtensions;
7178
private final ClassName className;
@@ -190,9 +197,18 @@ private MethodSpec closeMethod() {
190197
protected MethodSpec.Builder operationBody(MethodSpec.Builder builder, OperationModel opModel) {
191198

192199
builder.addModifiers(Modifier.PUBLIC)
193-
.addAnnotation(Override.class)
194-
.beginControlFlow("try")
195-
.addCode(ClientClassUtils.callApplySignerOverrideMethod(opModel))
200+
.addAnnotation(Override.class);
201+
202+
builder.addStatement("$1T $2N = $1T.create($3S)",
203+
MetricCollector.class, METRIC_COLLECTOR_NAME, "ApiCall");
204+
builder.beginControlFlow("try");
205+
206+
builder.addStatement("$N.reportMetric($T.$L, $S)", METRIC_COLLECTOR_NAME, CoreMetric.class, "SERVICE_ID",
207+
model.getMetadata().getServiceId());
208+
builder.addStatement("$N.reportMetric($T.$L, $S)", METRIC_COLLECTOR_NAME, CoreMetric.class, "OPERATION_NAME",
209+
opModel.getOperationName());
210+
211+
builder.addCode(ClientClassUtils.callApplySignerOverrideMethod(opModel))
196212
.addCode(ClientClassUtils.addEndpointTraitCode(opModel))
197213
.addCode(protocolSpec.responseHandler(model, opModel));
198214
protocolSpec.errorResponseHandler(opModel).ifPresent(builder::addCode);
@@ -224,8 +240,17 @@ protected MethodSpec.Builder operationBody(MethodSpec.Builder builder, Operation
224240
"() -> $N.exceptionOccurred(t))", paramName);
225241
}
226242

227-
return builder.addStatement("return $T.failedFuture(t)", CompletableFutureUtils.class)
228-
.endControlFlow();
243+
builder.addStatement("$T<$T> $N = $T.resolvePublisher(clientConfiguration, $N.overrideConfiguration().orElse(null))",
244+
Optional.class,
245+
MetricPublisher.class,
246+
PUBLISHER_NAME,
247+
MetricUtils.class,
248+
opModel.getInput().getVariableName())
249+
.addStatement("$N.ifPresent(p -> p.publish($N.collect()))", PUBLISHER_NAME, "apiCallMetricCollector")
250+
.addStatement("return $T.failedFuture(t)", CompletableFutureUtils.class)
251+
.endControlFlow();
252+
253+
return builder;
229254
}
230255

231256
@Override

‎codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/AsyncClientInterface.java

Lines changed: 0 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,6 @@
4747
import software.amazon.awssdk.core.SdkClient;
4848
import software.amazon.awssdk.core.async.AsyncRequestBody;
4949
import software.amazon.awssdk.core.async.AsyncResponseTransformer;
50-
import software.amazon.awssdk.metrics.MetricCollector;
5150
import software.amazon.awssdk.regions.providers.DefaultAwsRegionProviderChain;
5251

5352
public class AsyncClientInterface implements ClassSpec {
@@ -285,12 +284,6 @@ private MethodSpec traditionalMethod(OperationModel opModel) {
285284
.addParameter(requestType, opModel.getInput().getVariableName())
286285
.addJavadoc(opModel.getDocs(model, ClientType.ASYNC));
287286

288-
289-
String metricCollectorName = "apiCallMetricCollector";
290-
291-
builder.addStatement("$1T $2N = $1T.create($3S)",
292-
MetricCollector.class, metricCollectorName, "ApiCall");
293-
294287
if (opModel.hasStreamingInput()) {
295288
builder.addParameter(ClassName.get(AsyncRequestBody.class), "requestBody");
296289
} else if (opModel.hasEventStreamInput()) {

‎codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/specs/JsonProtocolSpec.java

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
import java.util.Optional;
2929
import java.util.concurrent.CompletableFuture;
3030
import javax.lang.model.element.Modifier;
31+
import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
3132
import software.amazon.awssdk.awscore.eventstream.EventStreamAsyncResponseTransformer;
3233
import software.amazon.awssdk.awscore.eventstream.EventStreamTaggedUnionPojoSupplier;
3334
import software.amazon.awssdk.awscore.eventstream.RestEventStreamAsyncResponseTransformer;
@@ -266,6 +267,8 @@ public CodeBlock asyncExecutionHandler(IntermediateModel intermediateModel, Oper
266267
asyncResponseTransformerVariable(isStreaming, isRestJson, opModel));
267268
String whenComplete = whenCompleteBody(opModel, customerResponseHandler);
268269
if (!whenComplete.isEmpty()) {
270+
builder.addStatement("$T requestOverrideConfig = $L.overrideConfiguration().orElse(null)",
271+
AwsRequestOverrideConfiguration.class, opModel.getInput().getVariableName());
269272
builder.add("executeFuture$L;", whenComplete);
270273
}
271274
if (opModel.hasEventStreamOutput()) {
@@ -326,7 +329,7 @@ private String whenCompleteBody(OperationModel operationModel, String responseHa
326329
return streamingOutputWhenComplete(responseHandlerName);
327330
} else {
328331
// Non streaming can just return the future as is
329-
return "";
332+
return publishMetricsWhenComplete();
330333
}
331334
}
332335

@@ -336,6 +339,7 @@ private String whenCompleteBody(OperationModel operationModel, String responseHa
336339
* {@link EventStreamAsyncResponseTransformer}. Failure is notified via the normal future (the one returned by the client
337340
* handler).
338341
*
342+
*
339343
* @param responseHandlerName Variable name of response handler customer passed in.
340344
* @return whenComplete to append to future.
341345
*/
@@ -347,12 +351,12 @@ private String eventStreamOutputWhenComplete(String responseHandlerName) {
347351
+ " } finally {"
348352
+ " future.completeExceptionally(e);"
349353
+ " }"
350-
+ " }%n"
351-
+ "})", responseHandlerName);
354+
+ " }"
355+
+ "%s"
356+
+ "})", responseHandlerName, publishMetrics());
352357
}
353358

354359

355-
356360
@Override
357361
public Optional<MethodSpec> createErrorResponseHandler() {
358362
ClassName httpResponseHandler = ClassName.get(HttpResponseHandler.class);

‎codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/specs/ProtocolSpec.java

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -155,7 +155,8 @@ default String streamingOutputWhenComplete(String responseHandlerName) {
155155
+ " runAndLogError(log, \"Exception thrown in exceptionOccurred callback, ignoring\", () "
156156
+ "-> %s.exceptionOccurred(e));%n"
157157
+ " }%n"
158-
+ "})", responseHandlerName);
158+
+ "%s"
159+
+ "})", responseHandlerName, publishMetrics());
159160

160161
}
161162

@@ -177,4 +178,16 @@ default TypeName executeFutureValueType(OperationModel opModel, PoetExtensions p
177178
default TypeName getPojoResponseType(OperationModel opModel, PoetExtensions poetExtensions) {
178179
return poetExtensions.getModelClass(opModel.getReturnType().getReturnType());
179180
}
181+
182+
default String publishMetricsWhenComplete() {
183+
return String.format(".whenComplete((r, e) -> {%n"
184+
+ "%s%n"
185+
+ "})", publishMetrics());
186+
}
187+
188+
default String publishMetrics() {
189+
return "Optional<MetricPublisher> metricPublisher = MetricUtils.resolvePublisher(clientConfiguration, "
190+
+ "requestOverrideConfig);\n"
191+
+ "metricPublisher.ifPresent(p -> p.publish(apiCallMetricCollector.collect()));";
192+
}
180193
}

‎codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/specs/QueryProtocolSpec.java

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import java.util.Optional;
2525
import java.util.concurrent.CompletableFuture;
2626
import javax.lang.model.element.Modifier;
27+
import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
2728
import software.amazon.awssdk.awscore.exception.AwsServiceException;
2829
import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
2930
import software.amazon.awssdk.codegen.model.intermediate.OperationModel;
@@ -154,9 +155,12 @@ public CodeBlock asyncExecutionHandler(IntermediateModel intermediateModel, Oper
154155
"errorResponseHandler",
155156
opModel.getInput().getVariableName(),
156157
opModel.hasStreamingOutput() ? ", asyncResponseTransformer" : "");
157-
158+
builder.addStatement("$T requestOverrideConfig = $L.overrideConfiguration().orElse(null)",
159+
AwsRequestOverrideConfiguration.class, opModel.getInput().getVariableName());
158160
if (opModel.hasStreamingOutput()) {
159161
builder.add("executeFuture$L;", streamingOutputWhenComplete("asyncResponseTransformer"));
162+
} else {
163+
builder.add("executeFuture$L;", publishMetricsWhenComplete());
160164
}
161165
builder.addStatement("return executeFuture");
162166
return builder.build();

‎codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/specs/XmlProtocolSpec.java

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import com.squareup.javapoet.ParameterizedTypeName;
2121
import com.squareup.javapoet.TypeName;
2222
import java.util.Optional;
23+
import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
2324
import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
2425
import software.amazon.awssdk.codegen.model.intermediate.OperationModel;
2526
import software.amazon.awssdk.codegen.poet.PoetExtensions;
@@ -169,8 +170,12 @@ public CodeBlock asyncExecutionHandler(IntermediateModel intermediateModel, Oper
169170
opModel.getInput().getVariableName(),
170171
opModel.hasStreamingOutput() ? ", asyncResponseTransformer" : "");
171172

173+
builder.addStatement("$T requestOverrideConfig = $L.overrideConfiguration().orElse(null)",
174+
AwsRequestOverrideConfiguration.class, opModel.getInput().getVariableName());
172175
if (opModel.hasStreamingOutput()) {
173176
builder.add("executeFuture$L;", streamingOutputWhenComplete("asyncResponseTransformer"));
177+
} else {
178+
builder.add("executeFuture$L;", publishMetricsWhenComplete());
174179
}
175180
builder.addStatement("return executeFuture");
176181
return builder.build();

0 commit comments

Comments
 (0)