fix: align log metadata with other SDKs

Author: zhongkechen

examples/src/main/java/software/amazon/lambda/durable/examples/general/LoggingExample.java

@@ -2,6 +2,8 @@
 // SPDX-License-Identifier: Apache-2.0
 package software.amazon.lambda.durable.examples.general;
 
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import software.amazon.lambda.durable.DurableContext;
 import software.amazon.lambda.durable.DurableHandler;
 import software.amazon.lambda.durable.examples.types.GreetingRequest;
@@ -13,15 +15,16 @@
  * in log entries via MDC. By default, logs are suppressed during replay to avoid duplicates.
  */
 public class LoggingExample extends DurableHandler<GreetingRequest, String> {
+    Logger logger = LoggerFactory.getLogger(LoggingExample.class);
 
     @Override
     public String handleRequest(GreetingRequest input, DurableContext context) {
         // Log at execution level (outside any step)
-        context.getLogger().info("Processing greeting for: {}", input.getName());
+        context.getLogger(logger).info("Processing greeting for: {}", input.getName());
 
         // Step 1: Create greeting - logs inside step include operation context
         var greeting = context.step("create-greeting", String.class, ctx -> {
-            ctx.getLogger().info("Creating greeting message");
+            ctx.getLogger(logger).info("Creating greeting message");
             return "Hello, " + input.getName();
         });
 

examples/src/test/java/software/amazon/lambda/durable/examples/CloudBasedIntegrationTest.java

@@ -253,6 +253,14 @@ void testWaitAtLeastInProcessExample() {
         assertTrue(asyncOp.getStepResult(String.class).contains("Processed: TestUser"));
     }
 
+    @Test
+    void testLoggingExample() {
+        var runner = CloudDurableTestRunner.create(
+                arn("logging-example"), GreetingRequest.class, String.class, lambdaClient);
+        var result = runner.run(new GreetingRequest("TestUser"));
+        assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
+    }
+
     @Test
     void testGenericTypesExample() {
         var runner = CloudDurableTestRunner.create(

sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java

@@ -22,6 +22,10 @@
 import software.amazon.lambda.durable.model.WaitForConditionResult;
 
 public interface DurableContext extends BaseContext {
+    static DurableContext getContext() {
+        return (DurableContext) BaseContext.getCurrentContext();
+    }
+
     /**
      * Executes a durable step with the given name and blocks until it completes.
      *

sdk/src/main/java/software/amazon/lambda/durable/ParallelDurableFuture.java

@@ -5,9 +5,10 @@
 import java.util.function.Function;
 import software.amazon.lambda.durable.config.ParallelBranchConfig;
 import software.amazon.lambda.durable.model.ParallelResult;
+import software.amazon.lambda.durable.model.SafeCloseable;
 
 /** User-facing context for managing parallel branch execution within a durable function. */
-public interface ParallelDurableFuture extends AutoCloseable, DurableFuture<ParallelResult> {
+public interface ParallelDurableFuture extends SafeCloseable, DurableFuture<ParallelResult> {
 
     /**
      * Registers and immediately starts a branch (respects maxConcurrency).
@@ -68,7 +69,4 @@ default <T> DurableFuture<T> branch(
      */
     <T> DurableFuture<T> branch(
             String name, TypeToken<T> resultType, Function<DurableContext, T> func, ParallelBranchConfig config);
-
-    /** Calls {@link #get()} if not already called. Guarantees that the context is closed. */
-    void close();
 }

sdk/src/main/java/software/amazon/lambda/durable/context/BaseContext.java

@@ -3,17 +3,36 @@
 package software.amazon.lambda.durable.context;
 
 import com.amazonaws.services.lambda.runtime.Context;
+import org.slf4j.Logger;
 import software.amazon.lambda.durable.DurableConfig;
 import software.amazon.lambda.durable.logging.DurableLogger;
 
-public interface BaseContext extends AutoCloseable {
+public interface BaseContext {
+    ThreadLocal<BaseContext> CONTEXT = new ThreadLocal<>();
+
+    /**
+     * Gets the current context (DurableContext or StepContext) for this thread.
+     *
+     * @return the current context or null if not set
+     */
+    static BaseContext getCurrentContext() {
+        return CONTEXT.get();
+    }
     /**
      * Gets a logger with additional information of the current execution context.
      *
      * @return a DurableLogger instance
      */
     DurableLogger getLogger();
 
+    /**
+     * Gets a logger with additional information of the current execution context.
+     *
+     * @param delegate the logger to wrap
+     * @return a DurableLogger instance
+     */
+    DurableLogger getLogger(Logger delegate);
+
     /**
      * Returns the AWS Lambda runtime context.
      *
@@ -46,7 +65,4 @@ public interface BaseContext extends AutoCloseable {
 
     /** Returns whether this context is currently in replay mode. */
     boolean isReplaying();
-
-    /** Closes this context. */
-    void close();
 }

sdk/src/main/java/software/amazon/lambda/durable/context/BaseContextImpl.java

@@ -3,11 +3,13 @@
 package software.amazon.lambda.durable.context;
 
 import com.amazonaws.services.lambda.runtime.Context;
+import org.slf4j.Logger;
 import software.amazon.lambda.durable.DurableConfig;
 import software.amazon.lambda.durable.execution.ExecutionManager;
 import software.amazon.lambda.durable.execution.ThreadType;
+import software.amazon.lambda.durable.logging.DurableLogger;
 
-public abstract class BaseContextImpl implements AutoCloseable, BaseContext {
+public abstract class BaseContextImpl implements BaseContext {
     private final ExecutionManager executionManager;
     private final DurableConfig durableConfig;
     private final Context lambdaContext;
@@ -109,4 +111,18 @@ public boolean isReplaying() {
     public void setExecutionMode() {
         this.isReplaying = false;
     }
+
+    /** Returns a durable logger for this context. */
+    public DurableLogger getLogger() {
+        return DurableLogger.INSTANCE;
+    }
+
+    /** Returns a durable logger for this context. */
+    public DurableLogger getLogger(Logger delegate) {
+        return new DurableLogger(delegate);
+    }
+
+    public static void setCurrentContext(BaseContext context) {
+        CONTEXT.set(context);
+    }
 }

sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java

@@ -10,7 +10,6 @@
 import java.util.function.BiConsumer;
 import java.util.function.BiFunction;
 import java.util.function.Function;
-import org.slf4j.LoggerFactory;
 import software.amazon.lambda.durable.DurableCallbackFuture;
 import software.amazon.lambda.durable.DurableConfig;
 import software.amazon.lambda.durable.DurableContext;
@@ -32,7 +31,6 @@
 import software.amazon.lambda.durable.execution.OperationIdGenerator;
 import software.amazon.lambda.durable.execution.SuspendExecutionException;
 import software.amazon.lambda.durable.execution.ThreadType;
-import software.amazon.lambda.durable.logging.DurableLogger;
 import software.amazon.lambda.durable.model.MapResult;
 import software.amazon.lambda.durable.model.OperationIdentifier;
 import software.amazon.lambda.durable.model.OperationSubType;
@@ -62,7 +60,6 @@ public class DurableContextImpl extends BaseContextImpl implements DurableContex
     private final OperationIdGenerator operationIdGenerator;
     private final DurableContextImpl parentContext;
     private final boolean isVirtual;
-    private volatile DurableLogger logger;
 
     /** Shared initialization — sets all fields. */
     private DurableContextImpl(
@@ -430,30 +427,6 @@ private static <T> T executeRetryLoop(
     }
 
     // =============== accessors ================
-    @Override
-    public DurableLogger getLogger() {
-        // lazy initialize logger
-        if (logger == null) {
-            synchronized (this) {
-                if (logger == null) {
-                    logger = new DurableLogger(LoggerFactory.getLogger(DurableContext.class), this);
-                }
-            }
-        }
-        return logger;
-    }
-
-    /**
-     * Clears the logger's thread properties. Called during context destruction to prevent memory leaks and ensure clean
-     * state for subsequent executions.
-     */
-    @Override
-    public void close() {
-        if (logger != null) {
-            logger.close();
-        }
-    }
-
     /**
      * Get the next operationId. Returns a globally unique operation ID by hashing a sequential operation counter. For
      * root contexts, the counter value is hashed directly (e.g. "1", "2", "3"). For child contexts, the values are

sdk/src/main/java/software/amazon/lambda/durable/context/StepContextImpl.java

@@ -3,12 +3,10 @@
 package software.amazon.lambda.durable.context;
 
 import com.amazonaws.services.lambda.runtime.Context;
-import org.slf4j.LoggerFactory;
 import software.amazon.lambda.durable.DurableConfig;
 import software.amazon.lambda.durable.StepContext;
 import software.amazon.lambda.durable.execution.ExecutionManager;
 import software.amazon.lambda.durable.execution.ThreadType;
-import software.amazon.lambda.durable.logging.DurableLogger;
 
 /**
  * Context available inside a step operation's user function.
@@ -17,7 +15,6 @@
  * {@link BaseContext} for thread lifecycle management.
  */
 public class StepContextImpl extends BaseContextImpl implements StepContext {
-    private volatile DurableLogger logger;
     private final int attempt;
 
     /**
@@ -46,25 +43,4 @@ protected StepContextImpl(
     public int getAttempt() {
         return attempt;
     }
-
-    @Override
-    public DurableLogger getLogger() {
-        // lazy initialize logger
-        if (logger == null) {
-            synchronized (this) {
-                if (logger == null) {
-                    logger = new DurableLogger(LoggerFactory.getLogger(StepContext.class), this);
-                }
-            }
-        }
-        return logger;
-    }
-
-    /** Closes the logger for this context. */
-    @Override
-    public void close() {
-        if (logger != null) {
-            logger.close();
-        }
-    }
 }

sdk/src/main/java/software/amazon/lambda/durable/execution/DurableExecutor.java

@@ -22,6 +22,7 @@
 import software.amazon.lambda.durable.exception.DurableOperationException;
 import software.amazon.lambda.durable.exception.IllegalDurableOperationException;
 import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException;
+import software.amazon.lambda.durable.logging.DurableLogger;
 import software.amazon.lambda.durable.model.DurableExecutionInput;
 import software.amazon.lambda.durable.model.DurableExecutionOutput;
 import software.amazon.lambda.durable.plugin.InvocationEndInfo;
@@ -69,9 +70,10 @@ public static <I, O> DurableExecutionOutput execute(
 
                         var userInput = extractUserInput(
                                 executionManager.getExecutionOperation(), config.getSerDes(), inputType);
-                        // use try-with-resources to clear logger properties
-                        try (var context =
-                                DurableContextImpl.createRootContext(executionManager, config, lambdaContext)) {
+                        var context = DurableContextImpl.createRootContext(executionManager, config, lambdaContext);
+                        DurableContextImpl.setCurrentContext(context);
+                        // use a try-with-resources to clear logger properties
+                        try (var ignored = DurableLogger.attachContext()) {
                             return handler.apply(userInput, context);
                         }
                     },

sdk/src/main/java/software/amazon/lambda/durable/execution/ExecutionManager.java

@@ -24,6 +24,7 @@
 import software.amazon.lambda.durable.DurableConfig;
 import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException;
 import software.amazon.lambda.durable.model.DurableExecutionInput;
+import software.amazon.lambda.durable.model.SafeCloseable;
 import software.amazon.lambda.durable.operation.BaseDurableOperation;
 
 /**
@@ -47,7 +48,7 @@
  *
  * @see InternalExecutor
  */
-public class ExecutionManager implements AutoCloseable {
+public class ExecutionManager implements SafeCloseable {
 
     private static final Logger logger = LoggerFactory.getLogger(ExecutionManager.class);