Support log sampling to DEBUG level based on config (#52)

Author: pankajagrawal16

docs/content/core/logging.mdx

@@ -63,6 +63,7 @@ Key | Type | Example | Description
 **level** | String | "INFO" | Logging level
 **coldStart** | Boolean | true| ColdStart value.
 **service** | String | "payment" | Service name defined. "service_undefined" will be used if unknown
+**samplingRate** | int |  0.1 | Debug logging sampling rate in percentage e.g. 10% in this case
 **message** | String |  "Collecting payment" | Log statement value. Unserializable JSON values will be casted to string
 **functionName**| String | "example-powertools-HelloWorldFunction-1P1Z6B39FLU73"
 **functionVersion**| String | "12"
@@ -161,4 +162,48 @@ public class App implements RequestHandler<APIGatewayProxyRequestEvent, APIGatew
         ...
     }
 }
+```
+
+## Sampling debug logs
+
+You can dynamically set a percentage of your logs to **DEBUG** level via env var `POWERTOOLS_LOGGER_SAMPLE_RATE` or
+via `samplingRate` attribute on annotation.
+
+Configuration on environment variable is given precedence over sampling rate configuration on annotation, provided it's in valid value range.
+
+```java:title=App.java
+package helloworld;
+
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import software.amazon.lambda.logging.PowertoolsLogger;
+import software.amazon.lambda.logging.PowertoolsLogging;
+...
+
+/**
+ * Handler for requests to Lambda function.
+ */
+public class App implements RequestHandler<APIGatewayProxyRequestEvent, APIGatewayProxyResponseEvent> {
+
+    Logger log = LogManager.getLogger();
+
+    @PowertoolsLogging(samplingRate = 0.5)
+    public APIGatewayProxyResponseEvent handleRequest(final APIGatewayProxyRequestEvent input, final Context context) {
+     ...
+    }
+}
+```
+
+You can also override sampling by setting `POWERTOOLS_LOGGER_SAMPLE_RATE` env var  - Here is an example using AWS Serverless Application Model (SAM)
+
+```yaml:title=template.yaml
+Resources:
+    HelloWorldFunction:
+        Type: AWS::Serverless::Function
+        Properties:
+        ...
+        Runtime: java8
+        Environment:
+            Variables:
+                POWERTOOLS_LOGGER_SAMPLE_RATE: 0.5 # highlight-line
 ```

example/HelloWorldFunction/src/main/java/helloworld/App.java

@@ -31,7 +31,7 @@ public class App implements RequestHandler<APIGatewayProxyRequestEvent, APIGatew
 
     Logger log = LogManager.getLogger();
 
-    @PowertoolsLogging(logEvent = true)
+    @PowertoolsLogging(logEvent = true, samplingRate = 0.7)
     @PowertoolsTracing(captureError = false, captureResponse = false)
     public APIGatewayProxyResponseEvent handleRequest(final APIGatewayProxyRequestEvent input, final Context context) {
         Map<String, String> headers = new HashMap<>();

example/template.yaml

@@ -39,7 +39,7 @@ Resources:
       Tracing: Active
       Environment:
         Variables:
-          PARAM1: VALUE
+          POWERTOOLS_LOGGER_SAMPLE_RATE: 0.7
       Events:
         HelloWorld:
           Type: Api

powertools-logging/src/main/java/software/amazon/lambda/powertools/logging/PowertoolsLogging.java

@@ -44,6 +44,8 @@
  * <ul>
  *     <li>coldStart - True if this is the first invocation of this Lambda execution environment; else False</li>
  *     <li>service - The value of the 'POWER_TOOLS_SERVICE_NAME' environment variable or 'service_undefined'</li>
+ *     <li>samplingRate - The value of the 'POWERTOOLS_LOGGER_SAMPLE_RATE' environment variable or value of samplingRate field or 0.
+ *     Valid value is from 0.0 to 1.0. Value outside this range is silently ignored.</li>
  * </ul>
  *
  * <p>These keys and values will be joined with the existing Log4J log event and written as JSON.</p>
@@ -53,11 +55,16 @@
  * <p>By default {@code PowertoolsLogging} will not log the event which has trigger the invoke of the Lambda function.
  * This can be enabled using {@code @PowertoolsLogging(logEvent = true)}.</p>
  *
+ * <p>By default {@code PowertoolsLogging} all debug loggs will follow log4j2 configuration unless configured via
+ * POWERTOOLS_LOGGER_SAMPLE_RATE environment variable {@code @PowertoolsLogging(samplingRate = <0.0-1.0>)}.</p>
+ *
  * <p>To append additional keys to each log entry you can use {@link PowertoolsLogger#appendKey(String, String)}</p>
  */
 @Retention(RetentionPolicy.RUNTIME)
 @Target(ElementType.METHOD)
 public @interface PowertoolsLogging {
 
     boolean logEvent() default false;
+
+    double samplingRate() default 0;
 }

powertools-logging/src/main/java/software/amazon/lambda/powertools/logging/internal/LambdaLoggingAspect.java

@@ -13,12 +13,21 @@
  */
 package software.amazon.lambda.powertools.logging.internal;
 
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.InputStreamReader;
+import java.io.OutputStreamWriter;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Random;
+
 import com.amazonaws.services.lambda.runtime.Context;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import org.apache.logging.log4j.Level;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
-import org.apache.logging.log4j.ThreadContext;
 import org.apache.logging.log4j.core.LoggerContext;
 import org.apache.logging.log4j.core.config.Configurator;
 import org.apache.logging.log4j.core.util.IOUtils;
@@ -28,15 +37,6 @@
 import org.aspectj.lang.annotation.Pointcut;
 import software.amazon.lambda.powertools.logging.PowertoolsLogging;
 
-import java.io.ByteArrayInputStream;
-import java.io.ByteArrayOutputStream;
-import java.io.IOException;
-import java.io.InputStream;
-import java.io.InputStreamReader;
-import java.io.OutputStreamWriter;
-import java.util.Map;
-import java.util.Optional;
-
 import static java.util.Optional.empty;
 import static java.util.Optional.of;
 import static software.amazon.lambda.powertools.core.internal.LambdaHandlerProcessor.coldStartDone;
@@ -45,22 +45,26 @@
 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.logging.PowertoolsLogger.appendKey;
+import static software.amazon.lambda.powertools.logging.PowertoolsLogger.appendKeys;
 
 @Aspect
 public final class LambdaLoggingAspect {
-    private static final ObjectMapper mapper = new ObjectMapper();
-    private static String LOG_LEVEL = System.getenv("LOG_LEVEL");
+    private static final Logger LOG = LogManager.getLogger(LambdaLoggingAspect.class);
+    private static final ObjectMapper MAPPER = new ObjectMapper();
+    private static final Random SAMPLER = new Random();
 
-    static {
-        resetLogLevels();
-    }
+    private static final String LOG_LEVEL = System.getenv("LOG_LEVEL");
+    private static final String SAMPLING_RATE = System.getenv("POWERTOOLS_LOGGER_SAMPLE_RATE");
 
-    private static void resetLogLevels() {
-        if (LOG_LEVEL != null) {
-            LoggerContext ctx = (LoggerContext) LogManager.getContext(false);
-            Configurator.setAllLevels(LogManager.getRootLogger().getName(), Level.getLevel(LOG_LEVEL));
-            ctx.updateLoggers();
+    private static Level LEVEL_AT_INITIALISATION;
+
+    static {
+        if (null != LOG_LEVEL) {
+            resetLogLevels(Level.getLevel(LOG_LEVEL));
         }
+
+        LEVEL_AT_INITIALISATION = LOG.getLevel();
     }
 
     @SuppressWarnings({"EmptyMethod"})
@@ -73,11 +77,13 @@ public Object around(ProceedingJoinPoint pjp,
                          PowertoolsLogging powertoolsLogging) throws Throwable {
         Object[] proceedArgs = pjp.getArgs();
 
+        setLogLevelBasedOnSamplingRate(pjp, powertoolsLogging);
+
         extractContext(pjp)
                 .ifPresent(context -> {
-                    ThreadContext.putAll(DefaultLambdaFields.values(context));
-                    ThreadContext.put("coldStart", null == isColdStart() ? "true" : "false");
-                    ThreadContext.put("service", serviceName());
+                    appendKeys(DefaultLambdaFields.values(context));
+                    appendKey("coldStart", null == isColdStart() ? "true" : "false");
+                    appendKey("service", serviceName());
                 });
 
 
@@ -91,7 +97,49 @@ public Object around(ProceedingJoinPoint pjp,
         return proceed;
     }
 
-    private Optional<Context> extractContext(ProceedingJoinPoint pjp) {
+    private static void resetLogLevels(Level logLevel) {
+        LoggerContext ctx = (LoggerContext) LogManager.getContext(false);
+        Configurator.setAllLevels(LogManager.getRootLogger().getName(), logLevel);
+        ctx.updateLoggers();
+    }
+
+    private void setLogLevelBasedOnSamplingRate(final ProceedingJoinPoint pjp,
+                                                final PowertoolsLogging powertoolsLogging) {
+        if (isHandlerMethod(pjp)) {
+            float sample = SAMPLER.nextFloat();
+            double samplingRate = samplingRate(powertoolsLogging);
+
+            if (samplingRate < 0 || samplingRate > 1) {
+                LOG.debug("Skipping sampling rate configuration because of invalid value. Sampling rate: {}", samplingRate);
+                return;
+            }
+
+            appendKey("samplingRate", String.valueOf(samplingRate));
+
+            if (samplingRate > sample) {
+                resetLogLevels(Level.DEBUG);
+
+                LOG.debug("Changed log level to DEBUG based on Sampling configuration. " +
+                        "Sampling Rate: {}, Sampler Value: {}.", samplingRate, sample);
+            } else if (LEVEL_AT_INITIALISATION != LOG.getLevel()) {
+                resetLogLevels(LEVEL_AT_INITIALISATION);
+            }
+        }
+    }
+
+    private double samplingRate(final PowertoolsLogging powertoolsLogging) {
+        if (null != SAMPLING_RATE) {
+            try {
+                return Double.parseDouble(SAMPLING_RATE);
+            } catch (NumberFormatException e) {
+                LOG.debug("Skipping sampling rate on environment variable configuration because of invalid " +
+                        "value. Sampling rate: {}", SAMPLING_RATE);
+            }
+        }
+        return powertoolsLogging.samplingRate();
+    }
+
+    private Optional<Context> extractContext(final ProceedingJoinPoint pjp) {
 
         if (isHandlerMethod(pjp)) {
             if (placedOnRequestHandler(pjp)) {
@@ -106,7 +154,7 @@ private Optional<Context> extractContext(ProceedingJoinPoint pjp) {
         return empty();
     }
 
-    private Object[] logEvent(ProceedingJoinPoint pjp) {
+    private Object[] logEvent(final ProceedingJoinPoint pjp) {
         Object[] args = pjp.getArgs();
 
         if (isHandlerMethod(pjp)) {
@@ -123,7 +171,7 @@ private Object[] logEvent(ProceedingJoinPoint pjp) {
         return args;
     }
 
-    private Object[] logFromInputStream(ProceedingJoinPoint pjp) {
+    private Object[] logFromInputStream(final ProceedingJoinPoint pjp) {
         Object[] args = pjp.getArgs();
 
         try (ByteArrayOutputStream out = new ByteArrayOutputStream();
@@ -136,7 +184,7 @@ private Object[] logFromInputStream(ProceedingJoinPoint pjp) {
             args[0] = new ByteArrayInputStream(bytes);
 
             Logger log = logger(pjp);
-            log.info(mapper.readValue(bytes, Map.class));
+            log.info(MAPPER.readValue(bytes, Map.class));
 
         } catch (IOException e) {
             Logger log = logger(pjp);
@@ -146,7 +194,7 @@ private Object[] logFromInputStream(ProceedingJoinPoint pjp) {
         return args;
     }
 
-    private Logger logger(ProceedingJoinPoint pjp) {
+    private Logger logger(final ProceedingJoinPoint pjp) {
         return LogManager.getLogger(pjp.getSignature().getDeclaringType());
     }
 }

powertools-logging/src/test/java/org/apache/logging/log4j/core/layout/LambdaJsonLayoutTest.java

@@ -26,10 +26,12 @@
 import com.amazonaws.services.lambda.runtime.RequestHandler;
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import org.apache.logging.log4j.Level;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.mockito.Mock;
 import software.amazon.lambda.powertools.logging.handlers.PowerLogToolEnabled;
+import software.amazon.lambda.powertools.logging.handlers.PowerLogToolSamplingEnabled;
 import software.amazon.lambda.powertools.logging.internal.LambdaLoggingAspect;
 
 import static java.util.Collections.emptyMap;
@@ -41,7 +43,7 @@
 
 class LambdaJsonLayoutTest {
 
-    private final RequestHandler<Object, Object> handler = new PowerLogToolEnabled();
+    private RequestHandler<Object, Object> handler = new PowerLogToolEnabled();
 
     @Mock
     private Context context;
@@ -52,8 +54,7 @@ void setUp() throws IOException, IllegalAccessException, NoSuchMethodException,
         setupContext();
         //Make sure file is cleaned up before running full stack logging regression
         FileChannel.open(Paths.get("target/logfile.json"), StandardOpenOption.WRITE).truncate(0).close();
-        writeStaticField(LambdaLoggingAspect.class, "LOG_LEVEL", "INFO", true);
-        resetLogLevel();
+        resetLogLevel(Level.INFO);
     }
 
     @Test
@@ -74,8 +75,7 @@ void shouldLogInStructuredFormat() throws IOException {
 
     @Test
     void shouldModifyLogLevelBasedOnEnvVariable() throws IllegalAccessException, IOException, NoSuchMethodException, InvocationTargetException {
-        writeStaticField(LambdaLoggingAspect.class, "LOG_LEVEL", "DEBUG", true);
-        resetLogLevel();
+        resetLogLevel(Level.DEBUG);
 
         handler.handleRequest("test", context);
 
@@ -92,10 +92,34 @@ void shouldModifyLogLevelBasedOnEnvVariable() throws IllegalAccessException, IOE
                 });
     }
 
-    private void resetLogLevel() throws NoSuchMethodException, IllegalAccessException, InvocationTargetException {
-        Method resetLogLevels = LambdaLoggingAspect.class.getDeclaredMethod("resetLogLevels");
+    @Test
+    void shouldModifyLogLevelBasedOnSamplingRule() throws IOException {
+        handler = new PowerLogToolSamplingEnabled();
+
+        handler.handleRequest("test", context);
+
+        assertThat(Files.lines(Paths.get("target/logfile.json")))
+                .hasSize(3)
+                .satisfies(line -> {
+                    assertThat(parseToMap(line.get(0)))
+                            .containsEntry("level", "DEBUG")
+                            .containsEntry("loggerName", LambdaLoggingAspect.class.getCanonicalName());
+
+                    assertThat(parseToMap(line.get(1)))
+                            .containsEntry("level", "INFO")
+                            .containsEntry("message", "Test event");
+
+                    assertThat(parseToMap(line.get(2)))
+                            .containsEntry("level", "DEBUG")
+                            .containsEntry("message", "Test debug event");
+                });
+    }
+
+    private void resetLogLevel(Level level) throws NoSuchMethodException, IllegalAccessException, InvocationTargetException {
+        Method resetLogLevels = LambdaLoggingAspect.class.getDeclaredMethod("resetLogLevels", Level.class);
         resetLogLevels.setAccessible(true);
-        resetLogLevels.invoke(null);
+        resetLogLevels.invoke(null, level);
+        writeStaticField(LambdaLoggingAspect.class, "LEVEL_AT_INITIALISATION", level, true);
     }
 
     private Map<String, Object> parseToMap(String stringAsJson) {

powertools-logging/src/test/java/software/amazon/lambda/powertools/logging/handlers/PowerLogToolSamplingEnabled.java

@@ -0,0 +1,32 @@
+/*
+ * 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.logging.handlers;
+
+import com.amazonaws.services.lambda.runtime.Context;
+import com.amazonaws.services.lambda.runtime.RequestHandler;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import software.amazon.lambda.powertools.logging.PowertoolsLogging;
+
+public class PowerLogToolSamplingEnabled implements RequestHandler<Object, Object> {
+    private final Logger LOG = LogManager.getLogger(PowerLogToolSamplingEnabled.class);
+
+    @Override
+    @PowertoolsLogging(samplingRate = 1.0)
+    public Object handleRequest(Object input, Context context) {
+        LOG.info("Test event");
+        LOG.debug("Test debug event");
+        return null;
+    }
+}