Add withLogging callback based interface to PowertoolsLogging.

Author: phipag

powertools-e2e-tests/handlers/logging-functional/src/main/java/software/amazon/lambda/powertools/e2e/Function.java

@@ -21,21 +21,20 @@
 import com.amazonaws.services.lambda.runtime.Context;
 import com.amazonaws.services.lambda.runtime.RequestHandler;
 
-import software.amazon.lambda.powertools.logging.Logging;
 import software.amazon.lambda.powertools.logging.PowertoolsLogging;
 
 public class Function implements RequestHandler<Input, String> {
     private static final Logger LOG = LoggerFactory.getLogger(Function.class);
 
     public String handleRequest(Input input, Context context) {
-        PowertoolsLogging.initializeLogging(context);
+        return PowertoolsLogging.withLogging(context, () -> {
+            input.getKeys().forEach(MDC::put);
+            LOG.info(input.getMessage());
 
-        input.getKeys().forEach(MDC::put);
-        LOG.info(input.getMessage());
+            // Flush buffer manually since we buffer at INFO level to test log buffering
+            PowertoolsLogging.flushBuffer();
 
-        // Flush buffer manually since we buffer at INFO level to test log buffering
-        PowertoolsLogging.flushBuffer();
-
-        return "OK";
+            return "OK";
+        });
     }
 }

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

@@ -30,6 +30,7 @@
 import java.util.Locale;
 import java.util.Random;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Supplier;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -140,6 +141,9 @@ public static void clearBuffer() {
      * This method should be called at the beginning of your Lambda handler to set up
      * logging context with Lambda function information, trace ID, and service name.
      * 
+     * <p>Important: Call {@link #clearState(boolean)} at the end of your handler or use
+     * {@link #withLogging(Context, Supplier)} to handle cleanup automatically.</p>
+     * 
      * @param context the Lambda context provided by AWS Lambda runtime
      */
     public static void initializeLogging(Context context) {
@@ -151,6 +155,9 @@ public static void initializeLogging(Context context) {
      * This method sets up logging context and optionally enables DEBUG logging
      * based on the provided sampling rate.
      * 
+     * <p>Important: Call {@link #clearState(boolean)} at the end of your handler or use
+     * {@link #withLogging(Context, double, Supplier)} to handle cleanup automatically.</p>
+     * 
      * @param context the Lambda context provided by AWS Lambda runtime
      * @param samplingRate sampling rate for DEBUG logging (0.0 to 1.0)
      */
@@ -163,6 +170,9 @@ public static void initializeLogging(Context context, double samplingRate) {
      * This method sets up logging context and extracts correlation ID from the event
      * using the provided JSON path.
      * 
+     * <p>Important: Call {@link #clearState(boolean)} at the end of your handler or use
+     * {@link #withLogging(Context, String, Object, Supplier)} to handle cleanup automatically.</p>
+     * 
      * @param context the Lambda context provided by AWS Lambda runtime
      * @param correlationIdPath JSON path to extract correlation ID from event
      * @param event the Lambda event object
@@ -177,7 +187,10 @@ public static void initializeLogging(Context context, String correlationIdPath,
      * configures sampling rate for DEBUG logging, and optionally extracts
      * correlation ID from the event.
      * 
-     * This method is tread-safe.
+     * <p>Important: Call {@link #clearState(boolean)} at the end of your handler or use
+     * {@link #withLogging(Context, double, String, Object, Supplier)} to handle cleanup automatically.</p>
+     * 
+     * <p>This method is thread-safe.</p>
      * 
      * @param context the Lambda context provided by AWS Lambda runtime
      * @param samplingRate sampling rate for DEBUG logging (0.0 to 1.0)
@@ -278,4 +291,79 @@ public static void clearState(boolean clearMdcState) {
         clearBuffer();
         SAMPLER.remove();
     }
+
+    /**
+     * Executes code with logging context initialized and automatically clears state.
+     * 
+     * @param context the Lambda context provided by AWS Lambda runtime
+     * @param supplier the code to execute with logging context
+     * @param <T> the return type
+     * @return the result of the supplier execution
+     */
+    public static <T> T withLogging(Context context, Supplier<T> supplier) {
+        initializeLogging(context);
+        try {
+            return supplier.get();
+        } finally {
+            clearState(true);
+        }
+    }
+
+    /**
+     * Executes code with logging context initialized with sampling rate and automatically clears state.
+     * 
+     * @param context the Lambda context provided by AWS Lambda runtime
+     * @param samplingRate sampling rate for DEBUG logging (0.0 to 1.0)
+     * @param supplier the code to execute with logging context
+     * @param <T> the return type
+     * @return the result of the supplier execution
+     */
+    public static <T> T withLogging(Context context, double samplingRate, Supplier<T> supplier) {
+        initializeLogging(context, samplingRate);
+        try {
+            return supplier.get();
+        } finally {
+            clearState(true);
+        }
+    }
+
+    /**
+     * Executes code with logging context initialized with correlation ID extraction and automatically clears state.
+     * 
+     * @param context the Lambda context provided by AWS Lambda runtime
+     * @param correlationIdPath JSON path to extract correlation ID from event
+     * @param event the Lambda event object
+     * @param supplier the code to execute with logging context
+     * @param <T> the return type
+     * @return the result of the supplier execution
+     */
+    public static <T> T withLogging(Context context, String correlationIdPath, Object event, Supplier<T> supplier) {
+        initializeLogging(context, correlationIdPath, event);
+        try {
+            return supplier.get();
+        } finally {
+            clearState(true);
+        }
+    }
+
+    /**
+     * Executes code with logging context initialized with full configuration and automatically clears state.
+     * 
+     * @param context the Lambda context provided by AWS Lambda runtime
+     * @param samplingRate sampling rate for DEBUG logging (0.0 to 1.0)
+     * @param correlationIdPath JSON path to extract correlation ID from event (can be null)
+     * @param event the Lambda event object (required if correlationIdPath is provided)
+     * @param supplier the code to execute with logging context
+     * @param <T> the return type
+     * @return the result of the supplier execution
+     */
+    public static <T> T withLogging(Context context, double samplingRate, String correlationIdPath, Object event,
+            Supplier<T> supplier) {
+        initializeLogging(context, samplingRate, correlationIdPath, event);
+        try {
+            return supplier.get();
+        } finally {
+            clearState(true);
+        }
+    }
 }

powertools-logging/src/test/java/software/amazon/lambda/powertools/logging/PowertoolsLoggingTest.java

@@ -425,6 +425,86 @@ void initializeLogging_concurrentCalls_shouldBeThreadSafe() throws InterruptedEx
         assertThat(coldStartCount).isEqualTo(1);
     }
 
+    @Test
+    void withLogging_basicUsage_shouldInitializeAndCleanup() {
+        // WHEN
+        String result = PowertoolsLogging.withLogging(context, () -> {
+            assertThat(MDC.get(PowertoolsLoggedFields.FUNCTION_NAME.getName())).isEqualTo("test-function");
+            return "test-result";
+        });
+
+        // THEN
+        assertThat(result).isEqualTo("test-result");
+        assertThat(MDC.getCopyOfContextMap()).isNull();
+        assertThat(testManager.isBufferCleared()).isTrue();
+    }
+
+    @Test
+    void withLogging_withSamplingRate_shouldSetSamplingRateAndCleanup() {
+        // WHEN
+        String result = PowertoolsLogging.withLogging(context, 0.5, () -> {
+            assertThat(MDC.get(PowertoolsLoggedFields.SAMPLING_RATE.getName())).isEqualTo("0.5");
+            return "sampled-result";
+        });
+
+        // THEN
+        assertThat(result).isEqualTo("sampled-result");
+        assertThat(MDC.getCopyOfContextMap()).isNull();
+    }
+
+    @Test
+    void withLogging_withCorrelationId_shouldExtractCorrelationIdAndCleanup() {
+        // GIVEN
+        Map<String, Object> event = Map.of("requestId", "correlation-123");
+
+        // WHEN
+        Integer result = PowertoolsLogging.withLogging(context, "requestId", event, () -> {
+            assertThat(MDC.get(PowertoolsLoggedFields.CORRELATION_ID.getName())).isEqualTo("correlation-123");
+            return 42;
+        });
+
+        // THEN
+        assertThat(result).isEqualTo(42);
+        assertThat(MDC.getCopyOfContextMap()).isNull();
+    }
+
+    @Test
+    void withLogging_withFullConfiguration_shouldSetAllFieldsAndCleanup() {
+        // GIVEN
+        Map<String, Object> event = Map.of("id", "full-correlation");
+
+        // WHEN
+        Boolean result = PowertoolsLogging.withLogging(context, 0.8, "id", event, () -> {
+            Map<String, String> mdcMap = MDC.getCopyOfContextMap();
+            assertThat(mdcMap)
+                    .containsEntry(PowertoolsLoggedFields.FUNCTION_NAME.getName(), "test-function")
+                    .containsEntry(PowertoolsLoggedFields.CORRELATION_ID.getName(), "full-correlation")
+                    .containsEntry(PowertoolsLoggedFields.SAMPLING_RATE.getName(), "0.8");
+            return true;
+        });
+
+        // THEN
+        assertThat(result).isTrue();
+        assertThat(MDC.getCopyOfContextMap()).isNull();
+    }
+
+    @Test
+    void withLogging_whenSupplierThrowsException_shouldStillCleanup() {
+        // WHEN & THEN
+        try {
+            PowertoolsLogging.withLogging(context, () -> {
+                assertThat(MDC.get(PowertoolsLoggedFields.FUNCTION_NAME.getName())).isEqualTo("test-function");
+                throw new RuntimeException("test exception");
+            });
+        } catch (RuntimeException e) {
+            assertThat(e.getMessage()).isEqualTo("test exception");
+        }
+
+        // THEN - cleanup should still happen
+        assertThat(MDC.getCopyOfContextMap()).isNull();
+        assertThat(testManager.isBufferCleared()).isTrue();
+    }
+
     private void reinitializeLogLevel() {
         try {
             Method initializeLogLevel = PowertoolsLogging.class.getDeclaredMethod("initializeLogLevel");