fix: align log metadata with other SDKs

Author: zhongkechen

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

@@ -6,7 +6,7 @@
 import software.amazon.lambda.durable.DurableConfig;
 import software.amazon.lambda.durable.logging.DurableLogger;
 
-public interface BaseContext extends AutoCloseable {
+public interface BaseContext {
     /**
      * Gets a logger with additional information of the current execution context.
      *
@@ -46,7 +46,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

@@ -6,8 +6,9 @@
 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 +110,8 @@ public boolean isReplaying() {
     public void setExecutionMode() {
         this.isReplaying = false;
     }
+
+    public DurableLogger getLogger() {
+        return DurableLogger.INSTANCE;
+    }
 }

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,9 @@ 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);
+                        // use a try-with-resources to clear logger properties
+                        try (var ignored = DurableLogger.attachContext(context)) {
                             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);
 

sdk/src/main/java/software/amazon/lambda/durable/logging/DurableLogger.java

@@ -3,39 +3,66 @@
 package software.amazon.lambda.durable.logging;
 
 import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import org.slf4j.MDC;
 import software.amazon.lambda.durable.DurableContext;
 import software.amazon.lambda.durable.StepContext;
-import software.amazon.lambda.durable.context.BaseContextImpl;
+import software.amazon.lambda.durable.context.BaseContext;
+import software.amazon.lambda.durable.model.SafeCloseable;
 
 /**
  * Logger wrapper that adds durable execution context to log entries via MDC and optionally suppresses logs during
  * replay.
  */
 public class DurableLogger {
-    static final String MDC_EXECUTION_ARN = "durableExecutionArn";
+    static final String MDC_DURABLE_EXECUTION_ARN = "durableExecutionArn";
+    static final String MDC_EXECUTION_ARN = "executionArn";
     static final String MDC_REQUEST_ID = "requestId";
     static final String MDC_OPERATION_ID = "operationId";
     static final String MDC_CONTEXT_ID = "contextId";
     static final String MDC_OPERATION_NAME = "operationName";
     static final String MDC_CONTEXT_NAME = "contextName";
     static final String MDC_ATTEMPT = "attempt";
 
+    public static final DurableLogger INSTANCE = new DurableLogger(LoggerFactory.getLogger(DurableLogger.class));
+    private static final SafeCloseable AUTO_CLOSER = DurableLogger::detachContext;
+    private static final ThreadLocal<BaseContext> CONTEXT = new ThreadLocal<>();
+
     private final Logger delegate;
-    private final BaseContextImpl context;
 
     /**
      * Creates a DurableLogger wrapping the given SLF4J logger with execution context MDC entries.
      *
      * @param delegate the SLF4J logger to wrap
-     * @param context the durable execution context providing MDC values
      */
-    public DurableLogger(Logger delegate, BaseContextImpl context) {
+    DurableLogger(Logger delegate) {
         this.delegate = delegate;
-        this.context = context;
+    }
+
+    public static SafeCloseable attachContext(BaseContext context) {
+        CONTEXT.set(context);
+        if (context != null) {
+            injectMdcProperties(context);
+        }
+        return AUTO_CLOSER;
+    }
+
+    public static void detachContext() {
+        if (CONTEXT.get() != null) {
+            CONTEXT.remove();
+            MDC.clear();
+        }
+    }
+
+    private static void injectMdcProperties(BaseContext context) {
+        var config = context.getDurableConfig().getLoggerConfig();
 
         // execution arn
-        MDC.put(MDC_EXECUTION_ARN, context.getExecutionArn());
+        if (config.oldKeyNames()) {
+            MDC.put(MDC_DURABLE_EXECUTION_ARN, context.getExecutionArn());
+        } else {
+            MDC.put(MDC_EXECUTION_ARN, context.getExecutionArn());
+        }
 
         // lambda request id
         var requestId =
@@ -47,10 +74,18 @@ public DurableLogger(Logger delegate, BaseContextImpl context) {
         if (context instanceof DurableContext) {
             // context thread - context id and name
             if (context.getContextId() != null) {
-                MDC.put(MDC_CONTEXT_ID, context.getContextId());
+                if (config.oldKeyNames()) {
+                    MDC.put(MDC_CONTEXT_ID, context.getContextId());
+                } else {
+                    MDC.put(MDC_OPERATION_ID, context.getContextId());
+                }
             }
             if (context.getContextName() != null) {
-                MDC.put(MDC_CONTEXT_NAME, context.getContextName());
+                if (config.oldKeyNames()) {
+                    MDC.put(MDC_CONTEXT_NAME, context.getContextName());
+                } else {
+                    MDC.put(MDC_OPERATION_NAME, context.getContextName());
+                }
             }
         } else if (context instanceof StepContext stepContext) {
             // In step context, context id is the operation id, context name is the operation name
@@ -63,11 +98,6 @@ public DurableLogger(Logger delegate, BaseContextImpl context) {
         }
     }
 
-    /** Clears all MDC entries. User set MDC entries will also be removed as the thread will not be used anymore. */
-    public void close() {
-        MDC.clear();
-    }
-
     public void trace(String format, Object... args) {
         log(() -> delegate.trace(format, args));
     }
@@ -92,13 +122,13 @@ public void error(String message, Throwable t) {
         log(() -> delegate.error(message, t));
     }
 
-    private boolean shouldSuppress() {
-        return context.getDurableConfig().getLoggerConfig().suppressReplayLogs()
-                && context.getExecutionManager().isReplaying();
+    private boolean shouldSuppress(BaseContext context) {
+        return context.getDurableConfig().getLoggerConfig().suppressReplayLogs() && context.isReplaying();
     }
 
     private void log(Runnable logAction) {
-        if (!shouldSuppress()) {
+        var threadLocalContext = CONTEXT.get();
+        if (threadLocalContext == null || !shouldSuppress(threadLocalContext)) {
             logAction.run();
         }
     }

sdk/src/main/java/software/amazon/lambda/durable/logging/LoggerConfig.java

@@ -3,15 +3,15 @@
 package software.amazon.lambda.durable.logging;
 
 /** Configuration for DurableLogger behavior. */
-public record LoggerConfig(boolean suppressReplayLogs) {
+public record LoggerConfig(boolean suppressReplayLogs, boolean oldKeyNames) {
 
     /** Default configuration: suppress logs during replay. */
     public static LoggerConfig defaults() {
-        return new LoggerConfig(true);
+        return new LoggerConfig(true, false);
     }
 
     /** Configuration that allows logs during replay. */
     public static LoggerConfig withReplayLogging() {
-        return new LoggerConfig(false);
+        return new LoggerConfig(false, false);
     }
 }

sdk/src/main/java/software/amazon/lambda/durable/model/SafeCloseable.java

@@ -0,0 +1,8 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package software.amazon.lambda.durable.model;
+
+/** A resource that can be safely closed, similar to {@link java.io.Closeable} but without throwing IOException. */
+public interface SafeCloseable extends AutoCloseable {
+    void close();
+}