Skip to content

Commit d85611a

Browse files
zhongkechenzhongkechen
committed
fix: align log metadata with other SDKs
1 parent da837f8 commit d85611a

12 files changed

Lines changed: 327 additions & 204 deletions

File tree

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

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
import software.amazon.lambda.durable.DurableConfig;
77
import software.amazon.lambda.durable.logging.DurableLogger;
88

9-
public interface BaseContext extends AutoCloseable {
9+
public interface BaseContext {
1010
/**
1111
* Gets a logger with additional information of the current execution context.
1212
*
@@ -46,7 +46,4 @@ public interface BaseContext extends AutoCloseable {
4646

4747
/** Returns whether this context is currently in replay mode. */
4848
boolean isReplaying();
49-
50-
/** Closes this context. */
51-
void close();
5249
}

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

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,9 @@
66
import software.amazon.lambda.durable.DurableConfig;
77
import software.amazon.lambda.durable.execution.ExecutionManager;
88
import software.amazon.lambda.durable.execution.ThreadType;
9+
import software.amazon.lambda.durable.logging.DurableLogger;
910

10-
public abstract class BaseContextImpl implements AutoCloseable, BaseContext {
11+
public abstract class BaseContextImpl implements BaseContext {
1112
private final ExecutionManager executionManager;
1213
private final DurableConfig durableConfig;
1314
private final Context lambdaContext;
@@ -109,4 +110,8 @@ public boolean isReplaying() {
109110
public void setExecutionMode() {
110111
this.isReplaying = false;
111112
}
113+
114+
public DurableLogger getLogger() {
115+
return DurableLogger.INSTANCE;
116+
}
112117
}

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

Lines changed: 0 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@
1010
import java.util.function.BiConsumer;
1111
import java.util.function.BiFunction;
1212
import java.util.function.Function;
13-
import org.slf4j.LoggerFactory;
1413
import software.amazon.lambda.durable.DurableCallbackFuture;
1514
import software.amazon.lambda.durable.DurableConfig;
1615
import software.amazon.lambda.durable.DurableContext;
@@ -32,7 +31,6 @@
3231
import software.amazon.lambda.durable.execution.OperationIdGenerator;
3332
import software.amazon.lambda.durable.execution.SuspendExecutionException;
3433
import software.amazon.lambda.durable.execution.ThreadType;
35-
import software.amazon.lambda.durable.logging.DurableLogger;
3634
import software.amazon.lambda.durable.model.MapResult;
3735
import software.amazon.lambda.durable.model.OperationIdentifier;
3836
import software.amazon.lambda.durable.model.OperationSubType;
@@ -62,7 +60,6 @@ public class DurableContextImpl extends BaseContextImpl implements DurableContex
6260
private final OperationIdGenerator operationIdGenerator;
6361
private final DurableContextImpl parentContext;
6462
private final boolean isVirtual;
65-
private volatile DurableLogger logger;
6663

6764
/** Shared initialization — sets all fields. */
6865
private DurableContextImpl(
@@ -430,30 +427,6 @@ private static <T> T executeRetryLoop(
430427
}
431428

432429
// =============== accessors ================
433-
@Override
434-
public DurableLogger getLogger() {
435-
// lazy initialize logger
436-
if (logger == null) {
437-
synchronized (this) {
438-
if (logger == null) {
439-
logger = new DurableLogger(LoggerFactory.getLogger(DurableContext.class), this);
440-
}
441-
}
442-
}
443-
return logger;
444-
}
445-
446-
/**
447-
* Clears the logger's thread properties. Called during context destruction to prevent memory leaks and ensure clean
448-
* state for subsequent executions.
449-
*/
450-
@Override
451-
public void close() {
452-
if (logger != null) {
453-
logger.close();
454-
}
455-
}
456-
457430
/**
458431
* Get the next operationId. Returns a globally unique operation ID by hashing a sequential operation counter. For
459432
* 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‎

Lines changed: 0 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -3,12 +3,10 @@
33
package software.amazon.lambda.durable.context;
44

55
import com.amazonaws.services.lambda.runtime.Context;
6-
import org.slf4j.LoggerFactory;
76
import software.amazon.lambda.durable.DurableConfig;
87
import software.amazon.lambda.durable.StepContext;
98
import software.amazon.lambda.durable.execution.ExecutionManager;
109
import software.amazon.lambda.durable.execution.ThreadType;
11-
import software.amazon.lambda.durable.logging.DurableLogger;
1210

1311
/**
1412
* Context available inside a step operation's user function.
@@ -17,7 +15,6 @@
1715
* {@link BaseContext} for thread lifecycle management.
1816
*/
1917
public class StepContextImpl extends BaseContextImpl implements StepContext {
20-
private volatile DurableLogger logger;
2118
private final int attempt;
2219

2320
/**
@@ -46,25 +43,4 @@ protected StepContextImpl(
4643
public int getAttempt() {
4744
return attempt;
4845
}
49-
50-
@Override
51-
public DurableLogger getLogger() {
52-
// lazy initialize logger
53-
if (logger == null) {
54-
synchronized (this) {
55-
if (logger == null) {
56-
logger = new DurableLogger(LoggerFactory.getLogger(StepContext.class), this);
57-
}
58-
}
59-
}
60-
return logger;
61-
}
62-
63-
/** Closes the logger for this context. */
64-
@Override
65-
public void close() {
66-
if (logger != null) {
67-
logger.close();
68-
}
69-
}
7046
}

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

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import software.amazon.lambda.durable.exception.DurableOperationException;
2323
import software.amazon.lambda.durable.exception.IllegalDurableOperationException;
2424
import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException;
25+
import software.amazon.lambda.durable.logging.DurableLogger;
2526
import software.amazon.lambda.durable.model.DurableExecutionInput;
2627
import software.amazon.lambda.durable.model.DurableExecutionOutput;
2728
import software.amazon.lambda.durable.plugin.InvocationEndInfo;
@@ -69,9 +70,9 @@ public static <I, O> DurableExecutionOutput execute(
6970

7071
var userInput = extractUserInput(
7172
executionManager.getExecutionOperation(), config.getSerDes(), inputType);
72-
// use try-with-resources to clear logger properties
73-
try (var context =
74-
DurableContextImpl.createRootContext(executionManager, config, lambdaContext)) {
73+
var context = DurableContextImpl.createRootContext(executionManager, config, lambdaContext);
74+
// use a try-with-resources to clear logger properties
75+
try (DurableLogger.DurableLoggerAutoCloser ignored = DurableLogger.attachContext(context)) {
7576
return handler.apply(userInput, context);
7677
}
7778
},

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

Lines changed: 50 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -3,39 +3,68 @@
33
package software.amazon.lambda.durable.logging;
44

55
import org.slf4j.Logger;
6+
import org.slf4j.LoggerFactory;
67
import org.slf4j.MDC;
78
import software.amazon.lambda.durable.DurableContext;
89
import software.amazon.lambda.durable.StepContext;
9-
import software.amazon.lambda.durable.context.BaseContextImpl;
10+
import software.amazon.lambda.durable.context.BaseContext;
1011

1112
/**
1213
* Logger wrapper that adds durable execution context to log entries via MDC and optionally suppresses logs during
1314
* replay.
1415
*/
1516
public class DurableLogger {
16-
static final String MDC_EXECUTION_ARN = "durableExecutionArn";
17+
static final String MDC_DURABLE_EXECUTION_ARN = "durableExecutionArn";
18+
static final String MDC_EXECUTION_ARN = "executionArn";
1719
static final String MDC_REQUEST_ID = "requestId";
1820
static final String MDC_OPERATION_ID = "operationId";
1921
static final String MDC_CONTEXT_ID = "contextId";
2022
static final String MDC_OPERATION_NAME = "operationName";
2123
static final String MDC_CONTEXT_NAME = "contextName";
2224
static final String MDC_ATTEMPT = "attempt";
2325

26+
public static final DurableLogger INSTANCE = new DurableLogger(LoggerFactory.getLogger(DurableLogger.class));
27+
private static final DurableLoggerAutoCloser AUTO_CLOSER = new DurableLoggerAutoCloser();
28+
private static final ThreadLocal<BaseContext> CONTEXT = new ThreadLocal<>();
29+
2430
private final Logger delegate;
25-
private final BaseContextImpl context;
31+
32+
public static class DurableLoggerAutoCloser implements AutoCloseable {
33+
@Override
34+
public void close() {
35+
DurableLogger.detachContext();
36+
}
37+
}
2638

2739
/**
2840
* Creates a DurableLogger wrapping the given SLF4J logger with execution context MDC entries.
2941
*
3042
* @param delegate the SLF4J logger to wrap
31-
* @param context the durable execution context providing MDC values
3243
*/
33-
public DurableLogger(Logger delegate, BaseContextImpl context) {
44+
DurableLogger(Logger delegate) {
3445
this.delegate = delegate;
35-
this.context = context;
46+
}
47+
48+
public static DurableLoggerAutoCloser attachContext(BaseContext context) {
49+
CONTEXT.set(context);
50+
injectMdcProperties(context);
51+
return AUTO_CLOSER;
52+
}
53+
54+
public static void detachContext() {
55+
CONTEXT.remove();
56+
MDC.clear();
57+
}
58+
59+
private static void injectMdcProperties(BaseContext context) {
60+
var config = context.getDurableConfig().getLoggerConfig();
3661

3762
// execution arn
38-
MDC.put(MDC_EXECUTION_ARN, context.getExecutionArn());
63+
if (config.oldKeyNames()) {
64+
MDC.put(MDC_DURABLE_EXECUTION_ARN, context.getExecutionArn());
65+
} else {
66+
MDC.put(MDC_EXECUTION_ARN, context.getExecutionArn());
67+
}
3968

4069
// lambda request id
4170
var requestId =
@@ -47,10 +76,18 @@ public DurableLogger(Logger delegate, BaseContextImpl context) {
4776
if (context instanceof DurableContext) {
4877
// context thread - context id and name
4978
if (context.getContextId() != null) {
50-
MDC.put(MDC_CONTEXT_ID, context.getContextId());
79+
if (config.oldKeyNames()) {
80+
MDC.put(MDC_CONTEXT_ID, context.getContextId());
81+
} else {
82+
MDC.put(MDC_OPERATION_ID, context.getContextId());
83+
}
5184
}
5285
if (context.getContextName() != null) {
53-
MDC.put(MDC_CONTEXT_NAME, context.getContextName());
86+
if (config.oldKeyNames()) {
87+
MDC.put(MDC_CONTEXT_NAME, context.getContextName());
88+
} else {
89+
MDC.put(MDC_OPERATION_NAME, context.getContextName());
90+
}
5491
}
5592
} else if (context instanceof StepContext stepContext) {
5693
// In step context, context id is the operation id, context name is the operation name
@@ -63,11 +100,6 @@ public DurableLogger(Logger delegate, BaseContextImpl context) {
63100
}
64101
}
65102

66-
/** Clears all MDC entries. User set MDC entries will also be removed as the thread will not be used anymore. */
67-
public void close() {
68-
MDC.clear();
69-
}
70-
71103
public void trace(String format, Object... args) {
72104
log(() -> delegate.trace(format, args));
73105
}
@@ -92,13 +124,13 @@ public void error(String message, Throwable t) {
92124
log(() -> delegate.error(message, t));
93125
}
94126

95-
private boolean shouldSuppress() {
96-
return context.getDurableConfig().getLoggerConfig().suppressReplayLogs()
97-
&& context.getExecutionManager().isReplaying();
127+
private boolean shouldSuppress(BaseContext context) {
128+
return context.getDurableConfig().getLoggerConfig().suppressReplayLogs() && context.isReplaying();
98129
}
99130

100131
private void log(Runnable logAction) {
101-
if (!shouldSuppress()) {
132+
var threadLocalContext = CONTEXT.get();
133+
if (threadLocalContext == null || !shouldSuppress(threadLocalContext)) {
102134
logAction.run();
103135
}
104136
}

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,15 +3,15 @@
33
package software.amazon.lambda.durable.logging;
44

55
/** Configuration for DurableLogger behavior. */
6-
public record LoggerConfig(boolean suppressReplayLogs) {
6+
public record LoggerConfig(boolean suppressReplayLogs, boolean oldKeyNames) {
77

88
/** Default configuration: suppress logs during replay. */
99
public static LoggerConfig defaults() {
10-
return new LoggerConfig(true);
10+
return new LoggerConfig(true, false);
1111
}
1212

1313
/** Configuration that allows logs during replay. */
1414
public static LoggerConfig withReplayLogging() {
15-
return new LoggerConfig(false);
15+
return new LoggerConfig(false, false);
1616
}
1717
}

‎sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException;
3333
import software.amazon.lambda.durable.execution.SuspendExecutionException;
3434
import software.amazon.lambda.durable.execution.ThreadType;
35+
import software.amazon.lambda.durable.logging.DurableLogger;
3536
import software.amazon.lambda.durable.model.DeserializedOperationResult;
3637
import software.amazon.lambda.durable.model.OperationIdentifier;
3738
import software.amazon.lambda.durable.util.ExceptionHelper;
@@ -130,7 +131,8 @@ private void executeChildContext() {
130131
// When this child is part of a ConcurrencyOperation (parentOperation != null),
131132
// we notify the parent BEFORE closing the child context. This ensures the parent
132133
// can trigger the next queued branch while the current child context is still valid.
133-
try (var childContext = getContext().createChildContext(contextId, getName(), isVirtual)) {
134+
var childContext = getContext().createChildContext(contextId, getName(), isVirtual);
135+
try (DurableLogger.DurableLoggerAutoCloser ignored = DurableLogger.attachContext(childContext)) {
134136
try {
135137
T result = function.apply(childContext);
136138

‎sdk/src/main/java/software/amazon/lambda/durable/operation/SerializableDurableOperation.java‎

Lines changed: 0 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -8,9 +8,7 @@
88
import software.amazon.lambda.durable.DurableFuture;
99
import software.amazon.lambda.durable.TypeToken;
1010
import software.amazon.lambda.durable.context.DurableContextImpl;
11-
import software.amazon.lambda.durable.exception.IllegalDurableOperationException;
1211
import software.amazon.lambda.durable.exception.SerDesException;
13-
import software.amazon.lambda.durable.execution.ThreadType;
1412
import software.amazon.lambda.durable.model.OperationIdentifier;
1513
import software.amazon.lambda.durable.serde.SerDes;
1614
import software.amazon.lambda.durable.util.ExceptionHelper;
@@ -77,22 +75,6 @@ protected SerializableDurableOperation(
7775
this.resultSerDes = resultSerDes;
7876
}
7977

80-
/**
81-
* Checks if it's called from a Step.
82-
*
83-
* @throws IllegalDurableOperationException if it's in a step
84-
*/
85-
private void validateCurrentThreadType() {
86-
ThreadType current = getCurrentThreadContext().threadType();
87-
if (current == ThreadType.STEP) {
88-
var message = String.format(
89-
"Nested %s operation is not supported on %s from within a %s execution.",
90-
getType(), getName(), current);
91-
// terminate execution and throw the exception
92-
throw terminateExecutionWithIllegalDurableOperationException(message);
93-
}
94-
}
95-
9678
/**
9779
* Deserializes a result string into the operation's result type.
9880
*

‎sdk/src/main/java/software/amazon/lambda/durable/operation/StepOperation.java‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException;
2323
import software.amazon.lambda.durable.execution.SuspendExecutionException;
2424
import software.amazon.lambda.durable.execution.ThreadType;
25+
import software.amazon.lambda.durable.logging.DurableLogger;
2526
import software.amazon.lambda.durable.model.OperationIdentifier;
2627
import software.amazon.lambda.durable.util.ExceptionHelper;
2728

@@ -104,7 +105,9 @@ private void executeStepLogic(int attempt) {
104105
// use a try-with-resources to
105106
// - add thread id/type to thread local when the step starts
106107
// - clear logger properties when the step finishes
107-
try (StepContext stepContext = getContext().createStepContext(getOperationId(), getName(), attempt)) {
108+
StepContext stepContext = getContext().createStepContext(getOperationId(), getName(), attempt);
109+
110+
try (DurableLogger.DurableLoggerAutoCloser ignored = DurableLogger.attachContext(stepContext)) {
108111
try {
109112
checkpointStarted();
110113

0 commit comments

Comments
 (0)