Skip to content

Commit 2525fc8

Browse files
authored
[alc] ContextLogger changes in RIC (#436)
* [alc] ContextLogger changes in RIC --------- Co-authored-by: Daniel Torok <[email protected]>
1 parent e87c680 commit 2525fc8

22 files changed

+532
-36
lines changed

aws-lambda-java-runtime-interface-client/pom.xml

+1-1
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,7 @@
5353
<dependency>
5454
<groupId>com.amazonaws</groupId>
5555
<artifactId>aws-lambda-java-core</artifactId>
56-
<version>1.2.2</version>
56+
<version>1.2.3</version>
5757
</dependency>
5858
<dependency>
5959
<groupId>com.amazonaws</groupId>

aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/AWSLambda.java

+8-1
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@
1616
import com.amazonaws.services.lambda.runtime.api.client.runtimeapi.LambdaRuntimeClient;
1717
import com.amazonaws.services.lambda.runtime.api.client.util.LambdaOutputStream;
1818
import com.amazonaws.services.lambda.runtime.api.client.util.UnsafeUtil;
19+
import com.amazonaws.services.lambda.runtime.logging.LogFormat;
20+
import com.amazonaws.services.lambda.runtime.logging.LogLevel;
1921
import com.amazonaws.services.lambda.runtime.serialization.PojoSerializer;
2022
import com.amazonaws.services.lambda.runtime.serialization.factories.GsonFactory;
2123
import com.amazonaws.services.lambda.runtime.serialization.factories.JacksonFactory;
@@ -187,7 +189,12 @@ public static void main(String[] args) {
187189

188190
private static void startRuntime(String handler) {
189191
try (LogSink logSink = createLogSink()) {
190-
startRuntime(handler, new LambdaContextLogger(logSink));
192+
LambdaLogger logger = new LambdaContextLogger(
193+
logSink,
194+
LogLevel.fromString(LambdaEnvironment.LAMBDA_LOG_LEVEL),
195+
LogFormat.fromString(LambdaEnvironment.LAMBDA_LOG_FORMAT)
196+
);
197+
startRuntime(handler, logger);
191198
} catch (Throwable t) {
192199
throw new Error(t);
193200
}

aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/EventHandlerLoader.java

+20
Original file line numberDiff line numberDiff line change
@@ -4,13 +4,15 @@
44

55
import com.amazonaws.services.lambda.runtime.ClientContext;
66
import com.amazonaws.services.lambda.runtime.Context;
7+
import com.amazonaws.services.lambda.runtime.LambdaLogger;
78
import com.amazonaws.services.lambda.runtime.LambdaRuntimeInternal;
89
import com.amazonaws.services.lambda.runtime.RequestHandler;
910
import com.amazonaws.services.lambda.runtime.RequestStreamHandler;
1011
import com.amazonaws.services.lambda.runtime.api.client.LambdaRequestHandler.UserFaultHandler;
1112
import com.amazonaws.services.lambda.runtime.api.client.api.LambdaClientContext;
1213
import com.amazonaws.services.lambda.runtime.api.client.api.LambdaCognitoIdentity;
1314
import com.amazonaws.services.lambda.runtime.api.client.api.LambdaContext;
15+
import com.amazonaws.services.lambda.runtime.api.client.logging.LambdaContextLogger;
1416
import com.amazonaws.services.lambda.runtime.api.client.runtimeapi.InvocationRequest;
1517
import com.amazonaws.services.lambda.runtime.api.client.util.UnsafeUtil;
1618
import com.amazonaws.services.lambda.runtime.serialization.PojoSerializer;
@@ -844,6 +846,22 @@ private void safeAddRequestIdToLog4j(String log4jContextClassName,
844846
}
845847
}
846848

849+
/**
850+
* Passes the LambdaContext to the logger so that the JSON formatter can include the requestId.
851+
*
852+
* We do casting here because both the LambdaRuntime and the LambdaLogger is in the core package,
853+
* and the setLambdaContext(context) is a method we don't want to publish for customers. That method is
854+
* only implemented on the internal LambdaContextLogger, so we check and cast to be able to call it.
855+
* @param context the LambdaContext
856+
*/
857+
private void safeAddContextToLambdaLogger(LambdaContext context) {
858+
LambdaLogger logger = com.amazonaws.services.lambda.runtime.LambdaRuntime.getLogger();
859+
if (logger instanceof LambdaContextLogger) {
860+
LambdaContextLogger contextLogger = (LambdaContextLogger) logger;
861+
contextLogger.setLambdaContext(context);
862+
}
863+
}
864+
847865
public ByteArrayOutputStream call(InvocationRequest request) throws Error, Exception {
848866
output.reset();
849867

@@ -871,6 +889,8 @@ public ByteArrayOutputStream call(InvocationRequest request) throws Error, Excep
871889
clientContext
872890
);
873891

892+
safeAddContextToLambdaLogger(context);
893+
874894
if (LambdaRuntimeInternal.getUseLog4jAppender()) {
875895
safeAddRequestIdToLog4j("org.apache.log4j.MDC", request, Object.class);
876896
safeAddRequestIdToLog4j("org.apache.logging.log4j.ThreadContext", request, String.class);

aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/LambdaEnvironment.java

+2
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,8 @@ public class LambdaEnvironment {
1212
public static final int MEMORY_LIMIT = parseInt(ENV_READER.getEnvOrDefault(AWS_LAMBDA_FUNCTION_MEMORY_SIZE, "128"));
1313
public static final String LOG_GROUP_NAME = ENV_READER.getEnv(AWS_LAMBDA_LOG_GROUP_NAME);
1414
public static final String LOG_STREAM_NAME = ENV_READER.getEnv(AWS_LAMBDA_LOG_STREAM_NAME);
15+
public static final String LAMBDA_LOG_LEVEL = ENV_READER.getEnvOrDefault(AWS_LAMBDA_LOG_LEVEL, "UNDEFINED");
16+
public static final String LAMBDA_LOG_FORMAT = ENV_READER.getEnvOrDefault(AWS_LAMBDA_LOG_FORMAT, "TEXT");
1517
public static final String FUNCTION_NAME = ENV_READER.getEnv(AWS_LAMBDA_FUNCTION_NAME);
1618
public static final String FUNCTION_VERSION = ENV_READER.getEnv(AWS_LAMBDA_FUNCTION_VERSION);
1719
}

aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/ReservedRuntimeEnvironmentVariables.java

+10
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,16 @@ public interface ReservedRuntimeEnvironmentVariables {
6666
*/
6767
String AWS_LAMBDA_LOG_STREAM_NAME = "AWS_LAMBDA_LOG_STREAM_NAME";
6868

69+
/**
70+
* The logging level set for the function.
71+
*/
72+
String AWS_LAMBDA_LOG_LEVEL = "AWS_LAMBDA_LOG_LEVEL";
73+
74+
/**
75+
* The logging format set for the function.
76+
*/
77+
String AWS_LAMBDA_LOG_FORMAT = "AWS_LAMBDA_LOG_FORMAT";
78+
6979
/**
7080
* Access key id obtained from the function's execution role.
7181
*/
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
/* Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved. */
2+
3+
package com.amazonaws.services.lambda.runtime.api.client.logging;
4+
5+
import static java.nio.charset.StandardCharsets.UTF_8;
6+
7+
import com.amazonaws.services.lambda.runtime.api.client.api.LambdaContext;
8+
import com.amazonaws.services.lambda.runtime.LambdaLogger;
9+
import com.amazonaws.services.lambda.runtime.logging.LogFormat;
10+
import com.amazonaws.services.lambda.runtime.logging.LogLevel;
11+
12+
/**
13+
* Provides default implementation of the convenience logger functions.
14+
* When extending AbstractLambdaLogger, only one function has to be overridden:
15+
* void logMessage(byte[] message, LogLevel logLevel);
16+
*/
17+
public abstract class AbstractLambdaLogger implements LambdaLogger {
18+
private final LogFiltering logFiltering;
19+
private final LogFormatter logFormatter;
20+
protected final LogFormat logFormat;
21+
22+
public AbstractLambdaLogger(LogLevel logLevel, LogFormat logFormat) {
23+
this.logFiltering = new LogFiltering(logLevel);
24+
25+
this.logFormat = logFormat;
26+
if (logFormat == LogFormat.JSON) {
27+
logFormatter = new JsonLogFormatter();
28+
} else {
29+
logFormatter = new TextLogFormatter();
30+
}
31+
}
32+
33+
protected abstract void logMessage(byte[] message, LogLevel logLevel);
34+
35+
protected void logMessage(String message, LogLevel logLevel) {
36+
logMessage(message.getBytes(UTF_8), logLevel);
37+
}
38+
39+
@Override
40+
public void log(String message, LogLevel logLevel) {
41+
if (logFiltering.isEnabled(logLevel)) {
42+
this.logMessage(logFormatter.format(message, logLevel), logLevel);
43+
}
44+
}
45+
46+
@Override
47+
public void log(byte[] message, LogLevel logLevel) {
48+
if (logFiltering.isEnabled(logLevel)) {
49+
// there is no formatting for byte[] messages
50+
this.logMessage(message, logLevel);
51+
}
52+
}
53+
54+
@Override
55+
public void log(String message) {
56+
this.log(message, LogLevel.UNDEFINED);
57+
}
58+
59+
@Override
60+
public void log(byte[] message) {
61+
this.log(message, LogLevel.UNDEFINED);
62+
}
63+
64+
public void setLambdaContext(LambdaContext lambdaContext) {
65+
this.logFormatter.setLambdaContext(lambdaContext);
66+
}
67+
}

aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/logging/FrameType.java

+30-2
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,40 @@
22

33
package com.amazonaws.services.lambda.runtime.api.client.logging;
44

5-
public enum FrameType {
5+
import com.amazonaws.services.lambda.runtime.logging.LogLevel;
6+
import com.amazonaws.services.lambda.runtime.logging.LogFormat;
67

7-
LOG(0xa55a0003);
8+
/**
9+
* The first 4 bytes of the framing protocol is the Frame Type, that's made of a magic number (3 bytes) and 1 byte of flags.
10+
* +-----------------------+
11+
* | Frame Type - 4 bytes |
12+
* +-----------------------+
13+
* | a5 5a 00 | flgs |
14+
* + - - - - - + - - - - - +
15+
* \ bit |
16+
* | view|
17+
* +---------+ +
18+
* | |
19+
* v byte 3 v F - free
20+
* +-+-+-+-+-+-+-+-+ J - { JsonLog = 0, PlainTextLog = 1 }
21+
* |F|F|F|L|l|l|T|J| T - { NoTimeStamp = 0, TimeStampPresent = 1 }
22+
* +-+-+-+-+-+-+-+-+ Lll -> Log Level in 3-bit binary (L-> most significant bit)
23+
*/
24+
public class FrameType {
25+
private static final int LOG_MAGIC = 0xa55a0000;
26+
private static final int OFFSET_LOG_FORMAT = 0;
27+
private static final int OFFSET_TIMESTAMP_PRESENT = 1;
28+
private static final int OFFSET_LOG_LEVEL = 2;
829

930
private final int val;
1031

32+
public static int getValue(LogLevel logLevel, LogFormat logFormat) {
33+
return LOG_MAGIC |
34+
(logLevel.ordinal() << OFFSET_LOG_LEVEL) |
35+
(1 << OFFSET_TIMESTAMP_PRESENT) |
36+
(logFormat.ordinal() << OFFSET_LOG_FORMAT);
37+
}
38+
1139
FrameType(int val) {
1240
this.val = val;
1341
}

aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/logging/FramedTelemetryLogSink.java

+14-6
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,9 @@
99
import java.nio.ByteOrder;
1010
import java.time.Instant;
1111

12+
import com.amazonaws.services.lambda.runtime.logging.LogLevel;
13+
import com.amazonaws.services.lambda.runtime.logging.LogFormat;
14+
1215
/**
1316
* FramedTelemetryLogSink implements the logging contract between runtimes and the platform. It implements a simple
1417
* framing protocol so message boundaries can be determined. Each frame can be visualized as follows:
@@ -38,16 +41,21 @@ public FramedTelemetryLogSink(FileDescriptor fd) throws IOException {
3841
}
3942

4043
@Override
41-
public synchronized void log(byte[] message) {
44+
public synchronized void log(LogLevel logLevel, LogFormat logFormat, byte[] message) {
4245
try {
43-
writeFrame(message);
46+
writeFrame(logLevel, logFormat, message);
4447
} catch (IOException e) {
4548
e.printStackTrace();
4649
}
4750
}
4851

49-
private void writeFrame(byte[] message) throws IOException {
50-
updateHeader(message.length);
52+
@Override
53+
public void log(byte[] message) {
54+
log(LogLevel.UNDEFINED, LogFormat.TEXT, message);
55+
}
56+
57+
private void writeFrame(LogLevel logLevel, LogFormat logFormat, byte[] message) throws IOException {
58+
updateHeader(logLevel, logFormat, message.length);
5159
this.logOutputStream.write(this.headerBuf.array());
5260
this.logOutputStream.write(message);
5361
}
@@ -60,9 +68,9 @@ private long timestamp() {
6068
/**
6169
* Updates the header ByteBuffer with the provided length. The header comprises the frame type and message length.
6270
*/
63-
private void updateHeader(int length) {
71+
private void updateHeader(LogLevel logLevel, LogFormat logFormat, int length) {
6472
this.headerBuf.clear();
65-
this.headerBuf.putInt(FrameType.LOG.getValue());
73+
this.headerBuf.putInt(FrameType.getValue(logLevel, logFormat));
6674
this.headerBuf.putInt(length);
6775
this.headerBuf.putLong(timestamp());
6876
this.headerBuf.flip();
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
/* Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved. */
2+
3+
package com.amazonaws.services.lambda.runtime.api.client.logging;
4+
5+
import com.amazonaws.services.lambda.runtime.api.client.api.LambdaContext;
6+
import com.amazonaws.services.lambda.runtime.logging.LogLevel;
7+
import com.amazonaws.services.lambda.runtime.serialization.PojoSerializer;
8+
import com.amazonaws.services.lambda.runtime.serialization.factories.GsonFactory;
9+
10+
import java.io.ByteArrayOutputStream;
11+
import java.nio.charset.StandardCharsets;
12+
import java.time.LocalDateTime;
13+
import java.time.ZoneId;
14+
import java.time.format.DateTimeFormatter;
15+
16+
public class JsonLogFormatter implements LogFormatter {
17+
private final PojoSerializer<StructuredLogMessage> serializer = GsonFactory.getInstance().getSerializer(StructuredLogMessage.class);
18+
private LambdaContext lambdaContext;
19+
20+
private static final DateTimeFormatter dateFormatter =
21+
DateTimeFormatter
22+
.ofPattern("yyyy-MM-dd'T'HH:mm:ss.SSS'Z'")
23+
.withZone(ZoneId.of("UTC"));
24+
25+
@Override
26+
public String format(String message, LogLevel logLevel) {
27+
ByteArrayOutputStream stream = new ByteArrayOutputStream();
28+
StructuredLogMessage msg = createLogMessage(message, logLevel);
29+
serializer.toJson(msg, stream);
30+
stream.write('\n');
31+
return new String(stream.toByteArray(), StandardCharsets.UTF_8);
32+
}
33+
34+
private StructuredLogMessage createLogMessage(String message, LogLevel logLevel) {
35+
StructuredLogMessage msg = new StructuredLogMessage();
36+
msg.timestamp = dateFormatter.format(LocalDateTime.now());
37+
msg.message = message;
38+
msg.level = logLevel;
39+
40+
if (lambdaContext != null) {
41+
msg.AWSRequestId = lambdaContext.getAwsRequestId();
42+
}
43+
return msg;
44+
}
45+
46+
47+
/**
48+
* Function to set the context for every invocation.
49+
* This way the logger will be able to attach additional information to the log packet.
50+
*/
51+
@Override
52+
public void setLambdaContext(LambdaContext context) {
53+
this.lambdaContext = context;
54+
}
55+
}

aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/logging/LambdaContextLogger.java

+9-13
Original file line numberDiff line numberDiff line change
@@ -2,33 +2,29 @@
22

33
package com.amazonaws.services.lambda.runtime.api.client.logging;
44

5-
import com.amazonaws.services.lambda.runtime.LambdaLogger;
5+
import com.amazonaws.services.lambda.runtime.api.client.api.LambdaContext;
6+
import com.amazonaws.services.lambda.runtime.logging.LogFormat;
7+
import com.amazonaws.services.lambda.runtime.logging.LogLevel;
68

79
import static java.nio.charset.StandardCharsets.UTF_8;
810

9-
public class LambdaContextLogger implements LambdaLogger {
11+
public class LambdaContextLogger extends AbstractLambdaLogger {
1012
// If a null string is passed in, replace it with "null",
1113
// replicating the behavior of System.out.println(null);
1214
private static final byte[] NULL_BYTES_VALUE = "null".getBytes(UTF_8);
1315

1416
private final transient LogSink sink;
1517

16-
public LambdaContextLogger(LogSink sink) {
18+
public LambdaContextLogger(LogSink sink, LogLevel logLevel, LogFormat logFormat) {
19+
super(logLevel, logFormat);
1720
this.sink = sink;
1821
}
1922

20-
public void log(byte[] message) {
23+
@Override
24+
protected void logMessage(byte[] message, LogLevel logLevel) {
2125
if (message == null) {
2226
message = NULL_BYTES_VALUE;
2327
}
24-
sink.log(message);
25-
}
26-
27-
public void log(String message) {
28-
if (message == null) {
29-
this.log(NULL_BYTES_VALUE);
30-
} else {
31-
this.log(message.getBytes(UTF_8));
32-
}
28+
sink.log(logLevel, this.logFormat, message);
3329
}
3430
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
/* Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved. */
2+
3+
package com.amazonaws.services.lambda.runtime.api.client.logging;
4+
5+
import com.amazonaws.services.lambda.runtime.logging.LogLevel;
6+
7+
public class LogFiltering {
8+
private final LogLevel minimumLogLevel;
9+
10+
public LogFiltering(LogLevel minimumLogLevel) {
11+
this.minimumLogLevel = minimumLogLevel;
12+
}
13+
14+
boolean isEnabled(LogLevel logLevel) {
15+
return (logLevel == LogLevel.UNDEFINED || logLevel.ordinal() >= minimumLogLevel.ordinal());
16+
}
17+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
/* Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved. */
2+
3+
package com.amazonaws.services.lambda.runtime.api.client.logging;
4+
5+
import com.amazonaws.services.lambda.runtime.api.client.api.LambdaContext;
6+
import com.amazonaws.services.lambda.runtime.logging.LogLevel;
7+
8+
public interface LogFormatter {
9+
String format(String message, LogLevel logLevel);
10+
11+
default void setLambdaContext(LambdaContext context) {
12+
};
13+
}

0 commit comments

Comments
 (0)