feat: support auto-trace

agent/pom.xml

@@ -32,12 +32,6 @@
             <organization>Lumigo</organization>
             <organizationUrl>https://lumigo.io/</organizationUrl>
         </developer>
-        <developer>
-            <name>Uri Parush</name>
-            <email>[email protected]</email>
-            <organization>Lumigo</organization>
-            <organizationUrl>https://lumigo.io</organizationUrl>
-        </developer>
     </developers>
 
     <organization>
@@ -104,9 +98,12 @@
                     <archive>
                         <manifestEntries>
                             <Agent-Class>io.lumigo.agent.Agent</Agent-Class>
-                            <Can-Redefine-Classes>false</Can-Redefine-Classes>
+                            <Main-Class>io.lumigo.agent.Agent</Main-Class>
+                            <Premain-Class>io.lumigo.agent.Agent</Premain-Class>
+                            <Can-Redefine-Classes>true</Can-Redefine-Classes>
                             <Can-Retransform-Classes>true</Can-Retransform-Classes>
                             <Can-Set-Native-Method-Prefix>false</Can-Set-Native-Method-Prefix>
+                            <Implementation-Vendor>Lumigo</Implementation-Vendor>
                         </manifestEntries>
                     </archive>
                 </configuration>

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

@@ -4,26 +4,45 @@
 import java.lang.instrument.Instrumentation;
 import java.lang.reflect.Method;
 import java.net.MalformedURLException;
+import java.net.URI;
 import java.net.URL;
 import java.net.URLClassLoader;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.util.LinkedList;
 import java.util.List;
+import java.util.jar.JarFile;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
 public class Agent {
 
+    private static boolean initialized = false;
+
+    public static String LUMIGO_JAVA_TRACER_PATH = "/opt/lumigo/lumigo-tracer.jar";
+
+    public static void premain(String agentArgs, Instrumentation inst) {
+        agentmain(agentArgs, inst);
+        // TODO handle switch off
+    }
+
     public static void agentmain(String agentArgs, Instrumentation inst) {
+        if (initialized) {
+            return;
+        }
+        initialized = true;
         try {
             URL[] urls;
             if ("lib".equalsIgnoreCase(agentArgs)) {
                 urls = getUrls();
             } else {
-                urls = new URL[] {new File("/var/task/").toURI().toURL()};
+                List<URL> jars = new LinkedList<>();
+                jars.add(new File("/var/task/").toURI().toURL());
+                jars.add(new File(LUMIGO_JAVA_TRACER_PATH).toURI().toURL());
+                urls = jars.toArray(new URL[jars.size()]);
             }
+            installTracerJar(inst);
             URLClassLoader newClassLoader = new URLClassLoader(urls, null);
             Thread.currentThread().setContextClassLoader(newClassLoader);
             final Class<?> loader =
@@ -35,6 +54,15 @@ public static void agentmain(String agentArgs, Instrumentation inst) {
         }
     }
 
+    private static void installTracerJar(Instrumentation inst) {
+        try (JarFile jar = new JarFile(new File(new File(LUMIGO_JAVA_TRACER_PATH).toURI()))){
+            inst.appendToSystemClassLoaderSearch(jar);
+        } catch (Exception e) {
+            System.out.println("Failed to append java tracer jar to system class loader: " + LUMIGO_JAVA_TRACER_PATH);
+            e.printStackTrace();
+        }
+    }
+
     public static URL[] getUrls() {
         List<URL> jars = new LinkedList<>();
         try (Stream<Path> paths = Files.walk(Paths.get("/var/task/lib"))) {

layer/generate.sh

@@ -0,0 +1,8 @@
+mkdir -p lumigo
+
+cp ../target/java-tracer-1.0.47.jar lumigo/lumigo-tracer.jar
+cp ../agent/target/lumigo-agent-1.0.47.jar lumigo/lumigo-agent.jar
+
+zip -r lumigo-java-layer.zip lumigo
+
+rm -rf lumigo

pom.xml

@@ -182,6 +182,7 @@
             <artifactId>json</artifactId>
             <version>20210307</version>
         </dependency>
+
         <!-- Testing dependencies -->
         <dependency>
             <groupId>org.junit.jupiter</groupId>
@@ -235,6 +236,19 @@
             </resource>
         </resources>
         <plugins>
+            <plugin>
+                <groupId>org.apache.maven.plugins</groupId>
+                <artifactId>maven-shade-plugin</artifactId>
+                <version>3.2.4</version>
+                <executions>
+                    <execution>
+                        <phase>package</phase>
+                        <goals>
+                            <goal>shade</goal>
+                        </goals>
+                    </execution>
+                </executions>
+            </plugin>
             <!-- Used to make it so the JavaDoc and sources are included. -->
             <plugin>
                 <groupId>org.apache.maven.plugins</groupId>

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

@@ -7,6 +7,9 @@
 import net.bytebuddy.agent.builder.AgentBuilder;
 import org.pmw.tinylog.Logger;
 
+import java.util.Arrays;
+
+@SuppressWarnings("unused")
 public class Loader {
     public static void instrument(java.lang.instrument.Instrumentation inst) {
         Logger.debug("Start Instrumentation");
@@ -19,28 +22,12 @@ public static void instrument(java.lang.instrument.Instrumentation inst) {
                 new ApacheKafkaProducerInstrumentation();
         ApacheKafkaConsumerInstrumentation apacheKafkaConsumerInstrumentation =
                 new ApacheKafkaConsumerInstrumentation();
+        AwsLambdaRequestHandlerInstrumentation awsLambdaRequestHandlerInstrumentation =
+                new AwsLambdaRequestHandlerInstrumentation();
         AgentBuilder builder =
                 new AgentBuilder.Default()
                         .disableClassFormatChanges()
                         .with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION)
-                        .ignore(
-                                not(nameStartsWith("com.amazonaws.http.AmazonHttpClient"))
-                                        .and(not(nameStartsWith("org.apache.http.impl.client")))
-                                        .and(
-                                                not(
-                                                        nameStartsWith(
-                                                                AmazonHttpClientV2Instrumentation
-                                                                        .INSTRUMENTATION_PACKAGE_PREFIX)))
-                                        .and(
-                                                not(
-                                                        nameStartsWith(
-                                                                ApacheKafkaProducerInstrumentation
-                                                                        .INSTRUMENTATION_PACKAGE_PREFIX)))
-                                        .and(
-                                                not(
-                                                        nameStartsWith(
-                                                                ApacheKafkaConsumerInstrumentation
-                                                                        .INSTRUMENTATION_PACKAGE_PREFIX))))
                         .type(apacheHttpInstrumentation.getTypeMatcher())
                         .transform(apacheHttpInstrumentation.getTransformer())
                         .type(amazonHttpClientInstrumentation.getTypeMatcher())
@@ -50,7 +37,9 @@ public static void instrument(java.lang.instrument.Instrumentation inst) {
                         .type(apacheKafkaInstrumentation.getTypeMatcher())
                         .transform(apacheKafkaInstrumentation.getTransformer())
                         .type(apacheKafkaConsumerInstrumentation.getTypeMatcher())
-                        .transform(apacheKafkaConsumerInstrumentation.getTransformer());
+                        .transform(apacheKafkaConsumerInstrumentation.getTransformer())
+                        .type(awsLambdaRequestHandlerInstrumentation.getTypeMatcher())
+                        .transform(awsLambdaRequestHandlerInstrumentation.getTransformer());
 
         builder.installOn(inst);
         Logger.debug("Finish Instrumentation");

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

@@ -43,7 +43,7 @@ public static void methodEnter(@Advice.Argument(0) final Request<?> request) {
                 String patchedRoot = spansContainer.getPatchedRoot();
                 request.getHeaders().put("X-Amzn-Trace-Id", patchedRoot);
                 startTimeMap.put(request.hashCode(), System.currentTimeMillis());
-            } catch (Exception e) {
+            } catch (Throwable e) {
                 Logger.error(e, "Failed to send data on http requests");
             }
         }

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

@@ -59,7 +59,7 @@ public static void methodEnter(@Advice.Argument(0) final HttpUriRequest request)
                     return;
                 }
                 startTimeMap.put(request.hashCode(), System.currentTimeMillis());
-            } catch (Exception e) {
+            } catch (Throwable e) {
                 Logger.error(e);
             }
         }

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

@@ -16,9 +16,6 @@
 import org.pmw.tinylog.Logger;
 
 public class ApacheKafkaConsumerInstrumentation implements LumigoInstrumentationApi {
-
-    public static final String INSTRUMENTATION_PACKAGE_PREFIX = "org.apache.kafka.clients.consumer";
-
     @Override
     public ElementMatcher<TypeDescription> getTypeMatcher() {
         return named("org.apache.kafka.clients.consumer.KafkaConsumer");
@@ -48,7 +45,7 @@ public static class ApacheKafkaConsumerAdvice {
         public static void methodEnter(@Advice.FieldValue("clientId") String clientId) {
             try {
                 startTimeMap.put(clientId, System.currentTimeMillis());
-            } catch (Exception e) {
+            } catch (Throwable e) {
                 Logger.error(e);
             }
         }

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

@@ -21,9 +21,6 @@
 import org.pmw.tinylog.Logger;
 
 public class ApacheKafkaProducerInstrumentation implements LumigoInstrumentationApi {
-
-    public static final String INSTRUMENTATION_PACKAGE_PREFIX = "org.apache.kafka.clients.producer";
-
     @Override
     public ElementMatcher<TypeDescription> getTypeMatcher() {
         return named("org.apache.kafka.clients.producer.KafkaProducer");
@@ -50,6 +47,7 @@ public AgentBuilder.Transformer.ForAdvice getTransformer() {
                         ApacheKafkaProducerAdvice.class.getName());
     }
 
+    @SuppressWarnings("unused")
     public static class ApacheKafkaProducerAdvice {
         public static final SpansContainer spansContainer = SpansContainer.getInstance();
 
@@ -78,7 +76,7 @@ public static <K, V> void methodEnter(
                                         .toString()
                                         .substring(0, 10)
                                         .getBytes(StandardCharsets.UTF_8));
-            } catch (Exception e) {
+            } catch (Throwable e) {
                 Logger.error(e);
             }
         }

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

@@ -0,0 +1,74 @@
+package io.lumigo.core.instrumentation.impl;
+
+import com.amazonaws.services.lambda.runtime.Context;
+
+import io.lumigo.core.SpansContainer;
+import io.lumigo.core.instrumentation.LumigoInstrumentationApi;
+import io.lumigo.core.instrumentation.agent.Loader;
+import io.lumigo.core.network.Reporter;
+import io.lumigo.core.utils.EnvUtil;
+import net.bytebuddy.agent.builder.AgentBuilder;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.implementation.bytecode.assign.Assigner.Typing;
+import net.bytebuddy.matcher.ElementMatcher;
+import org.pmw.tinylog.Logger;
+
+import static net.bytebuddy.matcher.ElementMatchers.*;
+
+public class AwsLambdaRequestHandlerInstrumentation implements LumigoInstrumentationApi {
+    @Override
+    public ElementMatcher<TypeDescription> getTypeMatcher() {
+        return hasSuperType(named("com.amazonaws.services.lambda.runtime.RequestHandler"))
+                // we don't want to instrument handlers that implement our interfaces because they are already instrumented
+                .and(not(hasSuperType(named("io.lumigo.handlers.LumigoRequestHandler")))
+                .and(not(hasSuperType(named("io.lumigo.handlers.LumigoRequestStreamHandler")))));
+    }
+
+    @Override
+    public AgentBuilder.Transformer.ForAdvice getTransformer() {
+        return new AgentBuilder.Transformer.ForAdvice()
+                .include(Loader.class.getClassLoader())
+                .advice(
+                        isMethod()
+                        .and(isPublic())
+                        .and(named("handleRequest"))
+                        .and(takesArgument(1, named("com.amazonaws.services.lambda.runtime.Context"))),
+                        HandleRequestAdvice.class.getName());
+    }
+
+    @SuppressWarnings("unused")
+    public static class HandleRequestAdvice {
+        @Advice.OnMethodEnter(suppress = Throwable.class)
+        public static void methodEnter(
+                @Advice.Argument(value = 0, typing = Typing.DYNAMIC) Object input,
+                @Advice.Argument(1) Context context,
+                @Advice.Local("lumigoSpansContainer") SpansContainer spansContainer
+        ) {
+            try {
+                spansContainer = SpansContainer.getInstance();
+                spansContainer.init(new EnvUtil().getEnv(), new Reporter(), context, input);
+                spansContainer.start();
+                Logger.debug("Finish sending start message and instrumentation");
+            } catch (Throwable e) {
+                Logger.error(e, "Failed to init span container");
+            }
+        }
+
+        @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+        public static void methodExit(
+                @Advice.Return(readOnly = false) Object returnValue,
+                @Advice.Thrown Throwable throwable,
+                @Advice.Local("lumigoSpansContainer") SpansContainer spansContainer) {
+            try {
+                if (throwable != null) {
+                    spansContainer.endWithException(throwable);
+                } else {
+                    spansContainer.end(returnValue);
+                }
+            } catch (Throwable e) {
+                Logger.error(e, "Failed to create end span");
+            }
+        }
+    }
+}