feat: support auto-trace with layer (#73)

Java Tracer - Support auto-trace with layer RD-13012

Author: nadav3396

.bumpversion.cfg

@@ -11,4 +11,6 @@ tag = True
 
 [bumpversion:file:src/main/java/io/lumigo/core/configuration/Configuration.java]
 
+[bumpversion:file:layer/generate.sh]
+
 [bumpversion:file:src/main/resources/lumigo-version.txt]

README.md

@@ -88,6 +88,13 @@ Find the latest version here (the format of the version will be n.n.n):
             }
         }
     ```
+  
+## Lambda Auto tracing with lambda layer
+
+* 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.
+
 
 ### Configuration
 
@@ -120,7 +127,7 @@ class MyFunction implements RequestHandler<String, String> {
 ### Support Java 11 and Above
 
 Add the environment variable `JAVA_TOOL_OPTIONS` to your Lambda functions and set it to
-`-Djdk.attach.allowAttachSelf=true` in addition to the manual code mentioned above.
+`-Djdk.attach.allowAttachSelf=true` in addition to the manual code mentioned above (This is not needed for the auto trace with lambda layer).
 
 ### Supported Instrumentation Libraries
 

agent/pom.xml

@@ -32,12 +32,6 @@
             <organization>Lumigo</organization>
             <organizationUrl>https://lumigo.io/</organizationUrl>
         </developer>
-        <developer>
-            <name>Uri Parush</name>
-            <email>uri@lumigo.io</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

@@ -11,19 +11,39 @@
 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-java/lumigo-tracer.jar";
+
+    public static void premain(String agentArgs, Instrumentation inst) {
+        if (!isKillSwitchOn()) {
+            agentmain(agentArgs, inst);
+        }
+    }
+
     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()};
+                urls =
+                        new URL[] {
+                            new File("/var/task/").toURI().toURL(),
+                            new File(LUMIGO_JAVA_TRACER_PATH).toURI().toURL()
+                        };
             }
+            installTracerJar(inst);
             URLClassLoader newClassLoader = new URLClassLoader(urls, null);
             Thread.currentThread().setContextClassLoader(newClassLoader);
             final Class<?> loader =
@@ -35,6 +55,14 @@ 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) {
+            e.printStackTrace();
+        }
+    }
+
     public static URL[] getUrls() {
         List<URL> jars = new LinkedList<>();
         try (Stream<Path> paths = Files.walk(Paths.get("/var/task/lib"))) {
@@ -56,4 +84,9 @@ public static URL[] getUrls() {
         }
         return jars.toArray(new URL[jars.size()]);
     }
+
+    public static boolean isKillSwitchOn() {
+        String value = System.getenv("LUMIGO_SWITCH_OFF");
+        return "true".equalsIgnoreCase(value);
+    }
 }

layers/LAYERS.md

@@ -0,0 +1,4 @@
+Java Layers 
+----
+| Region | ARN |
+| --- | --- |

pom.xml

@@ -116,28 +116,33 @@
             <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 -->
@@ -182,6 +187,7 @@
             <artifactId>json</artifactId>
             <version>20210307</version>
         </dependency>
+
         <!-- Testing dependencies -->
         <dependency>
             <groupId>org.junit.jupiter</groupId>
@@ -372,4 +378,32 @@
             </plugin>
         </plugins>
     </build>
+    <profiles>
+        <profile>
+            <id>shade</id>
+            <activation>
+                <property>
+                    <name>includeShade</name>
+                    <value>true</value>
+                </property>
+            </activation>
+            <build>
+                <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>
+                </plugins>
+            </build>
+        </profile>
+    </profiles>
 </project>

scripts/bd_to_prod.sh

@@ -40,6 +40,26 @@ mvn -f agent/pom.xml nexus-staging:release
 mvn -Dmaven.test.skip=true -Dfindbugs.skip=true clean deploy
 mvn nexus-staging:release
 
+echo "Creating lumigo-java-tracer layer"
+./scripts/prepare_layer_files.sh
+
+echo "Creating layer latest version arn table md file (LAYERS.md)"
+commit_version="$(git describe --abbrev=0 --tags)"
+../utils/common_bash/create_layer.sh \
+    --layer-name lumigo-java-tracer \
+    --region ALL \
+    --package-folder lumigo-java \
+    --version "$commit_version" \
+    --runtimes "java11 java17 java21"
+
+cd ../larn && npm i -g
+larn -r java11 -n layers/LAYERS.md --filter lumigo-java-tracer -p ~/java-tracer
+cd ../java-tracer
+
+git add layers/LAYERS.md
+git commit -m "docs: update layers md [skip ci]"
+git push origin master
+
 echo "Create release tag"
 push_tags
 

scripts/prepare_layer_files.sh

@@ -0,0 +1,11 @@
+#!/usr/bin/env bash
+
+mkdir -p lumigo-java
+
+MVN_DEFAULT_FLAGS="-Dhttp.keepAlive=false -Dmaven.wagon.http.pool=false -Dmaven.wagon.httpconnectionManager.ttlSeconds=120 -Dmaven.wagon.http.retryHandler.requestSentEnabled=true -Dmaven.wagon.http.retryHandler.count=10"
+
+mvn "$MVN_DEFAULT_FLAGS" -Dmaven.test.skip=true -Dfindbugs.skip=true -DincludeShade=true clean package --quiet
+mvn "$MVN_DEFAULT_FLAGS" -f agent/pom.xml clean package --quiet
+
+cp target/java-tracer-1.0.47.jar lumigo-java/lumigo-tracer.jar
+cp agent/target/lumigo-agent-1.0.47.jar lumigo-java/lumigo-agent.jar

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

@@ -1,12 +1,10 @@
 package io.lumigo.core.instrumentation.agent;
 
-import static net.bytebuddy.matcher.ElementMatchers.nameStartsWith;
-import static net.bytebuddy.matcher.ElementMatchers.not;
-
 import io.lumigo.core.instrumentation.impl.*;
 import net.bytebuddy.agent.builder.AgentBuilder;
 import org.pmw.tinylog.Logger;
 
+@SuppressWarnings("unused")
 public class Loader {
     public static void instrument(java.lang.instrument.Instrumentation inst) {
         Logger.debug("Start Instrumentation");
@@ -19,28 +17,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 +32,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");
             }
         }