fix(logger): correctly refresh sample rate (#3722)

Co-authored-by: Alexander Schueren <sha@amazon.com>

Author: dreamorosi

docs/core/logger.md

@@ -805,114 +805,41 @@ You can use values ranging from `0` to `1` (100%) when setting the `sampleRateVa
 
     This feature takes into account transient issues where additional debugging information can be useful.
 
-Sampling decision happens at the Logger initialization. When using the `injectLambdaContext` method either as a decorator or middleware, the sampling decision is refreshed at the beginning of each Lambda invocation for you, except for cold starts.
+Sampling decision happens at the Logger initialization. When using the `injectLambdaContext` method either as a decorator or Middy.js middleware, the sampling decision is refreshed at the beginning of each Lambda invocation for you, except for cold starts.
 
 If you're not using either of these, you'll need to manually call the `refreshSamplingRate()` function at the start of your handler to refresh the sampling decision for each invocation.
 
 === "handler.ts"
 
-    ```typescript hl_lines="6"
+    ```typescript hl_lines="5 9"
     --8<-- "examples/snippets/logger/logSampling.ts"
     ```
 
-=== "Example CloudWatch Logs excerpt - Invocation #1"
+    1. The log level must be set to a more verbose level than `DEBUG` for log sampling to kick in.
+    2. You need to call `logger.refreshSamplingRate()` at the start of your handler **only** if you're not using the `injectLambdaContext()` class method decorator or Middy.js middleware.
+
+=== "Example Logs Request #1 (not sampled)"
 
     ```json
-    {
-        "level": "ERROR",
-        "message": "This is an ERROR log",
-        "sampling_rate": "0.5",
-        "service": "serverlessAirline",
-        "timestamp": "2021-12-12T22:59:06.334Z",
-        "xray_trace_id": "abcdef123456abcdef123456abcdef123456"
-    }
-    {
-        "level": "DEBUG",
-        "message": "This is a DEBUG log that has 50% chance of being printed",
-        "sampling_rate": "0.5", 
-        "service": "serverlessAirline",
-        "timestamp": "2021-12-12T22:59:06.337Z",
-        "xray_trace_id": "abcdef123456abcdef123456abcdef123456"
-    }
-    {
-        "level": "INFO",
-        "message": "This is an INFO log that has 50% chance of being printed",
-        "sampling_rate": "0.5", 
-        "service": "serverlessAirline",
-        "timestamp": "2021-12-12T22:59:06.338Z",
-        "xray_trace_id": "abcdef123456abcdef123456abcdef123456"
-    }
-    {
-        "level": "WARN",
-        "message": "This is a WARN log that has 50% chance of being printed",
-        "sampling_rate": "0.5", 
-        "service": "serverlessAirline",
-        "timestamp": "2021-12-12T22:59:06.338Z",
-        "xray_trace_id": "abcdef123456abcdef123456abcdef123456"
-    }
+    --8<-- "examples/snippets/logger/samples/debugLogSamplingNotSampled.json"
     ```
 
-=== "Example CloudWatch Logs excerpt - Invocation #2"
+=== "Example Logs Request #2 (sampled)"
 
     ```json
-    {
-        "level": "ERROR",
-        "message": "This is an ERROR log",
-        "sampling_rate": "0.5",
-        "service": "serverlessAirline",
-        "timestamp": "2021-12-12T22:59:06.334Z",
-        "xray_trace_id": "abcdef123456abcdef123456abcdef123456"
-    }
+    --8<-- "examples/snippets/logger/samples/debugLogSamplingSampled.json"
     ```
 
-=== "Example CloudWatch Logs excerpt - Invocation #3"
+=== "Example Logs Request #3 (sampled)"
 
     ```json
-    {
-        "level": "ERROR",
-        "message": "This is an ERROR log",
-        "sampling_rate": "0.5",
-        "service": "serverlessAirline",
-        "timestamp": "2021-12-12T22:59:06.334Z",
-        "xray_trace_id": "abcdef123456abcdef123456abcdef123456"
-    }
-    {
-        "level": "DEBUG",
-        "message": "This is a DEBUG log that has 50% chance of being printed",
-        "sampling_rate": "0.5", 
-        "service": "serverlessAirline",
-        "timestamp": "2021-12-12T22:59:06.337Z",
-        "xray_trace_id": "abcdef123456abcdef123456abcdef123456"
-    }
-    {
-        "level": "INFO",
-        "message": "This is an INFO log that has 50% chance of being printed",
-        "sampling_rate": "0.5", 
-        "service": "serverlessAirline",
-        "timestamp": "2021-12-12T22:59:06.338Z",
-        "xray_trace_id": "abcdef123456abcdef123456abcdef123456"
-    }
-    {
-        "level": "WARN",
-        "message": "This is a WARN log that has 50% chance of being printed",
-        "sampling_rate": "0.5", 
-        "service": "serverlessAirline",
-        "timestamp": "2021-12-12T22:59:06.338Z",
-        "xray_trace_id": "abcdef123456abcdef123456abcdef123456"
-    }
+    --8<-- "examples/snippets/logger/samples/debugLogSamplingSampled.json"
     ```
 
-=== "Example CloudWatch Logs excerpt - Invocation #4"
+=== "Example Logs Request #4 (not sampled)"
 
     ```json
-    {
-        "level": "ERROR",
-        "message": "This is an ERROR log",
-        "sampling_rate": "0.5",
-        "service": "serverlessAirline",
-        "timestamp": "2021-12-12T22:59:06.334Z",
-        "xray_trace_id": "abcdef123456abcdef123456abcdef123456"
-    }
+    --8<-- "examples/snippets/logger/samples/debugLogSamplingNotSampled.json"
     ```
 
 ### Custom Log formatter

examples/snippets/logger/logSampling.ts

@@ -1,27 +1,16 @@
 import { Logger } from '@aws-lambda-powertools/logger';
 
-// Notice the log level set to 'ERROR'
 const logger = new Logger({
-  logLevel: 'ERROR',
+  logLevel: 'ERROR', // (1)!
   sampleRateValue: 0.5,
 });
 
-export const handler = async (
-  _event: unknown,
-  _context: unknown
-): Promise<void> => {
-  // Refresh sample rate calculation on runtime, only when not using injectLambdaContext
-  logger.refreshSampleRateCalculation();
-  // This log item (equal to log level 'ERROR') will be printed to standard output
-  // in all Lambda invocations
-  logger.error('This is an ERROR log');
+export const handler = async () => {
+  logger.refreshSampleRateCalculation(); // (2)!
 
-  // These log items (below the log level 'ERROR') have ~50% chance
-  // of being printed in a Lambda invocation
-  logger.debug('This is a DEBUG log that has 50% chance of being printed');
-  logger.info('This is an INFO log that has 50% chance of being printed');
-  logger.warn('This is a WARN log that has 50% chance of being printed');
+  logger.error('This log is always emitted');
 
-  // Optional: refresh sample rate calculation on runtime
-  // logger.refreshSampleRateCalculation();
+  logger.debug('This log has ~50% chance of being emitted');
+  logger.info('This log has ~50% chance of being emitted');
+  logger.warn('This log has ~50% chance of being emitted');
 };

examples/snippets/logger/samples/debugLogSamplingNotSampled.json

@@ -0,0 +1,8 @@
+{
+  "level": "ERROR",
+  "message": "This log is always emitted",
+  "sampling_rate": "0.5",
+  "service": "serverlessAirline",
+  "timestamp": "2021-12-12T22:59:06.334Z",
+  "xray_trace_id": "abcdef123456abcdef123456abcdef123456"
+}

examples/snippets/logger/samples/debugLogSamplingSampled.json

@@ -0,0 +1,34 @@
+[
+  {
+    "level": "ERROR",
+    "message": "This log is always emitted",
+    "sampling_rate": "0.5",
+    "service": "serverlessAirline",
+    "timestamp": "2021-12-12T22:59:06.334Z",
+    "xray_trace_id": "abcdef123456abcdef123456abcdef123456"
+  },
+  {
+    "level": "DEBUG",
+    "message": "This log has ~50% chance of being emitted",
+    "sampling_rate": "0.5",
+    "service": "serverlessAirline",
+    "timestamp": "2021-12-12T22:59:06.337Z",
+    "xray_trace_id": "abcdef123456abcdef123456abcdef123456"
+  },
+  {
+    "level": "INFO",
+    "message": "This log has ~50% chance of being emitted",
+    "sampling_rate": "0.5",
+    "service": "serverlessAirline",
+    "timestamp": "2021-12-12T22:59:06.338Z",
+    "xray_trace_id": "abcdef123456abcdef123456abcdef123456"
+  },
+  {
+    "level": "WARN",
+    "message": "This log has ~50% chance of being emitted",
+    "sampling_rate": "0.5",
+    "service": "serverlessAirline",
+    "timestamp": "2021-12-12T22:59:06.338Z",
+    "xray_trace_id": "abcdef123456abcdef123456abcdef123456"
+  }
+]

packages/logger/src/Logger.ts

@@ -140,7 +140,9 @@ class Logger extends Utility implements LoggerInterface {
   /**
    * Standard attributes managed by Powertools that will be logged in all log items.
    */
-  private powertoolsLogData: PowertoolsLogData = <PowertoolsLogData>{};
+  private powertoolsLogData: PowertoolsLogData = <PowertoolsLogData>{
+    sampleRateValue: 0,
+  };
   /**
    * Temporary log attributes that can be appended with `appendKeys()` method.
    */
@@ -213,6 +215,22 @@ class Logger extends Utility implements LoggerInterface {
    */
   #buffer?: CircularMap<string>;
 
+  /**
+   * The debug sampling rate configuration.
+   */
+  readonly #debugLogSampling = {
+    /**
+     * The sampling rate value used to determine if the log level should be set to DEBUG.
+     */
+    sampleRateValue: 0,
+    /**
+     * The number of times the debug sampling rate has been refreshed.
+     *
+     * We use this to determine if we should refresh it again.
+     */
+    refreshedTimes: 0,
+  };
+
   /**
    * Log level used by the current instance of Logger.
    *
@@ -302,7 +320,7 @@ class Logger extends Utility implements LoggerInterface {
         {
           logLevel: this.getLevelName(),
           serviceName: this.powertoolsLogData.serviceName,
-          sampleRateValue: this.powertoolsLogData.sampleRateValue,
+          sampleRateValue: this.#debugLogSampling.sampleRateValue,
           logFormatter: this.getLogFormatter(),
           customConfigService: this.getCustomConfigService(),
           environment: this.powertoolsLogData.environment,
@@ -547,8 +565,18 @@ class Logger extends Utility implements LoggerInterface {
    * This only works for warm starts, because we don't to avoid double sampling.
    */
   public refreshSampleRateCalculation(): void {
-    if (!this.coldStart) {
-      this.setInitialSampleRate(this.powertoolsLogData.sampleRateValue);
+    if (this.#debugLogSampling.refreshedTimes === 0) {
+      this.#debugLogSampling.refreshedTimes++;
+      return;
+    }
+    if (
+      this.#shouldEnableDebugSampling() &&
+      this.logLevel > LogLevelThreshold.TRACE
+    ) {
+      this.setLogLevel('DEBUG');
+      this.debug('Setting log level to DEBUG due to sampling rate');
+    } else {
+      this.setLogLevel(this.getLogLevelNameFromNumber(this.#initialLogLevel));
     }
   }
 
@@ -891,6 +919,16 @@ class Logger extends Utility implements LoggerInterface {
     }
   }
 
+  /**
+   * Make a new debug log sampling decision based on the sample rate value.
+   */
+  #shouldEnableDebugSampling() {
+    return (
+      this.#debugLogSampling.sampleRateValue &&
+      randomInt(0, 100) / 100 <= this.#debugLogSampling.sampleRateValue
+    );
+  }
+
   /**
    * Check if a given key is reserved and warn the user if it is.
    *
@@ -1142,30 +1180,24 @@ class Logger extends Utility implements LoggerInterface {
    * @param sampleRateValue - The sample rate value
    */
   private setInitialSampleRate(sampleRateValue?: number): void {
-    this.powertoolsLogData.sampleRateValue = 0;
     const constructorValue = sampleRateValue;
     const customConfigValue =
       this.getCustomConfigService()?.getSampleRateValue();
     const envVarsValue = this.getEnvVarsService().getSampleRateValue();
     for (const value of [constructorValue, customConfigValue, envVarsValue]) {
       if (this.isValidSampleRate(value)) {
+        this.#debugLogSampling.sampleRateValue = value;
         this.powertoolsLogData.sampleRateValue = value;
 
         if (
-          this.logLevel > LogLevelThreshold.DEBUG &&
-          value &&
-          randomInt(0, 100) / 100 <= value
+          this.#shouldEnableDebugSampling() &&
+          this.logLevel > LogLevelThreshold.TRACE
         ) {
-          // only change logLevel if higher than debug, i.e. don't change from e.g. tracing to debug
           this.setLogLevel('DEBUG');
           this.debug('Setting log level to DEBUG due to sampling rate');
-        } else {
-          this.setLogLevel(
-            this.getLogLevelNameFromNumber(this.#initialLogLevel)
-          );
         }
 
-        return;
+        break;
       }
     }
   }