diff --git a/pom.xml b/pom.xml
index 827c52cf0..3f20e8839 100644
--- a/pom.xml
+++ b/pom.xml
@@ -78,7 +78,13 @@
log4j-api
${log4j.version}
+
+ org.aspectj
+ aspectjrt
+ ${aspectj.version}
+
+
org.junit.jupiter
junit-jupiter-api
@@ -91,173 +97,117 @@
5.6.2
test
+
+ org.mockito
+ mockito-core
+ 3.3.3
+ test
+
org.aspectj
- aspectjrt
+ aspectjweaver
${aspectj.version}
+ test
+
+
+ org.assertj
+ assertj-core
+ 3.9.1
+ test
+
+
+
+ org.apache.maven.plugins
+ maven-compiler-plugin
+ 3.8.1
+
+ ${maven.compiler.source}
+ ${maven.compiler.target}
+ false
+
+
+
+
+ com.nickwongdev
+ aspectj-maven-plugin
+ 1.12.1
+
+ ${maven.compiler.source}
+ ${maven.compiler.target}
+ ${maven.compiler.target}
+ ignore
+ ${project.build.sourceEncoding}
+
+
+
+ process-sources
+
+ compile
+ test-compile
+
+
+
+
+
+ org.aspectj
+ aspectjtools
+ ${aspectj.version}
+
+
+
+
+ org.apache.maven.plugins
+ maven-surefire-plugin
+ 2.22.2
+
+
+ org.jacoco
+ jacoco-maven-plugin
+ 0.8.5
+
+
+
+ prepare-agent
+
+
+
+
+ report
+ test
+
+ report
+
+
+
+
+
+ org.apache.maven.plugins
+ maven-javadoc-plugin
+ 2.9.1
+
+ -Xdoclint:none
+ false
+
+
+
+ attach-javadocs
+
+ jar
+
+
+
+
+
+
+
-
- dev
-
- true
-
-
-
-
- org.apache.maven.plugins
- maven-compiler-plugin
- 3.8.1
-
- ${maven.compiler.source}
- ${maven.compiler.target}
- false
-
-
-
-
- com.nickwongdev
- aspectj-maven-plugin
- 1.12.1
-
- ${maven.compiler.source}
- ${maven.compiler.target}
- ${maven.compiler.target}
- ignore
- ${project.build.sourceEncoding}
-
-
-
- process-sources
-
- compile
- test-compile
-
-
-
-
-
- org.aspectj
- aspectjtools
- ${aspectj.version}
-
-
-
-
- org.apache.maven.plugins
- maven-surefire-plugin
- 2.22.2
-
-
- org.jacoco
- jacoco-maven-plugin
- 0.8.2
-
-
-
- prepare-agent
-
-
-
-
- report
- test
-
- report
-
-
-
-
-
- org.apache.maven.plugins
- maven-javadoc-plugin
- 2.9.1
-
- -Xdoclint:none
- false
-
-
-
- attach-javadocs
-
- jar
-
-
-
-
-
-
-
release
-
- org.apache.maven.plugins
- maven-compiler-plugin
- 3.8.1
-
- ${maven.compiler.source}
- ${maven.compiler.target}
- false
-
-
-
-
- com.nickwongdev
- aspectj-maven-plugin
- 1.12.1
-
- ${maven.compiler.source}
- ${maven.compiler.target}
- ${maven.compiler.target}
- ignore
- ${project.build.sourceEncoding}
-
-
-
- process-sources
-
- compile
- test-compile
-
-
-
-
-
- org.aspectj
- aspectjtools
- ${aspectj.version}
-
-
-
-
- org.apache.maven.plugins
- maven-surefire-plugin
- 2.22.1
-
-
- org.jacoco
- jacoco-maven-plugin
- 0.8.2
-
-
-
- prepare-agent
-
-
-
-
- report
- test
-
- report
-
-
-
-
org.apache.maven.plugins
maven-source-plugin
@@ -271,23 +221,6 @@
-
- org.apache.maven.plugins
- maven-javadoc-plugin
- 2.9.1
-
- -Xdoclint:none
- false
-
-
-
- attach-javadocs
-
- jar
-
-
-
-
org.apache.maven.plugins
maven-gpg-plugin
@@ -318,5 +251,4 @@
-
\ No newline at end of file
diff --git a/src/main/java/software/aws/lambda/logging/DefaultLambdaFields.java b/src/main/java/software/aws/lambda/logging/DefaultLambdaFields.java
index 1d5fc62ef..8fddc3b75 100644
--- a/src/main/java/software/aws/lambda/logging/DefaultLambdaFields.java
+++ b/src/main/java/software/aws/lambda/logging/DefaultLambdaFields.java
@@ -17,6 +17,10 @@ enum DefaultLambdaFields {
this.name = name;
}
+ public String getName() {
+ return name;
+ }
+
static Map values(Context context) {
Map hashMap = new HashMap<>();
diff --git a/src/main/java/software/aws/lambda/logging/LambdaAspect.java b/src/main/java/software/aws/lambda/logging/LambdaAspect.java
index d6f24aef7..ea13082a2 100644
--- a/src/main/java/software/aws/lambda/logging/LambdaAspect.java
+++ b/src/main/java/software/aws/lambda/logging/LambdaAspect.java
@@ -1,15 +1,22 @@
package software.aws.lambda.logging;
+import java.util.Optional;
+
import com.amazonaws.services.lambda.runtime.Context;
+import com.amazonaws.services.lambda.runtime.RequestHandler;
+import com.amazonaws.services.lambda.runtime.RequestStreamHandler;
import org.apache.logging.log4j.ThreadContext;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
+import static java.util.Optional.empty;
+import static java.util.Optional.of;
+
@Aspect
public final class LambdaAspect {
- private static Boolean IS_COLD_START = null;
+ static Boolean IS_COLD_START = null;
@Pointcut("@annotation(powerToolsLogging)")
public void callAt(PowerToolsLogging powerToolsLogging) {
@@ -19,17 +26,32 @@ public void callAt(PowerToolsLogging powerToolsLogging) {
public Object around(ProceedingJoinPoint pjp,
PowerToolsLogging powerToolsLogging) throws Throwable {
- // TODO JoinPoint that annotation is used on entry method of lambda or do we want it to work anywhere
- if(powerToolsLogging.injectContextInfo()) {
- if(pjp.getArgs().length == 2 && pjp.getArgs()[1] instanceof Context) {
- ThreadContext.putAll(DefaultLambdaFields.values((Context) pjp.getArgs()[1]));
- }
- }
-
- ThreadContext.put("coldStart", null == IS_COLD_START? "true" : "false");
+ extractContext(pjp)
+ .ifPresent(context -> {
+ ThreadContext.putAll(DefaultLambdaFields.values(context));
+ ThreadContext.put("coldStart", null == IS_COLD_START ? "true" : "false");
+ });
IS_COLD_START = false;
+
return pjp.proceed();
}
+
+ private Optional extractContext(ProceedingJoinPoint pjp) {
+
+ if ("handleRequest".equals(pjp.getSignature().getName())) {
+ if (RequestHandler.class.isAssignableFrom(pjp.getSignature().getDeclaringType())
+ && pjp.getArgs().length == 2 && pjp.getArgs()[1] instanceof Context) {
+ return of((Context) pjp.getArgs()[1]);
+ }
+
+ if (RequestStreamHandler.class.isAssignableFrom(pjp.getSignature().getDeclaringType())
+ && pjp.getArgs().length == 3 && pjp.getArgs()[2] instanceof Context) {
+ return of((Context) pjp.getArgs()[2]);
+ }
+ }
+
+ return empty();
+ }
}
diff --git a/src/main/java/software/aws/lambda/logging/PowerToolsLogging.java b/src/main/java/software/aws/lambda/logging/PowerToolsLogging.java
index aa09d6220..aba262569 100644
--- a/src/main/java/software/aws/lambda/logging/PowerToolsLogging.java
+++ b/src/main/java/software/aws/lambda/logging/PowerToolsLogging.java
@@ -8,6 +8,4 @@
@Retention(RetentionPolicy.RUNTIME)
@Target(ElementType.METHOD)
public @interface PowerToolsLogging {
-
- boolean injectContextInfo() default true;
}
diff --git a/src/test/java/software/aws/lambda/logging/LambdaAspectTest.java b/src/test/java/software/aws/lambda/logging/LambdaAspectTest.java
new file mode 100644
index 000000000..64480bf48
--- /dev/null
+++ b/src/test/java/software/aws/lambda/logging/LambdaAspectTest.java
@@ -0,0 +1,114 @@
+package software.aws.lambda.logging;
+
+import java.io.IOException;
+
+import com.amazonaws.services.lambda.runtime.Context;
+import com.amazonaws.services.lambda.runtime.RequestHandler;
+import com.amazonaws.services.lambda.runtime.RequestStreamHandler;
+import org.apache.logging.log4j.ThreadContext;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.mockito.Mock;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.when;
+import static org.mockito.MockitoAnnotations.initMocks;
+
+class LambdaAspectTest {
+
+ private RequestStreamHandler requestStreamHandler;
+ private RequestHandler