Add support for Context ID functionality

agent/pom.xml

@@ -0,0 +1,191 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project xmlns="http://maven.apache.org/POM/4.0.0"
+         xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+    <modelVersion>4.0.0</modelVersion>
+
+    <groupId>io.github.dpsoft</groupId>
+    <artifactId>ap-agent</artifactId>
+    <version>0.1.2-SNAPSHOT</version>
+    <packaging>jar</packaging>
+
+    <properties>
+        <me.bechberger.ap-loader.version>2.9-4</me.bechberger.ap-loader.version>
+        <io.vavr.version>0.10.3</io.vavr.version>
+        <com.typesafe.config.version>1.4.2</com.typesafe.config.version>
+        <org.tinylog.version>2.5.0</org.tinylog.version>
+        <me.bechberger.version>0.0.2-SNAPSHOT</me.bechberger.version>
+        <java.version>11</java.version>
+        <maven.compiler.source>11</maven.compiler.source>
+        <maven.compiler.target>11</maven.compiler.target>
+    </properties>
+
+    <repositories>
+        <repository>
+            <id>oss-snapshots</id>
+            <name>Oss Snapshots</name>
+            <url>https://s01.oss.sonatype.org/content/repositories/snapshots/</url>
+        </repository>
+        <repository>
+            <id>maven-central</id>
+            <name>Maven central</name>
+            <url>https://repo.maven.apache.org/maven2/</url>
+        </repository>
+        <repository>
+            <id>local</id>
+            <name>Maven local</name>
+            <url>file:////home/diego/.m2/repository</url>
+        </repository>
+    </repositories>
+
+    <dependencies>
+        <dependency>
+            <groupId>me.bechberger</groupId>
+            <artifactId>ap-loader-all</artifactId>
+            <version>${me.bechberger.ap-loader.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>io.vavr</groupId>
+            <artifactId>vavr</artifactId>
+            <version>${io.vavr.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>com.typesafe</groupId>
+            <artifactId>config</artifactId>
+            <version>${com.typesafe.config.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>org.tinylog</groupId>
+            <artifactId>tinylog-api</artifactId>
+            <version>${org.tinylog.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>org.tinylog</groupId>
+            <artifactId>tinylog-impl</artifactId>
+            <version>${org.tinylog.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>me.bechberger</groupId>
+            <artifactId>jfrtofp</artifactId>
+            <version>${me.bechberger.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>io.github.dpsoft</groupId>
+            <artifactId>context-api</artifactId>
+            <version>0.1.2-SNAPSHOT</version>
+            <scope>provided</scope>
+        </dependency>
+    </dependencies>
+    <build>
+        <resources>
+            <resource>
+                <directory>src/main/resources</directory>
+                <filtering>true</filtering>
+                <includes>
+                    <include>**/build-info.properties</include>
+                </includes>
+            </resource>
+            <resource>
+                <directory>src/main/resources</directory>
+                <filtering>false</filtering>
+                <excludes>
+                    <exclude>**/build-info.properties</exclude>
+                </excludes>
+            </resource>
+        </resources>
+        <plugins>
+            <plugin>
+                <groupId>org.apache.maven.plugins</groupId>
+                <artifactId>maven-shade-plugin</artifactId>
+                <version>3.4.1</version>
+                <executions>
+                    <execution>
+                        <phase>package</phase>
+                        <goals>
+                            <goal>shade</goal>
+                        </goals>
+                        <configuration>
+                            <createDependencyReducedPom>true</createDependencyReducedPom>
+                            <dependencyReducedPomLocation>${project.build.directory}/dependency-reduced-pom.xml</dependencyReducedPomLocation>
+                            <minimizeJar>true</minimizeJar>
+                            <artifactSet>
+                                <excludes>
+                                    <exclude>org.codehaus.mojo</exclude>
+                                </excludes>
+                            </artifactSet>
+                            <transformers>
+                                <transformer
+                                    implementation="org.apache.maven.plugins.shade.resource.ManifestResourceTransformer">
+                                    <manifestEntries>
+                                        <Premain-Class>io.github.dpsoft.ap.Agent</Premain-Class>
+                                    </manifestEntries>
+                                </transformer>
+                                <transformer
+                                    implementation="org.apache.maven.plugins.shade.resource.ServicesResourceTransformer"/>
+                            </transformers>
+                            <filters>
+                                <filter>
+                                    <artifact>*:*</artifact>
+                                    <excludes>
+                                        <exclude>META-INF/maven/**</exclude>
+                                        <exclude>META-INF/*.SF</exclude>
+                                        <exclude>META-INF/*.DSA</exclude>
+                                        <exclude>META-INF/*.RSA</exclude>
+                                        <exclude>META-INF/*.md</exclude>
+                                        <exclude>META-INF/LICENSE</exclude>
+                                    </excludes>
+                                </filter>
+                            </filters>
+                            <relocations>
+                                <relocation>
+                                    <pattern>com.typesafe.config</pattern>
+                                    <shadedPattern>io.github.dpsoft.ap.libs.com.typesafe.config</shadedPattern>
+                                </relocation>
+                                <relocation>
+                                    <pattern>io.vavr</pattern>
+                                    <shadedPattern>io.github.dpsoft.ap.libs.io.vavr</shadedPattern>
+                                </relocation>
+                                <relocation>
+                                    <pattern>org.tinylog</pattern>
+                                    <shadedPattern>io.github.dpsoft.ap.libs.org.tinylog</shadedPattern>
+                                </relocation>
+                                <relocation>
+                                    <pattern>one</pattern>
+                                    <shadedPattern>io.github.dpsoft.ap.libs.one</shadedPattern>
+                                </relocation>
+                                <relocation>
+                                    <pattern>dev.dirs</pattern>
+                                    <shadedPattern>io.github.dpsoft.ap.libs.dev.dirs</shadedPattern>
+                                </relocation>
+                                <relocation>
+                                    <pattern>me.bechberger</pattern>
+                                    <shadedPattern>io.github.dpsoft.ap.libs.me.bechberger</shadedPattern>
+                                </relocation>
+                                <relocation>
+                                    <pattern>kotlin</pattern>
+                                    <shadedPattern>io.github.dpsoft.ap.libs.kotlin</shadedPattern>
+                                </relocation>
+                                <relocation>
+                                    <pattern>kotlinx</pattern>
+                                    <shadedPattern>io.github.dpsoft.ap.libs.kotlinx</shadedPattern>
+                                </relocation>
+                                <relocation>
+                                    <pattern>org.jetbrains</pattern>
+                                    <shadedPattern>io.github.dpsoft.ap.libs.org.jetbrains</shadedPattern>
+                                </relocation>
+                                <relocation>
+                                    <pattern>org.jline</pattern>
+                                    <shadedPattern>io.github.dpsoft.ap.libs.org.jline</shadedPattern>
+                                </relocation>
+                                <relocation>
+                                    <pattern>org.objectweb</pattern>
+                                    <shadedPattern>io.github.dpsoft.ap.libs.org.objectweb</shadedPattern>
+                                </relocation>
+                            </relocations>
+                        </configuration>
+                    </execution>
+                </executions>
+            </plugin>
+        </plugins>
+    </build>
+</project>

src/main/java/io/github/dpsoft/ap/Agent.java → agent/src/main/java/io/github/dpsoft/ap/Agent.java

@@ -1,14 +1,14 @@
 package io.github.dpsoft.ap;
 
-import io.github.dpsoft.ap.handler.AsyncProfilerHandler;
 import io.github.dpsoft.ap.util.Banner;
 import io.github.dpsoft.ap.util.Runner;
 import io.github.dpsoft.ap.util.Server;
+import io.github.dpsoft.ap.handler.AsyncProfilerHandler;
 
 import java.lang.instrument.Instrumentation;
 
 public final class Agent {
-    public static void premain(String args, Instrumentation inst)  {
+    public static void premain(String args, Instrumentation instrumentation)  {
         Runner.runWith((profiler, configuration) -> {
 
             Banner.show(configuration);

src/main/java/io/github/dpsoft/ap/config/AgentConfiguration.java → agent/src/main/java/io/github/dpsoft/ap/config/AgentConfiguration.java

@@ -14,6 +14,7 @@ public final class AgentConfiguration {
     public final Handler handler;
     public final boolean showBanner;
 
+
     private AgentConfiguration(Config config) {
         this.server = new Server(config);
         this.handler = new Handler(config);

agent/src/main/java/io/github/dpsoft/ap/context/APContextListener.java

@@ -0,0 +1,18 @@
+package io.github.dpsoft.ap.context;
+
+import io.github.dpsoft.ap.context.api.context.Context;
+import io.github.dpsoft.ap.context.api.storage.ContextStorageListener;
+
+/**
+ * This listener is used to store the current context id in AsyncProfiler and restore it when the context is restored.
+ */
+public final class APContextListener implements ContextStorageListener {
+    @Override
+    public void onContextStored(Context context) {
+        System.out.println("stored::AsyncProfiler.setContextId(" + context.get(Context.ContextID) + ")");
+    }
+    @Override
+    public void onContextRestored(Context context) {
+        System.out.println("restored::AsyncProfiler.setContextId(" + context.get(Context.ContextID) + ")");
+    }
+}

agent/src/main/resources/META-INF/services/io.github.dpsoft.ap.context.api.storage.ContextStorageListener

@@ -0,0 +1 @@
+io.github.dpsoft.ap.context.APContextListener

src/main/resources/reference.conf → agent/src/main/resources/reference.conf

@@ -9,6 +9,7 @@ ap-agent {
     port: 8080
   }
 
+  # The agent will recive requests on this path
   handler {
     go-mode: false
     go-context: "/debug/pprof/profile"

agent/src/test/resources/application.conf

@@ -0,0 +1,10 @@
+ap-agent {
+
+  #The agent banner will be shown on the console
+  show-banner = false
+
+  # Bytecode instrumentation is disabled by default
+  context-instrumenter {
+    enabled = true
+  }
+}

context-api/pom.xml

@@ -0,0 +1,23 @@
+<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+    <modelVersion>4.0.0</modelVersion>
+
+    <groupId>io.github.dpsoft</groupId>
+    <artifactId>context-api</artifactId>
+    <version>0.1.2-SNAPSHOT</version>
+    <packaging>jar</packaging>
+
+    <dependencies>
+        <dependency>
+            <groupId>org.junit.jupiter</groupId>
+            <artifactId>junit-jupiter</artifactId>
+            <version>5.9.2</version>
+            <scope>test</scope>
+        </dependency>
+    </dependencies>
+
+    <properties>
+        <java.version>11</java.version>
+        <maven.compiler.source>11</maven.compiler.source>
+        <maven.compiler.target>11</maven.compiler.target>
+    </properties>
+</project>

context-api/src/main/java/io/github/dpsoft/ap/context/api/ContextHandler.java

@@ -0,0 +1,42 @@
+package io.github.dpsoft.ap.context.api;
+
+
+import io.github.dpsoft.ap.context.api.context.Context;
+import io.github.dpsoft.ap.context.api.storage.ContextStorage;
+import io.github.dpsoft.ap.context.api.storage.Storage;
+
+import java.util.function.Supplier;
+
+/**
+ * This class is used to store and retrieve the current context.
+ * Its implementation is Noop by default, but it can be replaced by an agent in runtime.
+ */
+public class ContextHandler  {
+    /**
+     * Returns the current context.
+     */
+    public static Context currentContext() { return ContextStorage.INSTANCE.currentContext();}
+    /**
+     * Stores the current context and returns a scope that can be used to restore it.
+     */
+    public static Storage.Scope storeContext(Context context) { return ContextStorage.INSTANCE.storeContext(context);}
+
+    /**
+     * Runs the given supplier with the given context.
+     * The context is restored after the supplier returns.
+     */
+    public static <T> T runWithContext(Context context, Supplier<T> supplier) {
+        try (var ignored = storeContext(context)) {
+            return supplier.get();
+        }
+    }
+    /**
+     * Runs the given runnable with the given context.
+     * The context is restored after the runnable returns.
+     */
+    public static void runWitContext(Context context, Runnable runnable) {
+        try (var ignored = storeContext(context)) {
+            runnable.run();
+        }
+    }
+}

context-api/src/main/java/io/github/dpsoft/ap/context/api/context/Context.java

@@ -0,0 +1,66 @@
+package io.github.dpsoft.ap.context.api.context;
+
+import java.util.HashMap;
+import java.util.Map;
+
+public final class Context {
+    public static final Context EMPTY = new Context(Map.of(), Labels.EMPTY);
+
+    public static final Context.Key<Long> ContextID = new Context.Key<>("contextId", 0L);
+
+    private final Map<String, Object> underlying;
+
+    private final Labels labels;
+
+    public static Context of(Labels labels) {
+        return new Context(Map.of(),  labels);
+    }
+
+    public static <T> Context of(Context.Key<T> key, T value) {
+        return new Context(Map.of(key.name, value), Labels.EMPTY);
+    }
+
+    public static <T> Context of(Context.Key<T> key, T value, Labels labels) {
+        return new Context(Map.of(key.name, value), labels);
+    }
+
+    private Context(Map<String, Object> underling, Labels labels) {
+        this.labels = labels;
+        this.underlying = underling;
+    }
+
+    public <T> Context withEntry(Context.Key<T> key, T value) {
+        var x = new HashMap<String, Object>(underlying.size() + 1);
+        x.putAll(underlying);
+        x.put(key.name, value);
+
+        return new Context(x, labels);
+    }
+
+    public <T> Context withTag(String key, String value) {
+        return new Context(underlying, labels.withLabel(key, value));
+    }
+
+    public Labels labels() {
+        return labels;
+    }
+
+
+    public <T> T get(Context.Key<T> key) {
+        return (T) underlying.getOrDefault(key.name, key.emptyValue);
+    }
+
+    public static <T> Context.Key<T> key(String name, T emptyValue) {
+        return new Context.Key<>(name, emptyValue);
+    }
+
+    public static class Key<T> {
+        public final String name;
+        public final T emptyValue;
+
+        public Key(String name, T emptyValue) {
+            this.name = name;
+            this.emptyValue = emptyValue;
+        }
+    }
+}