fix: AWS SdkV2 support with auto-trace

README.md

@@ -94,7 +94,7 @@ Find the latest version here (the format of the version will be n.n.n):
 * Add to your lambda a new layer with the arn from here
 * Add environment variable `JAVA_TOOL_OPTIONS` and set it to `-javaagent:/opt/lumigo-java/lumigo-agent.jar` (This is instead of the flag for more than java11 support)
 * Add the `LUMIGO_TRACER_TOKEN` env var.
-
+* NOTE: because of lambda internal implementation you need to specify the handler with the method name (e.g. `com.example.Handler::handleRequest`) 
 
 ### Configuration
 
@@ -131,7 +131,7 @@ Add the environment variable `JAVA_TOOL_OPTIONS` to your Lambda functions and se
 
 ### Supported Instrumentation Libraries
 
-- Aws SDK V1
+- Aws SDK V1 (Supported only from dependency and not from the auto trace with lambda layer)
 - Aws SDK V2
 - Apache HTTP Client
 - Apache Kafka

agent/src/main/java/io/lumigo/agent/Agent.java

@@ -37,19 +37,19 @@ public static void agentmain(String agentArgs, Instrumentation inst) {
             if ("lib".equalsIgnoreCase(agentArgs)) {
                 urls = getUrls();
             } else {
-                urls =
-                        new URL[] {
-                            new File("/var/task/").toURI().toURL(),
-                            new File(LUMIGO_JAVA_TRACER_PATH).toURI().toURL()
-                        };
+                List<URL> jars = new LinkedList<>();
+                jars.add(new File("/var/task/").toURI().toURL());
+                if (isAutoTrace()) {
+                    jars.add(new File(LUMIGO_JAVA_TRACER_PATH).toURI().toURL());
+                    installTracerJar(inst);
+                }
+                urls = jars.toArray(new URL[jars.size()]);
             }
-            installTracerJar(inst);
-            URLClassLoader newClassLoader = new URLClassLoader(urls, null);
-            Thread.currentThread().setContextClassLoader(newClassLoader);
-            final Class<?> loader =
-                    newClassLoader.loadClass("io.lumigo.core.instrumentation.agent.Loader");
-            final Method instrument = loader.getMethod("instrument", Instrumentation.class);
-            instrument.invoke(null, inst);
+            URLClassLoader agentClassLoader = new AgentClassLoader(urls, ClassLoader.getSystemClassLoader());
+            final Class<?> instrumentationLoader =
+                    agentClassLoader.loadClass("io.lumigo.core.instrumentation.agent.Loader");
+            final Method instrument = instrumentationLoader.getMethod("instrument", Instrumentation.class, ClassLoader.class);
+            instrument.invoke(null, inst, agentClassLoader);
         } catch (Exception e) {
             e.printStackTrace();
         }
@@ -89,4 +89,9 @@ public static boolean isKillSwitchOn() {
         String value = System.getenv("LUMIGO_SWITCH_OFF");
         return "true".equalsIgnoreCase(value);
     }
+
+    public static boolean isAutoTrace() {
+        String value = System.getenv("JAVA_TOOL_OPTIONS");
+        return !value.contains("allowAttachSelf=true");
+    }
 }

agent/src/main/java/io/lumigo/agent/AgentClassLoader.java

@@ -0,0 +1,72 @@
+package io.lumigo.agent;
+
+import java.net.URL;
+import java.net.URLClassLoader;
+
+/**
+ * This class is refered from :
+ * https://github.com/gaoxingliang/tracing-research/blob/main/bootstrap/src/main/java/com/zoomphant/agent/trace/boost/StandaloneAgentClassloader.java
+ */
+public class AgentClassLoader extends URLClassLoader {
+
+    private final ClassLoader additionalClassloader;
+
+    public AgentClassLoader(URL[] urls, ClassLoader additionalClassloader) {
+        super(urls, ClassLoader.getSystemClassLoader().getParent());
+        this.additionalClassloader = additionalClassloader;
+    }
+
+    @Override
+    protected synchronized Class<?> loadClass(String name, boolean resolve) throws ClassNotFoundException {
+        Class<?> loadedClass = findLoadedClass(name);
+        if (loadedClass != null) {
+            return loadedClass;
+        }
+
+        if (name != null && (name.startsWith("sun.") || name.startsWith("java."))) {
+            return super.loadClass(name, resolve);
+        }
+
+        ClassLoader platformClassLoader = safeGetPlatformClassLoader();
+        if (platformClassLoader != null) {
+            // try load from platform class loader
+            try {
+                loadedClass = platformClassLoader.loadClass(name);
+                return loadedClass;
+            } catch (Exception ignore) {
+            }
+        }
+        if (additionalClassloader != null) {
+            // try load from additional classloader
+            try {
+                loadedClass = additionalClassloader.loadClass(name);
+                return loadedClass;
+            } catch (Exception ignore) {
+            }
+        }
+        try {
+            loadedClass = this.getParent().loadClass(name);
+            return loadedClass;
+        } catch (Exception e) {
+        }
+        try {
+            Class<?> aClass = findClass(name);
+            if (resolve) {
+                resolveClass(aClass);
+            }
+            return aClass;
+        } catch (Exception e) {
+            // ignore
+        }
+        return super.loadClass(name, resolve);
+    }
+
+    private ClassLoader safeGetPlatformClassLoader() {
+        // get platform class loader with reflection
+        try {
+            return (ClassLoader) ClassLoader.class.getDeclaredMethod("getPlatformClassLoader").invoke(null);
+        } catch (Exception e) {
+            return null;
+        }
+    }
+}

pom.xml

@@ -84,7 +84,7 @@
         <dependency>
             <groupId>com.amazonaws</groupId>
             <artifactId>aws-java-sdk-s3</artifactId>
-            <version>1.12.261</version>
+            <version>1.11.505</version>
             <scope>provided</scope>
         </dependency>
         <dependency>
@@ -99,53 +99,47 @@
             <version>1.11.551</version>
             <scope>provided</scope>
         </dependency>
-        <dependency>
-            <groupId>com.squareup.okhttp3</groupId>
-            <artifactId>okhttp</artifactId>
-            <version>3.14.2</version>
-        </dependency>
 
         <!-- AWS sdk V2 dependencies -->
         <dependency>
             <groupId>software.amazon.awssdk</groupId>
-            <artifactId>core</artifactId>
+            <artifactId>lambda</artifactId>
             <version>2.25.45</version>
-            <type>pom</type>
         </dependency>
         <dependency>
             <groupId>software.amazon.awssdk</groupId>
             <artifactId>dynamodb</artifactId>
             <version>2.25.45</version>
-            <scope>provided</scope>
         </dependency>
         <dependency>
             <groupId>software.amazon.awssdk</groupId>
             <artifactId>sqs</artifactId>
             <version>2.25.45</version>
-            <scope>provided</scope>
         </dependency>
         <dependency>
             <groupId>software.amazon.awssdk</groupId>
             <artifactId>sns</artifactId>
             <version>2.25.45</version>
-            <scope>provided</scope>
         </dependency>
         <dependency>
             <groupId>software.amazon.awssdk</groupId>
             <artifactId>kinesis</artifactId>
             <version>2.25.45</version>
-            <scope>provided</scope>
         </dependency>
 
         <!-- Kafka dependencies -->
         <dependency>
             <groupId>org.apache.kafka</groupId>
             <artifactId>kafka-clients</artifactId>
             <version>3.1.0</version>
-            <scope>provided</scope>
         </dependency>
 
         <!-- Tracer dependencies -->
+        <dependency>
+            <groupId>com.squareup.okhttp3</groupId>
+            <artifactId>okhttp</artifactId>
+            <version>3.14.2</version>
+        </dependency>
         <dependency>
             <groupId>com.fasterxml.jackson.core</groupId>
             <artifactId>jackson-core</artifactId>

src/main/java/io/lumigo/core/instrumentation/LumigoInstrumentationApi.java

@@ -8,5 +8,5 @@ public interface LumigoInstrumentationApi {
 
     ElementMatcher<TypeDescription> getTypeMatcher();
 
-    AgentBuilder.Transformer.ForAdvice getTransformer();
+    AgentBuilder.Transformer.ForAdvice getTransformer(ClassLoader classLoader);
 }

src/main/java/io/lumigo/core/instrumentation/agent/Loader.java

@@ -6,7 +6,7 @@
 
 @SuppressWarnings("unused")
 public class Loader {
-    public static void instrument(java.lang.instrument.Instrumentation inst) {
+    public static void instrument(java.lang.instrument.Instrumentation inst, ClassLoader classLoader) {
         Logger.debug("Start Instrumentation");
         ApacheHttpInstrumentation apacheHttpInstrumentation = new ApacheHttpInstrumentation();
         AmazonHttpClientInstrumentation amazonHttpClientInstrumentation =
@@ -24,17 +24,20 @@ public static void instrument(java.lang.instrument.Instrumentation inst) {
                         .disableClassFormatChanges()
                         .with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION)
                         .type(apacheHttpInstrumentation.getTypeMatcher())
-                        .transform(apacheHttpInstrumentation.getTransformer())
+                        .transform(apacheHttpInstrumentation.getTransformer(classLoader))
                         .type(amazonHttpClientInstrumentation.getTypeMatcher())
-                        .transform(amazonHttpClientInstrumentation.getTransformer())
+                        .transform(amazonHttpClientInstrumentation.getTransformer(classLoader))
                         .type(amazonHttpClientV2Instrumentation.getTypeMatcher())
-                        .transform(amazonHttpClientV2Instrumentation.getTransformer())
+                        .transform(amazonHttpClientV2Instrumentation.getTransformer(classLoader))
                         .type(apacheKafkaInstrumentation.getTypeMatcher())
-                        .transform(apacheKafkaInstrumentation.getTransformer())
+                        .transform(apacheKafkaInstrumentation.getTransformer(classLoader))
                         .type(apacheKafkaConsumerInstrumentation.getTypeMatcher())
-                        .transform(apacheKafkaConsumerInstrumentation.getTransformer())
+                        .transform(apacheKafkaConsumerInstrumentation.getTransformer(classLoader))
                         .type(awsLambdaRequestHandlerInstrumentation.getTypeMatcher())
-                        .transform(awsLambdaRequestHandlerInstrumentation.getTransformer());
+                        .transform(awsLambdaRequestHandlerInstrumentation.getTransformer(classLoader))
+                        .with(AgentBuilder.Listener.StreamWriting.toSystemError())
+                        .with(AgentBuilder.InstallationListener.StreamWriting.toSystemError())
+                ;
 
         builder.installOn(inst);
         Logger.debug("Finish Instrumentation");

src/main/java/io/lumigo/core/instrumentation/impl/AmazonHttpClientInstrumentation.java

@@ -23,10 +23,13 @@ public ElementMatcher<TypeDescription> getTypeMatcher() {
     }
 
     @Override
-    public AgentBuilder.Transformer.ForAdvice getTransformer() {
+    public AgentBuilder.Transformer.ForAdvice getTransformer(ClassLoader classLoader) {
         return new AgentBuilder.Transformer.ForAdvice()
-                .include(Loader.class.getClassLoader())
-                .advice(isMethod().and(named("execute")), AmazonHttpClientAdvice.class.getName());
+                .include(classLoader)
+                .advice(
+                        isMethod().and(named("execute")),
+                        AmazonHttpClientInstrumentation.class.getName()
+                                + "$AmazonHttpClientAdvice");
     }
 
     public static class AmazonHttpClientAdvice {

src/main/java/io/lumigo/core/instrumentation/impl/AmazonHttpClientV2Instrumentation.java

@@ -12,10 +12,10 @@
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 import org.pmw.tinylog.Logger;
+import software.amazon.awssdk.core.client.builder.SdkDefaultClientBuilder;
 import software.amazon.awssdk.core.interceptor.Context;
 import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
 import software.amazon.awssdk.core.interceptor.ExecutionInterceptor;
-import software.amazon.awssdk.http.SdkHttpRequest;
 
 public class AmazonHttpClientV2Instrumentation implements LumigoInstrumentationApi {
 
@@ -28,12 +28,13 @@ public ElementMatcher<TypeDescription> getTypeMatcher() {
     }
 
     @Override
-    public AgentBuilder.Transformer.ForAdvice getTransformer() {
+    public AgentBuilder.Transformer.ForAdvice getTransformer(ClassLoader classLoader) {
         return new AgentBuilder.Transformer.ForAdvice()
-                .include(Loader.class.getClassLoader())
+                .include(classLoader)
                 .advice(
                         isMethod().and(named("resolveExecutionInterceptors")),
-                        AmazonHttpClientV2Advice.class.getName());
+                        AmazonHttpClientV2Instrumentation.class.getName()
+                                + "$AmazonHttpClientV2Advice");
     }
 
     @SuppressWarnings("unused")
@@ -61,20 +62,11 @@ public static class TracingExecutionInterceptor implements ExecutionInterceptor
             public void beforeExecution(
                     final Context.BeforeExecution context,
                     final ExecutionAttributes executionAttributes) {
-                startTimeMap.put(context.request().hashCode(), System.currentTimeMillis());
-            }
-
-            @Override
-            public SdkHttpRequest modifyHttpRequest(
-                    Context.ModifyHttpRequest context, ExecutionAttributes executionAttributes) {
                 try {
-                    SdkHttpRequest.Builder requestBuilder = context.httpRequest().toBuilder();
-                    requestBuilder.appendHeader("X-Amzn-Trace-Id", spansContainer.getPatchedRoot());
-                    return requestBuilder.build();
+                    startTimeMap.put(context.request().hashCode(), System.currentTimeMillis());
                 } catch (Throwable e) {
-                    Logger.debug("Unable to inject trace header", e);
+                    Logger.error(e, "Failed save trace context");
                 }
-                return context.httpRequest();
             }
 
             @Override

src/main/java/io/lumigo/core/instrumentation/impl/ApacheHttpInstrumentation.java

@@ -23,9 +23,9 @@ public ElementMatcher<TypeDescription> getTypeMatcher() {
     }
 
     @Override
-    public AgentBuilder.Transformer.ForAdvice getTransformer() {
+    public AgentBuilder.Transformer.ForAdvice getTransformer(ClassLoader classLoader) {
         return new AgentBuilder.Transformer.ForAdvice()
-                .include(Loader.class.getClassLoader())
+                .include(classLoader)
                 .advice(
                         isMethod()
                                 .and(named("execute"))
@@ -36,7 +36,7 @@ public AgentBuilder.Transformer.ForAdvice getTransformer() {
                                                                 0,
                                                                 named(
                                                                         "org.apache.http.client.methods.HttpUriRequest")))),
-                        ApacheHttpAdvice.class.getName());
+                        ApacheHttpInstrumentation.class.getName() + "$ApacheHttpAdvice");
     }
 
     public static class ApacheHttpAdvice {

src/main/java/io/lumigo/core/instrumentation/impl/ApacheKafkaConsumerInstrumentation.java

@@ -22,9 +22,9 @@ public ElementMatcher<TypeDescription> getTypeMatcher() {
     }
 
     @Override
-    public AgentBuilder.Transformer.ForAdvice getTransformer() {
+    public AgentBuilder.Transformer.ForAdvice getTransformer(ClassLoader classLoader) {
         return new AgentBuilder.Transformer.ForAdvice()
-                .include(Loader.class.getClassLoader())
+                .include(classLoader)
                 .advice(
                         isMethod()
                                 .and(isPublic())
@@ -34,9 +34,11 @@ public AgentBuilder.Transformer.ForAdvice getTransformer() {
                                         returns(
                                                 named(
                                                         "org.apache.kafka.clients.consumer.ConsumerRecords"))),
-                        ApacheKafkaConsumerAdvice.class.getName());
+                        ApacheKafkaConsumerInstrumentation.class.getName()
+                                + "$ApacheKafkaConsumerAdvice");
     }
 
+    @SuppressWarnings("unused")
     public static class ApacheKafkaConsumerAdvice {
         public static final SpansContainer spansContainer = SpansContainer.getInstance();
         public static final LRUCache<String, Long> startTimeMap = new LRUCache<>(1000);

src/main/java/io/lumigo/core/instrumentation/impl/ApacheKafkaProducerInstrumentation.java

@@ -27,9 +27,9 @@ public ElementMatcher<TypeDescription> getTypeMatcher() {
     }
 
     @Override
-    public AgentBuilder.Transformer.ForAdvice getTransformer() {
+    public AgentBuilder.Transformer.ForAdvice getTransformer(ClassLoader classLoader) {
         return new AgentBuilder.Transformer.ForAdvice()
-                .include(Loader.class.getClassLoader())
+                .include(classLoader)
                 .advice(
                         isMethod()
                                 .and(isPublic())
@@ -44,7 +44,8 @@ public AgentBuilder.Transformer.ForAdvice getTransformer() {
                                                                 1,
                                                                 named(
                                                                         "org.apache.kafka.clients.producer.Callback")))),
-                        ApacheKafkaProducerAdvice.class.getName());
+                        ApacheKafkaProducerInstrumentation.class.getName()
+                                + "$ApacheKafkaProducerAdvice");
     }
 
     @SuppressWarnings("unused")