Make cold start tracking thread-safe.

Author: phipag

powertools-logging/powertools-logging-log4j/src/test/java/org/apache/logging/log4j/layout/template/json/resolver/PowerToolsResolverFactoryTest.java

@@ -24,6 +24,7 @@
 import java.nio.file.NoSuchFileException;
 import java.nio.file.Paths;
 import java.nio.file.StandardOpenOption;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
@@ -48,7 +49,7 @@ void setUp() throws IllegalAccessException, IOException {
         MDC.clear();
         // Reset cold start state
         writeStaticField(LambdaHandlerProcessor.class, "isColdStart", null, true);
-        writeStaticField(PowertoolsLogging.class, "hasBeenInitialized", false, true);
+        writeStaticField(PowertoolsLogging.class, "hasBeenInitialized", new AtomicBoolean(false), true);
 
         context = new TestLambdaContext();
         // Make sure file is cleaned up before running tests

powertools-logging/powertools-logging-logback/src/test/java/software/amazon/lambda/powertools/logging/internal/LambdaEcsEncoderTest.java

@@ -25,6 +25,7 @@
 import java.nio.file.NoSuchFileException;
 import java.nio.file.Paths;
 import java.nio.file.StandardOpenOption;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
@@ -57,7 +58,7 @@ void setUp() throws IllegalAccessException, IOException {
         MDC.clear();
         // Reset cold start state
         writeStaticField(LambdaHandlerProcessor.class, "isColdStart", null, true);
-        writeStaticField(PowertoolsLogging.class, "hasBeenInitialized", false, true);
+        writeStaticField(PowertoolsLogging.class, "hasBeenInitialized", new AtomicBoolean(false), true);
 
         context = new TestLambdaContext();
         // Make sure file is cleaned up before running tests

powertools-logging/powertools-logging-logback/src/test/java/software/amazon/lambda/powertools/logging/internal/LambdaJsonEncoderTest.java

@@ -43,6 +43,8 @@
 import java.util.Collections;
 import java.util.Date;
 import java.util.TimeZone;
+import java.util.concurrent.atomic.AtomicBoolean;
+
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeEach;
@@ -76,7 +78,7 @@ void setUp() throws IllegalAccessException, IOException {
         MDC.clear();
         // Reset cold start state
         writeStaticField(LambdaHandlerProcessor.class, "isColdStart", null, true);
-        writeStaticField(PowertoolsLogging.class, "hasBeenInitialized", false, true);
+        writeStaticField(PowertoolsLogging.class, "hasBeenInitialized", new AtomicBoolean(false), true);
 
         context = new TestLambdaContext();
         // Make sure file is cleaned up before running tests

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

@@ -29,6 +29,7 @@
 import java.util.Arrays;
 import java.util.Locale;
 import java.util.Random;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -65,7 +66,7 @@
 public final class PowertoolsLogging {
     private static final Logger LOG = LoggerFactory.getLogger(PowertoolsLogging.class);
     private static final ThreadLocal<Random> SAMPLER = ThreadLocal.withInitial(Random::new);
-    private static volatile boolean hasBeenInitialized = false;
+    private static AtomicBoolean hasBeenInitialized = new AtomicBoolean(false);
 
     static {
         initializeLogLevel();
@@ -176,16 +177,14 @@ 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.
+     * 
      * @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)
      */
     public static void initializeLogging(Context context, double samplingRate, String correlationIdPath, Object event) {
-        if (hasBeenInitialized) {
-            coldStartDone();
-        }
-        hasBeenInitialized = true;
 
         addLambdaContextToLoggingContext(context);
         setLogLevelBasedOnSamplingRate(samplingRate);
@@ -196,12 +195,17 @@ public static void initializeLogging(Context context, double samplingRate, Strin
         }
     }
 
-    private static void addLambdaContextToLoggingContext(Context context) {
+    // Synchronized since isColdStart() is a globally managed constant in LambdaHandlerProcessor
+    private static synchronized void addLambdaContextToLoggingContext(Context context) {
         if (context != null) {
             PowertoolsLoggedFields.setValuesFromLambdaContext(context).forEach(MDC::put);
-            MDC.put(FUNCTION_COLD_START.getName(), isColdStart() ? "true" : "false");
-            MDC.put(SERVICE.getName(), serviceName());
         }
+
+        MDC.put(FUNCTION_COLD_START.getName(), isColdStart() ? "true" : "false");
+        if (hasBeenInitialized.compareAndSet(false, true)) {
+            coldStartDone();
+        }
+        MDC.put(SERVICE.getName(), serviceName());
     }
 
     private static void setLogLevelBasedOnSamplingRate(double samplingRate) {

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

@@ -27,6 +27,7 @@
 import java.nio.file.Paths;
 import java.nio.file.StandardOpenOption;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
@@ -68,7 +69,7 @@ void setUp() throws IllegalAccessException, IOException {
 
         // Reset cold start state
         writeStaticField(LambdaHandlerProcessor.class, "isColdStart", null, true);
-        writeStaticField(PowertoolsLogging.class, "hasBeenInitialized", false, true);
+        writeStaticField(PowertoolsLogging.class, "hasBeenInitialized", new AtomicBoolean(false), true);
 
         try {
             FileChannel.open(Paths.get("target/logfile.json"), StandardOpenOption.WRITE).truncate(0).close();
@@ -313,7 +314,7 @@ void initializeLogging_withEnvVarAndParameter_shouldUseEnvVarPrecedence() throws
     @Test
     void initializeLogging_calledTwice_shouldMarkColdStartDoneOnSecondCall() throws IllegalAccessException {
         // GIVEN
-        writeStaticField(PowertoolsLogging.class, "hasBeenInitialized", false, true);
+        writeStaticField(PowertoolsLogging.class, "hasBeenInitialized", new AtomicBoolean(false), true);
 
         // WHEN - First call
         PowertoolsLogging.initializeLogging(context);
@@ -371,21 +372,24 @@ void initializeLogging_concurrentCalls_shouldBeThreadSafe() throws InterruptedEx
         int threadCount = 10;
         Thread[] threads = new Thread[threadCount];
         String[] samplingRates = new String[threadCount];
+        boolean[] coldStarts = new boolean[threadCount];
         boolean[] success = new boolean[threadCount];
 
         // WHEN - Multiple threads call initializeLogging with alternating sampling rates
         for (int i = 0; i < threadCount; i++) {
             final int threadIndex = i;
             final double samplingRate = (i % 2 == 0) ? 1.0 : 0.0; // Alternate between 1.0 and 0.0
-            
+
             threads[i] = new Thread(() -> {
                 try {
                     PowertoolsLogging.initializeLogging(context, samplingRate);
-                    
-                    // Capture the sampling rate set in MDC (thread-local)
+
+                    // Capture the sampling rate and cold start values set in MDC (thread-local)
                     samplingRates[threadIndex] = MDC.get(PowertoolsLoggedFields.SAMPLING_RATE.getName());
+                    coldStarts[threadIndex] = Boolean
+                            .parseBoolean(MDC.get(PowertoolsLoggedFields.FUNCTION_COLD_START.getName()));
                     success[threadIndex] = true;
-                    
+
                     // Clean up thread-local state
                     PowertoolsLogging.clearState(true);
                 } catch (Exception e) {
@@ -408,12 +412,17 @@ void initializeLogging_concurrentCalls_shouldBeThreadSafe() throws InterruptedEx
         for (boolean result : success) {
             assertThat(result).isTrue();
         }
-        
-        // THEN - Each thread should have its own sampling rate in MDC
+
+        // THEN - Each thread should have its own sampling rate in MDC and exactly one invocation was a cold start
+        int coldStartCount = 0;
         for (int i = 0; i < threadCount; i++) {
-            String expectedRate = (i % 2 == 0) ? "1.0" : "0.0";
-            assertThat(samplingRates[i]).as("Thread %d should have sampling rate %s", i, expectedRate).isEqualTo(expectedRate);
+            String expectedSamplingRate = (i % 2 == 0) ? "1.0" : "0.0";
+            assertThat(samplingRates[i]).as("Thread %d should have sampling rate %s", i, expectedSamplingRate)
+                    .isEqualTo(expectedSamplingRate);
+
+            coldStartCount += coldStarts[i] ? 1 : 0;
         }
+        assertThat(coldStartCount).isEqualTo(1);
     }
 
     private void reinitializeLogLevel() {

powertools-logging/src/test/java/software/amazon/lambda/powertools/logging/internal/LambdaLoggingAspectTest.java

@@ -39,6 +39,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
@@ -99,7 +100,7 @@ void setUp() throws IllegalAccessException, IOException {
 
         // Reset cold start state
         writeStaticField(LambdaHandlerProcessor.class, "isColdStart", null, true);
-        writeStaticField(PowertoolsLogging.class, "hasBeenInitialized", false, true);
+        writeStaticField(PowertoolsLogging.class, "hasBeenInitialized", new AtomicBoolean(false), true);
 
         context = new TestLambdaContext();
         requestHandler = new PowertoolsLogEnabled();