Addressed review comments for logger

Author: maschnetwork

examples/pom.xml

@@ -40,12 +40,12 @@
         <dependency>
             <groupId>org.apache.logging.log4j</groupId>
             <artifactId>log4j-slf4j2-impl</artifactId>
-            <version>2.24.3</version>
+            <version>2.25.3</version>
         </dependency>
         <dependency>
             <groupId>org.apache.logging.log4j</groupId>
             <artifactId>log4j-core</artifactId>
-            <version>2.24.3</version>
+            <version>2.25.3</version>
         </dependency>
 
         <!-- Testing -->

examples/template.yaml

@@ -186,6 +186,28 @@ Resources:
       DockerContext: ../
       DockerTag: durable-examples
 
+  LoggingExampleFunction:
+    Type: AWS::Serverless::Function
+    Properties:
+      PackageType: Image
+      FunctionName: logging-example
+      ImageConfig:
+        Command: ["com.amazonaws.lambda.durable.examples.LoggingExample::handleRequest"]
+      DurableConfig:
+        ExecutionTimeout: 300
+        RetentionPeriodInDays: 7
+      Policies:
+        - Statement:
+            - Effect: Allow
+              Action:
+                - lambda:CheckpointDurableExecutions
+                - lambda:GetDurableExecutionState
+              Resource: !Sub "arn:aws:lambda:${AWS::Region}:${AWS::AccountId}:function:logging-example"
+    Metadata:
+      Dockerfile: examples/Dockerfile
+      DockerContext: ../
+      DockerTag: durable-examples
+
 Outputs:
   SimpleStepExampleFunction:
     Description: Simple Step Example Function ARN
@@ -250,3 +272,11 @@ Outputs:
   CustomConfigExampleFunctionName:
     Description: Custom Config Example Function Name
     Value: !Ref CustomConfigExampleFunction
+
+  LoggingExampleFunction:
+    Description: Logging Example Function ARN
+    Value: !GetAtt LoggingExampleFunction.Arn
+
+  LoggingExampleFunctionName:
+    Description: Logging Example Function Name
+    Value: !Ref LoggingExampleFunction

sdk/src/main/java/com/amazonaws/lambda/durable/execution/ExecutionManager.java

@@ -16,6 +16,7 @@
 import java.util.concurrent.Executor;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Phaser;
+import java.util.concurrent.atomic.AtomicReference;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import software.amazon.awssdk.services.lambda.model.Operation;
@@ -46,10 +47,10 @@ public class ExecutionManager {
     private final String executionOperationId;
     private volatile String checkpointToken;
     private final String durableExecutionArn;
-    private volatile ExecutionMode executionMode;
+    private final AtomicReference<ExecutionMode> executionMode;
 
     // ===== Current Operation Context (for logging) =====
-    private static final ThreadLocal<OperationContext> currentOperation = new ThreadLocal<>();
+    private static final ThreadLocal<OperationContext> currentOperation = new InheritableThreadLocal<>();
 
     // ===== Thread Coordination =====
     private final Map<String, ThreadType> activeThreads = Collections.synchronizedMap(new HashMap<>());
@@ -76,7 +77,7 @@ public ExecutionManager(
         loadAllOperations(initialExecutionState);
 
         // Start in REPLAY mode if we have more than just the initial EXECUTION operation
-        this.executionMode = operations.size() > 1 ? ExecutionMode.REPLAY : ExecutionMode.EXECUTION;
+        this.executionMode = new AtomicReference<>(operations.size() > 1 ? ExecutionMode.REPLAY : ExecutionMode.EXECUTION);
 
         this.managedExecutor = executor;
 
@@ -114,7 +115,7 @@ public String getDurableExecutionArn() {
     }
 
     public boolean isReplaying() {
-        return executionMode == ExecutionMode.REPLAY;
+        return executionMode.get() == ExecutionMode.REPLAY;
     }
 
     public String getCheckpointToken() {
@@ -152,9 +153,10 @@ public Operation getOperation(String operationId) {
      */
     public Operation getOperationAndUpdateReplayState(String operationId) {
         var existing = operations.get(operationId);
-        if (executionMode == ExecutionMode.REPLAY && existing == null) {
-            executionMode = ExecutionMode.EXECUTION;
-            logger.debug("Transitioned to EXECUTION mode at operation '{}'", operationId);
+        if (existing == null) {
+            if (executionMode.compareAndSet(ExecutionMode.REPLAY, ExecutionMode.EXECUTION)) {
+                logger.debug("Transitioned to EXECUTION mode at operation '{}'", operationId);
+            }
         }
         return existing;
     }
@@ -337,6 +339,7 @@ public CompletableFuture<Void> getSuspendExecutionFuture() {
 
     public void shutdown() {
         checkpointBatcher.shutdown();
+        currentOperation.remove();
     }
 
     private boolean isTerminalStatus(OperationStatus status) {

sdk/src/test/java/com/amazonaws/lambda/durable/logging/DurableLoggerTest.java

@@ -2,7 +2,6 @@
 // SPDX-License-Identifier: Apache-2.0
 package com.amazonaws.lambda.durable.logging;
 
-import static org.junit.jupiter.api.Assertions.*;
 import static org.mockito.Mockito.*;
 
 import com.amazonaws.lambda.durable.execution.ExecutionManager;