feat(logger): refresh sample rate calculation before each invocation (#3672)

Author: am29d

docs/core/logger.md

@@ -634,7 +634,7 @@ Once a child logger is created, the logger and its parent will act as separate i
 
 ### Sampling debug logs
 
-Use sampling when you want to dynamically change your log level to **DEBUG** based on a **percentage of your concurrent/cold start invocations**.
+Use sampling when you want to dynamically change your log level to **DEBUG** based on a **percentage of your invocations**.
 
 You can use values ranging from `0` to `1` (100%) when setting the `sampleRateValue` constructor option or `POWERTOOLS_LOGGER_SAMPLE_RATE` env var.
 
@@ -643,10 +643,9 @@ 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. This means sampling may happen significantly more or less than depending on your traffic patterns, for example a steady low number of invocations and thus few cold starts.
+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.
 
-!!! note
-    Open a [feature request](https://github.com/aws-powertools/powertools-lambda-typescript/issues/new?assignees=&labels=type%2Ffeature-request%2Ctriage&projects=aws-powertools%2F7&template=feature_request.yml&title=Feature+request%3A+TITLE) if you want Logger to calculate sampling for every invocation
+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"
 

examples/snippets/logger/logSampling.ts

@@ -10,6 +10,8 @@ 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');

packages/commons/src/Utility.ts

@@ -51,7 +51,7 @@
  * ```
  */
 export class Utility {
-  private coldStart = true;
+  protected coldStart = true;
   private readonly defaultServiceName: string = 'service_undefined';
 
   /**

packages/logger/src/Logger.ts

@@ -419,6 +419,7 @@ class Logger extends Utility implements LoggerInterface {
         context,
         callback
       ) {
+        loggerRef.refreshSampleRateCalculation();
         Logger.injectLambdaContextBefore(loggerRef, event, context, options);
 
         let result: unknown;
@@ -493,9 +494,13 @@ class Logger extends Utility implements LoggerInterface {
    * This method allows recalculating the initial sampling decision for changing
    * the log level to DEBUG based on a sample rate value used during initialization,
    * potentially yielding a different outcome.
+   *
+   * This only works for warm starts, because we don't to avoid double sampling.
    */
   public refreshSampleRateCalculation(): void {
-    this.setInitialSampleRate(this.powertoolsLogData.sampleRateValue);
+    if (!this.coldStart) {
+      this.setInitialSampleRate(this.powertoolsLogData.sampleRateValue);
+    }
   }
 
   /**

packages/logger/src/middleware/middy.ts

@@ -87,6 +87,9 @@ const injectLambdaContext = (
       if (isResetStateEnabled) {
         setCleanupFunction(request);
       }
+
+      logger.refreshSampleRateCalculation();
+
       Logger.injectLambdaContextBefore(
         logger,
         request.event,

packages/logger/tests/unit/injectLambdaContext.test.ts

@@ -184,4 +184,39 @@ describe('Inject Lambda Context', () => {
       })
     );
   });
+
+  it('refreshes sample rate calculation before each invocation using decorator for warm start only', async () => {
+    // Prepare
+    const logger = new Logger({ sampleRateValue: 0.5 });
+    const refreshSpy = vi.spyOn(logger, 'refreshSampleRateCalculation');
+
+    class Lambda {
+      @logger.injectLambdaContext()
+      public async handler(_event: unknown, _context: Context): Promise<void> {
+        logger.info('test');
+      }
+    }
+    const lambda = new Lambda();
+    // Act
+    await lambda.handler({}, {} as Context);
+
+    // Assess
+    expect(refreshSpy).toHaveBeenCalledTimes(1);
+  });
+
+  it('refreshes sample rate calculation before each invocation using middleware for warm start only', async () => {
+    // Prepare
+    const logger = new Logger({ sampleRateValue: 0.5 });
+    const refreshSpy = vi.spyOn(logger, 'refreshSampleRateCalculation');
+
+    const handler = middy(async () => {
+      logger.info('Hello, world!');
+    }).use(injectLambdaContext(logger));
+
+    // Act
+    await handler(event, context);
+
+    // Assess
+    expect(refreshSpy).toHaveBeenCalledTimes(1);
+  });
 });

packages/logger/tests/unit/sampling.test.ts

@@ -129,13 +129,15 @@ describe('Log sampling', () => {
     expect(logger.getLevelName()).toBe(LogLevel.INFO);
   });
 
-  it('when sample rate calculation is refreshed, it respects probability sampling and change log level to DEBUG ', () => {
+  it('when sample rate calculation is refreshed, it respects probability sampling and change log level to DEBUG when not in coldStart ', () => {
     // Prepare
     const logger = new Logger({
       logLevel: LogLevel.ERROR,
       sampleRateValue: 0.1, // 10% probability
     });
 
+    logger.getColdStart(); // Set coldStart to false
+
     let logLevelChangedToDebug = 0;
     const numOfIterations = 1000;
     const minExpected = numOfIterations * 0.05; // Min expected based on 5% probability