Add metric and span support for CBOR serialize/deserialize (#566)

aws · Jan 28, 2025 · 0681fa9 · 0681fa9
1 parent d708d1d
commit 0681fa9
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/...c/main/java/software/amazon/smithy/go/codegen/protocol/DeserializeResponseMiddleware.java b/...c/main/java/software/amazon/smithy/go/codegen/protocol/DeserializeResponseMiddleware.java
@@ -18,6 +18,7 @@
 import static software.amazon.smithy.go.codegen.GoStackStepMiddlewareGenerator.createDeserializeStepMiddleware;
 import static software.amazon.smithy.go.codegen.GoWriter.emptyGoTemplate;
 import static software.amazon.smithy.go.codegen.GoWriter.goTemplate;
+import static software.amazon.smithy.go.codegen.SmithyGoDependency.SMITHY_TRACING;
 import static software.amazon.smithy.go.codegen.integration.ProtocolGenerator.getDeserializeMiddlewareName;
 
 import software.amazon.smithy.go.codegen.GoStdlibTypes;
@@ -61,6 +62,10 @@ public void accept(GoWriter writer) {
 
     private GoWriter.Writable generateHandleDeserialize() {
         return goTemplate("""
+                _, span := $startSpan:T(ctx, "OperationDeserializer")
+                endTimer := startMetricTimer(ctx, "client.call.deserialization_duration")
+                defer endTimer()
+                defer span.End()
                 out, metadata, err = next.HandleDeserialize(ctx, in)
                 if err != nil {
                     return out, metadata, err
@@ -76,6 +81,7 @@ private GoWriter.Writable generateHandleDeserialize() {
                 return out, metadata, nil
                 """,
                 MapUtils.of(
+                        "startSpan", SMITHY_TRACING.func("StartSpan"),
                         "response", generator.getApplicationProtocol().getResponseType(),
                         "deserialize", generateDeserialize(),
                         "errorf", GoStdlibTypes.Fmt.Errorf

diff --git a/.../src/main/java/software/amazon/smithy/go/codegen/protocol/SerializeRequestMiddleware.java b/.../src/main/java/software/amazon/smithy/go/codegen/protocol/SerializeRequestMiddleware.java
@@ -18,6 +18,7 @@
 import static software.amazon.smithy.go.codegen.GoStackStepMiddlewareGenerator.createSerializeStepMiddleware;
 import static software.amazon.smithy.go.codegen.GoWriter.emptyGoTemplate;
 import static software.amazon.smithy.go.codegen.GoWriter.goTemplate;
+import static software.amazon.smithy.go.codegen.SmithyGoDependency.SMITHY_TRACING;
 
 import software.amazon.smithy.go.codegen.GoStdlibTypes;
 import software.amazon.smithy.go.codegen.GoWriter;
@@ -53,7 +54,6 @@ public void accept(GoWriter writer) {
         var name = ProtocolGenerator.getSerializeMiddlewareName(operation.getId(), ctx.getService(),
                 generator.getProtocolName());
         var middleware = createSerializeStepMiddleware(name, ProtocolUtils.OPERATION_SERIALIZER_MIDDLEWARE_ID);
-
         writer.write(middleware.asWritable(generateHandleSerialize(), emptyGoTemplate()));
     }
 
@@ -63,6 +63,10 @@ public void accept(GoWriter writer) {
 
     private GoWriter.Writable generateHandleSerialize() {
         return goTemplate("""
+                _, span := $startSpan:T(ctx, "OperationSerializer")
+                endTimer := startMetricTimer(ctx, "client.call.serialization_duration")
+                defer endTimer()
+                defer span.End()
                 input, ok := in.Parameters.($input:P)
                 if !ok {
                     return out, metadata, $errorf:T("unexpected input type %T", in.Parameters)
@@ -81,6 +85,7 @@ private GoWriter.Writable generateHandleSerialize() {
                 return next.HandleSerialize(ctx, in)
                 """,
                 MapUtils.of(
+                        "startSpan", SMITHY_TRACING.func("StartSpan"),
                         "input", ctx.getSymbolProvider().toSymbol(input),
                         "request", generator.getApplicationProtocol().getRequestType(),
                         "route", generateRouteRequest(),