Skip to content

Commit

Permalink
[alc] ContextLogger changes in RIC (#436)
Browse files Browse the repository at this point in the history
* [alc] ContextLogger changes in RIC

---------

Co-authored-by: Daniel Torok <torokd@amazon.com>
  • Loading branch information
dtorok and dtorok committed Aug 29, 2023
1 parent e87c680 commit 2525fc8
Show file tree
Hide file tree
Showing 22 changed files with 532 additions and 36 deletions.
2 changes: 1 addition & 1 deletion aws-lambda-java-runtime-interface-client/pom.xml
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,7 @@
<dependency>
<groupId>com.amazonaws</groupId>
<artifactId>aws-lambda-java-core</artifactId>
<version>1.2.2</version>
<version>1.2.3</version>
</dependency>
<dependency>
<groupId>com.amazonaws</groupId>
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,8 @@
import com.amazonaws.services.lambda.runtime.api.client.runtimeapi.LambdaRuntimeClient;
import com.amazonaws.services.lambda.runtime.api.client.util.LambdaOutputStream;
import com.amazonaws.services.lambda.runtime.api.client.util.UnsafeUtil;
import com.amazonaws.services.lambda.runtime.logging.LogFormat;
import com.amazonaws.services.lambda.runtime.logging.LogLevel;
import com.amazonaws.services.lambda.runtime.serialization.PojoSerializer;
import com.amazonaws.services.lambda.runtime.serialization.factories.GsonFactory;
import com.amazonaws.services.lambda.runtime.serialization.factories.JacksonFactory;
Expand Down Expand Up @@ -187,7 +189,12 @@ public static void main(String[] args) {

private static void startRuntime(String handler) {
try (LogSink logSink = createLogSink()) {
startRuntime(handler, new LambdaContextLogger(logSink));
LambdaLogger logger = new LambdaContextLogger(
logSink,
LogLevel.fromString(LambdaEnvironment.LAMBDA_LOG_LEVEL),
LogFormat.fromString(LambdaEnvironment.LAMBDA_LOG_FORMAT)
);
startRuntime(handler, logger);
} catch (Throwable t) {
throw new Error(t);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4,13 +4,15 @@

import com.amazonaws.services.lambda.runtime.ClientContext;
import com.amazonaws.services.lambda.runtime.Context;
import com.amazonaws.services.lambda.runtime.LambdaLogger;
import com.amazonaws.services.lambda.runtime.LambdaRuntimeInternal;
import com.amazonaws.services.lambda.runtime.RequestHandler;
import com.amazonaws.services.lambda.runtime.RequestStreamHandler;
import com.amazonaws.services.lambda.runtime.api.client.LambdaRequestHandler.UserFaultHandler;
import com.amazonaws.services.lambda.runtime.api.client.api.LambdaClientContext;
import com.amazonaws.services.lambda.runtime.api.client.api.LambdaCognitoIdentity;
import com.amazonaws.services.lambda.runtime.api.client.api.LambdaContext;
import com.amazonaws.services.lambda.runtime.api.client.logging.LambdaContextLogger;
import com.amazonaws.services.lambda.runtime.api.client.runtimeapi.InvocationRequest;
import com.amazonaws.services.lambda.runtime.api.client.util.UnsafeUtil;
import com.amazonaws.services.lambda.runtime.serialization.PojoSerializer;
Expand Down Expand Up @@ -844,6 +846,22 @@ private void safeAddRequestIdToLog4j(String log4jContextClassName,
}
}

/**
* Passes the LambdaContext to the logger so that the JSON formatter can include the requestId.
*
* We do casting here because both the LambdaRuntime and the LambdaLogger is in the core package,
* and the setLambdaContext(context) is a method we don't want to publish for customers. That method is
* only implemented on the internal LambdaContextLogger, so we check and cast to be able to call it.
* @param context the LambdaContext
*/
private void safeAddContextToLambdaLogger(LambdaContext context) {
LambdaLogger logger = com.amazonaws.services.lambda.runtime.LambdaRuntime.getLogger();
if (logger instanceof LambdaContextLogger) {
LambdaContextLogger contextLogger = (LambdaContextLogger) logger;
contextLogger.setLambdaContext(context);
}
}

public ByteArrayOutputStream call(InvocationRequest request) throws Error, Exception {
output.reset();

Expand Down Expand Up @@ -871,6 +889,8 @@ public ByteArrayOutputStream call(InvocationRequest request) throws Error, Excep
clientContext
);

safeAddContextToLambdaLogger(context);

if (LambdaRuntimeInternal.getUseLog4jAppender()) {
safeAddRequestIdToLog4j("org.apache.log4j.MDC", request, Object.class);
safeAddRequestIdToLog4j("org.apache.logging.log4j.ThreadContext", request, String.class);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,8 @@ public class LambdaEnvironment {
public static final int MEMORY_LIMIT = parseInt(ENV_READER.getEnvOrDefault(AWS_LAMBDA_FUNCTION_MEMORY_SIZE, "128"));
public static final String LOG_GROUP_NAME = ENV_READER.getEnv(AWS_LAMBDA_LOG_GROUP_NAME);
public static final String LOG_STREAM_NAME = ENV_READER.getEnv(AWS_LAMBDA_LOG_STREAM_NAME);
public static final String LAMBDA_LOG_LEVEL = ENV_READER.getEnvOrDefault(AWS_LAMBDA_LOG_LEVEL, "UNDEFINED");
public static final String LAMBDA_LOG_FORMAT = ENV_READER.getEnvOrDefault(AWS_LAMBDA_LOG_FORMAT, "TEXT");
public static final String FUNCTION_NAME = ENV_READER.getEnv(AWS_LAMBDA_FUNCTION_NAME);
public static final String FUNCTION_VERSION = ENV_READER.getEnv(AWS_LAMBDA_FUNCTION_VERSION);
}
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,16 @@ public interface ReservedRuntimeEnvironmentVariables {
*/
String AWS_LAMBDA_LOG_STREAM_NAME = "AWS_LAMBDA_LOG_STREAM_NAME";

/**
* The logging level set for the function.
*/
String AWS_LAMBDA_LOG_LEVEL = "AWS_LAMBDA_LOG_LEVEL";

/**
* The logging format set for the function.
*/
String AWS_LAMBDA_LOG_FORMAT = "AWS_LAMBDA_LOG_FORMAT";

/**
* Access key id obtained from the function's execution role.
*/
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,67 @@
/* Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved. */

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

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

import com.amazonaws.services.lambda.runtime.api.client.api.LambdaContext;
import com.amazonaws.services.lambda.runtime.LambdaLogger;
import com.amazonaws.services.lambda.runtime.logging.LogFormat;
import com.amazonaws.services.lambda.runtime.logging.LogLevel;

/**
* Provides default implementation of the convenience logger functions.
* When extending AbstractLambdaLogger, only one function has to be overridden:
* void logMessage(byte[] message, LogLevel logLevel);
*/
public abstract class AbstractLambdaLogger implements LambdaLogger {
private final LogFiltering logFiltering;
private final LogFormatter logFormatter;
protected final LogFormat logFormat;

public AbstractLambdaLogger(LogLevel logLevel, LogFormat logFormat) {
this.logFiltering = new LogFiltering(logLevel);

this.logFormat = logFormat;
if (logFormat == LogFormat.JSON) {
logFormatter = new JsonLogFormatter();
} else {
logFormatter = new TextLogFormatter();
}
}

protected abstract void logMessage(byte[] message, LogLevel logLevel);

protected void logMessage(String message, LogLevel logLevel) {
logMessage(message.getBytes(UTF_8), logLevel);
}

@Override
public void log(String message, LogLevel logLevel) {
if (logFiltering.isEnabled(logLevel)) {
this.logMessage(logFormatter.format(message, logLevel), logLevel);
}
}

@Override
public void log(byte[] message, LogLevel logLevel) {
if (logFiltering.isEnabled(logLevel)) {
// there is no formatting for byte[] messages
this.logMessage(message, logLevel);
}
}

@Override
public void log(String message) {
this.log(message, LogLevel.UNDEFINED);
}

@Override
public void log(byte[] message) {
this.log(message, LogLevel.UNDEFINED);
}

public void setLambdaContext(LambdaContext lambdaContext) {
this.logFormatter.setLambdaContext(lambdaContext);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -2,12 +2,40 @@

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

public enum FrameType {
import com.amazonaws.services.lambda.runtime.logging.LogLevel;
import com.amazonaws.services.lambda.runtime.logging.LogFormat;

LOG(0xa55a0003);
/**
* 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.
* +-----------------------+
* | Frame Type - 4 bytes |
* +-----------------------+
* | a5 5a 00 | flgs |
* + - - - - - + - - - - - +
* \ bit |
* | view|
* +---------+ +
* | |
* v byte 3 v F - free
* +-+-+-+-+-+-+-+-+ J - { JsonLog = 0, PlainTextLog = 1 }
* |F|F|F|L|l|l|T|J| T - { NoTimeStamp = 0, TimeStampPresent = 1 }
* +-+-+-+-+-+-+-+-+ Lll -> Log Level in 3-bit binary (L-> most significant bit)
*/
public class FrameType {
private static final int LOG_MAGIC = 0xa55a0000;
private static final int OFFSET_LOG_FORMAT = 0;
private static final int OFFSET_TIMESTAMP_PRESENT = 1;
private static final int OFFSET_LOG_LEVEL = 2;

private final int val;

public static int getValue(LogLevel logLevel, LogFormat logFormat) {
return LOG_MAGIC |
(logLevel.ordinal() << OFFSET_LOG_LEVEL) |
(1 << OFFSET_TIMESTAMP_PRESENT) |
(logFormat.ordinal() << OFFSET_LOG_FORMAT);
}

FrameType(int val) {
this.val = val;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,9 @@
import java.nio.ByteOrder;
import java.time.Instant;

import com.amazonaws.services.lambda.runtime.logging.LogLevel;
import com.amazonaws.services.lambda.runtime.logging.LogFormat;

/**
* FramedTelemetryLogSink implements the logging contract between runtimes and the platform. It implements a simple
* framing protocol so message boundaries can be determined. Each frame can be visualized as follows:
Expand Down Expand Up @@ -38,16 +41,21 @@ public FramedTelemetryLogSink(FileDescriptor fd) throws IOException {
}

@Override
public synchronized void log(byte[] message) {
public synchronized void log(LogLevel logLevel, LogFormat logFormat, byte[] message) {
try {
writeFrame(message);
writeFrame(logLevel, logFormat, message);
} catch (IOException e) {
e.printStackTrace();
}
}

private void writeFrame(byte[] message) throws IOException {
updateHeader(message.length);
@Override
public void log(byte[] message) {
log(LogLevel.UNDEFINED, LogFormat.TEXT, message);
}

private void writeFrame(LogLevel logLevel, LogFormat logFormat, byte[] message) throws IOException {
updateHeader(logLevel, logFormat, message.length);
this.logOutputStream.write(this.headerBuf.array());
this.logOutputStream.write(message);
}
Expand All @@ -60,9 +68,9 @@ private long timestamp() {
/**
* Updates the header ByteBuffer with the provided length. The header comprises the frame type and message length.
*/
private void updateHeader(int length) {
private void updateHeader(LogLevel logLevel, LogFormat logFormat, int length) {
this.headerBuf.clear();
this.headerBuf.putInt(FrameType.LOG.getValue());
this.headerBuf.putInt(FrameType.getValue(logLevel, logFormat));
this.headerBuf.putInt(length);
this.headerBuf.putLong(timestamp());
this.headerBuf.flip();
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,55 @@
/* Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved. */

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

import com.amazonaws.services.lambda.runtime.api.client.api.LambdaContext;
import com.amazonaws.services.lambda.runtime.logging.LogLevel;
import com.amazonaws.services.lambda.runtime.serialization.PojoSerializer;
import com.amazonaws.services.lambda.runtime.serialization.factories.GsonFactory;

import java.io.ByteArrayOutputStream;
import java.nio.charset.StandardCharsets;
import java.time.LocalDateTime;
import java.time.ZoneId;
import java.time.format.DateTimeFormatter;

public class JsonLogFormatter implements LogFormatter {
private final PojoSerializer<StructuredLogMessage> serializer = GsonFactory.getInstance().getSerializer(StructuredLogMessage.class);
private LambdaContext lambdaContext;

private static final DateTimeFormatter dateFormatter =
DateTimeFormatter
.ofPattern("yyyy-MM-dd'T'HH:mm:ss.SSS'Z'")
.withZone(ZoneId.of("UTC"));

@Override
public String format(String message, LogLevel logLevel) {
ByteArrayOutputStream stream = new ByteArrayOutputStream();
StructuredLogMessage msg = createLogMessage(message, logLevel);
serializer.toJson(msg, stream);
stream.write('\n');
return new String(stream.toByteArray(), StandardCharsets.UTF_8);
}

private StructuredLogMessage createLogMessage(String message, LogLevel logLevel) {
StructuredLogMessage msg = new StructuredLogMessage();
msg.timestamp = dateFormatter.format(LocalDateTime.now());
msg.message = message;
msg.level = logLevel;

if (lambdaContext != null) {
msg.AWSRequestId = lambdaContext.getAwsRequestId();
}
return msg;
}


/**
* Function to set the context for every invocation.
* This way the logger will be able to attach additional information to the log packet.
*/
@Override
public void setLambdaContext(LambdaContext context) {
this.lambdaContext = context;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -2,33 +2,29 @@

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

import com.amazonaws.services.lambda.runtime.LambdaLogger;
import com.amazonaws.services.lambda.runtime.api.client.api.LambdaContext;
import com.amazonaws.services.lambda.runtime.logging.LogFormat;
import com.amazonaws.services.lambda.runtime.logging.LogLevel;

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

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

private final transient LogSink sink;

public LambdaContextLogger(LogSink sink) {
public LambdaContextLogger(LogSink sink, LogLevel logLevel, LogFormat logFormat) {
super(logLevel, logFormat);
this.sink = sink;
}

public void log(byte[] message) {
@Override
protected void logMessage(byte[] message, LogLevel logLevel) {
if (message == null) {
message = NULL_BYTES_VALUE;
}
sink.log(message);
}

public void log(String message) {
if (message == null) {
this.log(NULL_BYTES_VALUE);
} else {
this.log(message.getBytes(UTF_8));
}
sink.log(logLevel, this.logFormat, message);
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
/* Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved. */

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

import com.amazonaws.services.lambda.runtime.logging.LogLevel;

public class LogFiltering {
private final LogLevel minimumLogLevel;

public LogFiltering(LogLevel minimumLogLevel) {
this.minimumLogLevel = minimumLogLevel;
}

boolean isEnabled(LogLevel logLevel) {
return (logLevel == LogLevel.UNDEFINED || logLevel.ordinal() >= minimumLogLevel.ordinal());
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
/* Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved. */

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

import com.amazonaws.services.lambda.runtime.api.client.api.LambdaContext;
import com.amazonaws.services.lambda.runtime.logging.LogLevel;

public interface LogFormatter {
String format(String message, LogLevel logLevel);

default void setLambdaContext(LambdaContext context) {
};
}

0 comments on commit 2525fc8

Please sign in to comment.