feat(tracing): ability to override object mapper used for serialization

Author: pankajagrawal16

powertools-tracing/src/main/java/software/amazon/lambda/powertools/tracing/TracingUtils.java

@@ -13,11 +13,11 @@
  */
 package software.amazon.lambda.powertools.tracing;
 
+import java.util.function.Consumer;
 import com.amazonaws.xray.AWSXRay;
 import com.amazonaws.xray.entities.Entity;
 import com.amazonaws.xray.entities.Subsegment;
-
-import java.util.function.Consumer;
+import com.fasterxml.jackson.databind.ObjectMapper;
 
 import static software.amazon.lambda.powertools.core.internal.LambdaHandlerProcessor.serviceName;
 
@@ -27,6 +27,7 @@
  *
  */
 public final class TracingUtils {
+    private static ObjectMapper objectMapper;
 
     /**
      * Put an annotation to the current subsegment with a String value.
@@ -155,4 +156,18 @@ public static void withSubsegment(String namespace, String name, Consumer<Subseg
             AWSXRay.endSubsegment();
         }
     }
+
+    /**
+     * Sets the instance of ObjectMapper object which is used for serialising object response when
+     * {@code @Tracing(captureMode=CaptureMode.RESPONSE)}.
+     *
+     * @param objectMapper Custom implementation of object mapper to be used for serializing response
+     */
+    public static void defaultObjectMapper(ObjectMapper objectMapper) {
+        TracingUtils.objectMapper = objectMapper;
+    }
+
+    public static ObjectMapper objectMapper() {
+        return objectMapper;
+    }
 }

powertools-tracing/src/main/java/software/amazon/lambda/powertools/tracing/internal/LambdaTracingAspect.java

@@ -14,7 +14,6 @@
 package software.amazon.lambda.powertools.tracing.internal;
 
 import java.util.function.Supplier;
-
 import com.amazonaws.xray.AWSXRay;
 import com.amazonaws.xray.entities.Subsegment;
 import org.aspectj.lang.ProceedingJoinPoint;
@@ -30,6 +29,7 @@
 import static software.amazon.lambda.powertools.core.internal.LambdaHandlerProcessor.placedOnRequestHandler;
 import static software.amazon.lambda.powertools.core.internal.LambdaHandlerProcessor.placedOnStreamHandler;
 import static software.amazon.lambda.powertools.core.internal.LambdaHandlerProcessor.serviceName;
+import static software.amazon.lambda.powertools.tracing.TracingUtils.objectMapper;
 
 @Aspect
 public final class LambdaTracingAspect {
@@ -59,7 +59,7 @@ public Object around(ProceedingJoinPoint pjp,
         try {
             Object methodReturn = pjp.proceed(proceedArgs);
             if (captureResponse) {
-                segment.putMetadata(namespace(tracing), pjp.getSignature().getName() + " response", methodReturn);
+                segment.putMetadata(namespace(tracing), pjp.getSignature().getName() + " response", null != objectMapper() ? objectMapper().writeValueAsString(methodReturn): methodReturn);
             }
 
             if (placedOnHandlerMethod(pjp)) {
@@ -69,7 +69,7 @@ public Object around(ProceedingJoinPoint pjp,
             return methodReturn;
         } catch (Exception e) {
             if (captureError) {
-                segment.putMetadata(namespace(tracing), pjp.getSignature().getName() + " error", e);
+                segment.putMetadata(namespace(tracing), pjp.getSignature().getName() + " error", null != objectMapper() ? objectMapper().writeValueAsString(e) : e);
             }
             throw e;
         } finally {

powertools-tracing/src/test/java/software/amazon/lambda/powertools/tracing/handlers/PowerTracerToolEnabledForResponseWithCustomMapper.java

@@ -0,0 +1,88 @@
+/*
+ * Copyright 2020 Amazon.com, Inc. or its affiliates.
+ * Licensed under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License 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.lambda.powertools.tracing.handlers;
+
+import java.io.IOException;
+import com.amazonaws.services.lambda.runtime.Context;
+import com.amazonaws.services.lambda.runtime.RequestHandler;
+import com.fasterxml.jackson.core.JsonGenerator;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.SerializerProvider;
+import com.fasterxml.jackson.databind.module.SimpleModule;
+import com.fasterxml.jackson.databind.ser.std.StdSerializer;
+import software.amazon.lambda.powertools.tracing.Tracing;
+import software.amazon.lambda.powertools.tracing.TracingUtils;
+
+import static software.amazon.lambda.powertools.tracing.CaptureMode.RESPONSE;
+
+public class PowerTracerToolEnabledForResponseWithCustomMapper implements RequestHandler<Object, Object> {
+    static {
+        ObjectMapper objectMapper = new ObjectMapper();
+        SimpleModule simpleModule = new SimpleModule();
+        simpleModule.addSerializer(ChildClass.class, new ChildSerializer());
+        objectMapper.registerModule(simpleModule);
+
+        TracingUtils.defaultObjectMapper(objectMapper);
+    }
+    @Override
+    @Tracing(namespace = "lambdaHandler", captureMode = RESPONSE)
+    public Object handleRequest(Object input, Context context) {
+        ParentClass parentClass = new ParentClass("parent");
+        ChildClass childClass = new ChildClass("child", parentClass);
+        parentClass.setC(childClass);
+        return parentClass;
+    }
+
+    public class ParentClass {
+        public String name;
+        public ChildClass c;
+
+        public ParentClass(String name) {
+            this.name = name;
+        }
+
+        public void setC(ChildClass c) {
+            this.c = c;
+        }
+    }
+
+    public class ChildClass {
+        public String name;
+        public ParentClass p;
+
+        public ChildClass(String name, ParentClass p) {
+            this.name = name;
+            this.p = p;
+        }
+    }
+
+    public static class ChildSerializer extends StdSerializer<ChildClass> {
+
+        public ChildSerializer() {
+            this(null);
+        }
+
+        public ChildSerializer(Class<ChildClass> t) {
+            super(t);
+        }
+
+        @Override
+        public void serialize(ChildClass value, JsonGenerator jgen, SerializerProvider provider) throws IOException {
+            jgen.writeStartObject();
+            jgen.writeStringField("name", value.name);
+            jgen.writeStringField("p", value.p.name);
+            jgen.writeEndObject();
+        }
+    }
+}

powertools-tracing/src/test/java/software/amazon/lambda/powertools/tracing/internal/LambdaTracingAspectTest.java

@@ -16,7 +16,6 @@
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
-
 import com.amazonaws.services.lambda.runtime.Context;
 import com.amazonaws.services.lambda.runtime.RequestHandler;
 import com.amazonaws.services.lambda.runtime.RequestStreamHandler;
@@ -34,6 +33,7 @@
 import software.amazon.lambda.powertools.tracing.handlers.PowerTracerToolEnabledExplicitlyForResponseAndError;
 import software.amazon.lambda.powertools.tracing.handlers.PowerTracerToolEnabledForError;
 import software.amazon.lambda.powertools.tracing.handlers.PowerTracerToolEnabledForResponse;
+import software.amazon.lambda.powertools.tracing.handlers.PowerTracerToolEnabledForResponseWithCustomMapper;
 import software.amazon.lambda.powertools.tracing.handlers.PowerTracerToolEnabledForStream;
 import software.amazon.lambda.powertools.tracing.handlers.PowerTracerToolEnabledForStreamWithNoMetaData;
 import software.amazon.lambda.powertools.tracing.handlers.PowerTracerToolEnabledWithException;
@@ -43,6 +43,7 @@
 
 import static org.apache.commons.lang3.reflect.FieldUtils.writeStaticField;
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatNoException;
 import static org.assertj.core.api.Assertions.catchThrowable;
 import static org.mockito.Mockito.mockStatic;
 import static org.mockito.Mockito.when;
@@ -277,6 +278,28 @@ void shouldCaptureTracesIfExplicitlyEnabledAndEnvironmentVariableIsDisabled() {
         }
     }
 
+    @Test
+    void shouldCaptureTracesForSelfReferencingReturnTypesViaCustomMapper() {
+        requestHandler = new PowerTracerToolEnabledForResponseWithCustomMapper();
+
+        requestHandler.handleRequest(new Object(), context);
+
+        assertThat(AWSXRay.getTraceEntity().getSubsegmentsCopy())
+                .hasSize(1)
+                .allSatisfy(subsegment -> {
+                    assertThat(subsegment.getMetadata())
+                            .hasSize(1)
+                            .containsKey("lambdaHandler");
+
+                    assertThat(subsegment.getMetadata().get("lambdaHandler"))
+                            .hasFieldOrPropertyWithValue("handleRequest response", "{\"name\":\"parent\",\"c\":{\"name\":\"child\",\"p\":\"parent\"}}");
+                });
+
+        assertThatNoException().isThrownBy(AWSXRay::endSegment);
+
+        AWSXRay.beginSegment(LambdaTracingAspectTest.class.getName());
+    }
+
     @Test
     void shouldCaptureTracesIfExplicitlyEnabledBothAndEnvironmentVariableIsDisabled() {
         try (MockedStatic<SystemWrapper> mocked = mockStatic(SystemWrapper.class)) {

spotbugs-exclude.xml

@@ -71,6 +71,10 @@
                 <Class name="software.amazon.lambda.powertools.logging.LoggingUtils"/>
                 <Method name="defaultObjectMapper"/>
             </And>
+            <And>
+                <Class name="software.amazon.lambda.powertools.tracing.TracingUtils"/>
+                <Method name="defaultObjectMapper"/>
+            </And>
             <And>
                 <Class name="software.amazon.lambda.powertools.sqs.SqsUtils"/>
                 <Method name="overrideSqsClient"/>
@@ -88,6 +92,10 @@
                 <Class name="software.amazon.lambda.powertools.logging.LoggingUtils"/>
                 <Method name="objectMapper"/>
             </And>
+            <And>
+                <Class name="software.amazon.lambda.powertools.tracing.TracingUtils"/>
+                <Method name="objectMapper"/>
+            </And>
             <And>
                 <Class name="software.amazon.lambda.powertools.sqs.SqsUtils"/>
                 <Method name="objectMapper"/>