feat: support for env variable in tracing capture modes (#249)

Author: pankajagrawal16

docs/content/core/tracing.mdx

@@ -61,24 +61,52 @@ public class App implements RequestHandler<APIGatewayProxyRequestEvent, APIGatew
 }
 ```
 
-By default this annotation will automatically record method responses and exceptions.
 If you want to customize segment name that appears in traces, use:
-`@Tracing(segmentName="yourCustomName")`
+
+```java:title=CustomSegmentName.java
+public class App implements RequestHandler<APIGatewayProxyRequestEvent, APIGatewayProxyResponseEvent> {
+
+    @Tracing(segmentName="yourCustomName")
+    public APIGatewayProxyResponseEvent handleRequest(APIGatewayProxyRequestEvent input, Context context) {
+    ...
+    }
+```
+
+By default, this annotation will automatically record method responses and exceptions. You can change the default behavior by setting
+the environment variables `TRACING_CAPTURE_RESPONSE` and `TRACING_CAPTURE_ERROR` as needed. Optionally, you can override behavior by
+different supported `captureMode` to record response, exception or both.
 
 <Note type="warning">
   <strong>Returning sensitive information from your Lambda handler or functions, where Tracer is used?</strong>
   <br/><br/>
-  You can disable Tracer from capturing their responses and exception as tracing metadata with <strong><code>captureResponse=false</code></strong> and <strong><code>captureError=false</code></strong>
+  You can disable annotation from capturing their responses and exception as tracing metadata with <strong><code>captureMode=DISABLED </code></strong>
+   or globally by setting environment variables <strong><code>TRACING_CAPTURE_RESPONSE</code></strong> and <strong><code>TRACING_CAPTURE_ERROR</code></strong> to <strong><code>false</code></strong>.
 </Note><br/>
 
 ```java:title=HandlerWithoutCapturingResponseOrError.java
 public class App implements RequestHandler<APIGatewayProxyRequestEvent, APIGatewayProxyResponseEvent> {
 
-    @Tracing(captureError = false, captureResponse = false)
+    @Tracing(captureMode=CaptureMode.DISABLED)
     public APIGatewayProxyResponseEvent handleRequest(APIGatewayProxyRequestEvent input, Context context) {
     ...
     }
 ```
+Globally:
+
+```yaml:title=template.yaml
+Resources:
+    HelloWorldFunction:
+        Type: AWS::Serverless::Function
+        Properties:
+        ...
+        Runtime: java8
+
+        Tracing: Active
+        Environment:
+            Variables:
+                TRACING_CAPTURE_RESPONSE: false # highlight-line
+                TRACING_CAPTURE_ERROR: false # highlight-line
+```
 
 ### Annotations
 

docs/content/index.mdx

@@ -123,7 +123,9 @@ Environment variable | Description | Utility
 **POWERTOOLS_SERVICE_NAME** | Sets service name used for tracing namespace, metrics dimension and structured logging | All
 **POWERTOOLS_METRICS_NAMESPACE** | Sets namespace used for metrics | [Metrics](./core/metrics)
 **POWERTOOLS_LOGGER_SAMPLE_RATE** | Debug log sampling | [Logging](./core/logging)
-**LOG_LEVEL** | Sets logging level | [Logging](./core/logger)
+**LOG_LEVEL** | Sets logging level | [Logging](./core/logging)
+**TRACING_CAPTURE_RESPONSE** | Enables/Disables tracing mode to capture method response | [Tracing](./core/tracing)
+**TRACING_CAPTURE_ERROR** | Enables/Disables tracing mode to capture method error | [Tracing](./core/tracing)
 
 ## Tenets
 

powertools-tracing/pom.xml

@@ -100,6 +100,11 @@
             <artifactId>mockito-core</artifactId>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>org.mockito</groupId>
+            <artifactId>mockito-inline</artifactId>
+            <scope>test</scope>
+        </dependency>
         <dependency>
             <groupId>org.aspectj</groupId>
             <artifactId>aspectjweaver</artifactId>

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

@@ -0,0 +1,31 @@
+package software.amazon.lambda.powertools.tracing;
+
+public enum CaptureMode {
+    /**
+     * Enables annotation to capture only response. If this mode is explicitly overridden
+     * on {@link Tracing} annotation, it will override value of environment variable TRACING_CAPTURE_RESPONSE
+     */
+    RESPONSE,
+    /**
+     * Enabled annotation to capture only error from the method. If this mode is explicitly overridden
+     * on {@link Tracing} annotation, it will override value of environment variable TRACING_CAPTURE_ERROR
+     */
+    ERROR,
+    /**
+     * Enabled annotation to capture both response error from the method. If this mode is explicitly overridden
+     * on {@link Tracing} annotation, it will override value of environment variables TRACING_CAPTURE_RESPONSE
+     * and TRACING_CAPTURE_ERROR
+     */
+    RESPONSE_AND_ERROR,
+    /**
+     * Disables annotation to capture both response and error from the method. If this mode is explicitly overridden
+     * on {@link Tracing} annotation, it will override values of environment variable TRACING_CAPTURE_RESPONSE
+     * and TRACING_CAPTURE_ERROR
+     */
+    DISABLED,
+    /**
+     * Enables/Disables annotation to capture response and error from the method based on the value of
+     * environment variable TRACING_CAPTURE_RESPONSE and TRACING_CAPTURE_ERROR
+     */
+    ENVIRONMENT_VAR
+}

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

@@ -48,7 +48,20 @@
 @Target(ElementType.METHOD)
 public @interface Tracing {
     String namespace() default "";
+    /**
+     * @deprecated As of release 1.2.0, replaced by captureMode()
+     * in order to support different modes and support via
+     * environment variables
+     */
+    @Deprecated
     boolean captureResponse() default true;
+    /**
+     * @deprecated As of release 1.2.0, replaced by captureMode()
+     * in order to support different modes and support via
+     * environment variables
+     */
+    @Deprecated
     boolean captureError() default true;
     String segmentName() default "";
+    CaptureMode captureMode() default CaptureMode.ENVIRONMENT_VAR;
 }

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

@@ -13,9 +13,10 @@
  */
 package software.amazon.lambda.powertools.tracing.internal;
 
+import java.util.function.Supplier;
+
 import com.amazonaws.xray.AWSXRay;
 import com.amazonaws.xray.entities.Subsegment;
-import java.util.function.Supplier;
 import org.aspectj.lang.ProceedingJoinPoint;
 import org.aspectj.lang.annotation.Around;
 import org.aspectj.lang.annotation.Aspect;
@@ -32,7 +33,6 @@
 
 @Aspect
 public final class LambdaTracingAspect {
-
     @SuppressWarnings({"EmptyMethod"})
     @Pointcut("@annotation(tracing)")
     public void callAt(Tracing tracing) {
@@ -52,16 +52,19 @@ public Object around(ProceedingJoinPoint pjp,
             segment.putAnnotation("ColdStart", isColdStart());
         }
 
+        boolean captureResponse = captureResponse(tracing);
+        boolean captureError = captureError(tracing);
+
         try {
             Object methodReturn = pjp.proceed(proceedArgs);
-            if (tracing.captureResponse()) {
+            if (captureResponse) {
                 segment.putMetadata(namespace(tracing), pjp.getSignature().getName() + " response", methodReturn);
             }
 
             coldStartDone();
             return methodReturn;
         } catch (Exception e) {
-            if (tracing.captureError()) {
+            if (captureError) {
                 segment.putMetadata(namespace(tracing), pjp.getSignature().getName() + " error", e);
             }
             throw e;
@@ -72,6 +75,34 @@ public Object around(ProceedingJoinPoint pjp,
         }
     }
 
+    private boolean captureResponse(Tracing powerToolsTracing) {
+        switch (powerToolsTracing.captureMode()) {
+            case ENVIRONMENT_VAR:
+                Boolean captureResponse = environmentVariable("TRACING_CAPTURE_RESPONSE");
+                return null != captureResponse ? captureResponse : powerToolsTracing.captureResponse();
+            case RESPONSE:
+            case RESPONSE_AND_ERROR:
+                return true;
+            case DISABLED:
+            default:
+                return false;
+        }
+    }
+
+    private boolean captureError(Tracing powerToolsTracing) {
+        switch (powerToolsTracing.captureMode()) {
+            case ENVIRONMENT_VAR:
+                Boolean captureError = environmentVariable("TRACING_CAPTURE_ERROR");
+                return null != captureError ? captureError : powerToolsTracing.captureError();
+            case ERROR:
+            case RESPONSE_AND_ERROR:
+                return true;
+            case DISABLED:
+            default:
+                return false;
+        }
+    }
+
     private String customSegmentNameOrDefault(Tracing powerToolsTracing, Supplier<String> defaultSegmentName) {
         String segmentName = powerToolsTracing.segmentName();
         return segmentName.isEmpty() ? defaultSegmentName.get() : segmentName;
@@ -85,4 +116,9 @@ private boolean placedOnHandlerMethod(ProceedingJoinPoint pjp) {
         return isHandlerMethod(pjp)
                 && (placedOnRequestHandler(pjp) || placedOnStreamHandler(pjp));
     }
+
+    private Boolean environmentVariable(String tracing_capture_response) {
+        return null != SystemWrapper.getenv(tracing_capture_response)
+                ? Boolean.valueOf(SystemWrapper.getenv(tracing_capture_response)) : null;
+    }
 }

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

@@ -0,0 +1,10 @@
+package software.amazon.lambda.powertools.tracing.internal;
+
+public class SystemWrapper {
+    public SystemWrapper() {
+    }
+
+    public static String getenv(String name) {
+        return System.getenv(name);
+    }
+}

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

@@ -0,0 +1,29 @@
+/*
+ * 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 com.amazonaws.services.lambda.runtime.Context;
+import com.amazonaws.services.lambda.runtime.RequestHandler;
+import software.amazon.lambda.powertools.tracing.Tracing;
+
+import static software.amazon.lambda.powertools.tracing.CaptureMode.RESPONSE_AND_ERROR;
+
+public class PowerTracerToolEnabledExplicitlyForResponseAndError implements RequestHandler<Object, Object> {
+
+    @Override
+    @Tracing(namespace = "lambdaHandler", captureMode = RESPONSE_AND_ERROR)
+    public Object handleRequest(Object input, Context context) {
+        return null;
+    }
+}

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

@@ -0,0 +1,29 @@
+/*
+ * 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 com.amazonaws.services.lambda.runtime.Context;
+import com.amazonaws.services.lambda.runtime.RequestHandler;
+import software.amazon.lambda.powertools.tracing.Tracing;
+
+import static software.amazon.lambda.powertools.tracing.CaptureMode.ERROR;
+
+public class PowerTracerToolEnabledForError implements RequestHandler<Object, Object> {
+
+    @Override
+    @Tracing(namespace = "lambdaHandler", captureMode = ERROR)
+    public Object handleRequest(Object input, Context context) {
+        throw new RuntimeException("I am failing!");
+    }
+}

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

@@ -0,0 +1,29 @@
+/*
+ * 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 com.amazonaws.services.lambda.runtime.Context;
+import com.amazonaws.services.lambda.runtime.RequestHandler;
+import software.amazon.lambda.powertools.tracing.Tracing;
+
+import static software.amazon.lambda.powertools.tracing.CaptureMode.RESPONSE;
+
+public class PowerTracerToolEnabledForResponse implements RequestHandler<Object, Object> {
+
+    @Override
+    @Tracing(namespace = "lambdaHandler", captureMode = RESPONSE)
+    public Object handleRequest(Object input, Context context) {
+        return null;
+    }
+}

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

@@ -13,17 +13,19 @@
  */
 package software.amazon.lambda.powertools.tracing.handlers;
 
+import java.io.InputStream;
+import java.io.OutputStream;
+
 import com.amazonaws.services.lambda.runtime.Context;
 import com.amazonaws.services.lambda.runtime.RequestStreamHandler;
 import software.amazon.lambda.powertools.tracing.Tracing;
 
-import java.io.InputStream;
-import java.io.OutputStream;
+import static software.amazon.lambda.powertools.tracing.CaptureMode.DISABLED;
 
 public class PowerTracerToolEnabledForStreamWithNoMetaData implements RequestStreamHandler {
 
     @Override
-    @Tracing(captureResponse = false, captureError = false)
+    @Tracing(captureMode = DISABLED)
     public void handleRequest(InputStream input, OutputStream output, Context context) {
 
     }

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

@@ -17,10 +17,12 @@
 import com.amazonaws.services.lambda.runtime.RequestHandler;
 import software.amazon.lambda.powertools.tracing.Tracing;
 
+import static software.amazon.lambda.powertools.tracing.CaptureMode.DISABLED;
+
 public class PowerTracerToolEnabledWithNoMetaData implements RequestHandler<Object, Object> {
 
     @Override
-    @Tracing(captureResponse = false, captureError = false)
+    @Tracing(captureMode = DISABLED)
     public Object handleRequest(Object input, Context context) {
         return null;
     }

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

@@ -0,0 +1,29 @@
+/*
+ * 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 com.amazonaws.services.lambda.runtime.Context;
+import com.amazonaws.services.lambda.runtime.RequestHandler;
+import software.amazon.lambda.powertools.tracing.Tracing;
+
+import static software.amazon.lambda.powertools.tracing.CaptureMode.DISABLED;
+
+public class PowerTracerToolEnabledWithNoMetaDataDeprecated implements RequestHandler<Object, Object> {
+
+    @Override
+    @Tracing(captureResponse = false, captureError = false)
+    public Object handleRequest(Object input, Context context) {
+        return null;
+    }
+}