Replace the use of threadlocal with @Enter annotation

Signed-off-by: Paolo Di Tommaso <[email protected]>

src/main/java/io/seqera/debug/AllocationContext.java

@@ -9,45 +9,39 @@
  */
 public class AllocationContext {
 
-    static final ThreadLocal<AllocationContext> instance = new ThreadLocal<>();
-
     /**
      * The name of method used to allocate the memory
      */
-    String name;
+    final String name;
 
     /**
      * The size of the allocated memoty in bytes
      */
-    long size;
+    final long size;
 
     /**
      * The instant when the allocation happened
      */
-    Instant createdAt;
+    final Instant createdAt;
 
     /**
      * The stack trace at the allocation point
      */
-    String stackTrace;
+    final String stackTrace;
 
     /**
      * The address of the memory allocated
      */
-    long address;
+    final long address;
 
-    AllocationContext(String name, long size, Instant ts, String stackTrace) {
+    AllocationContext(String name, long size, long address, Instant ts, String stackTrace) {
         this.name = name;
         this.size = size;
+        this.address = address;
         this.createdAt = ts;
         this.stackTrace = stackTrace;
     }
 
-    public AllocationContext withAddress(long address) {
-        this.address = address;
-        return this;
-    }
-
     @Override
     public String toString() {
         return "AllocationContext{" +
@@ -73,14 +67,13 @@ public static String dumpStack() {
         return result.toString();
     }
 
-    public static void create(String name, long size) {
-        AllocationContext context = new AllocationContext(name, size, Instant.now(), dumpStack());
-        instance.set(context);
+    public static AllocationContext create(String name, long size, long address) {
+        return new AllocationContext(
+                name,
+                size,
+                address,
+                Instant.now(),
+                dumpStack());
     }
 
-    public static AllocationContext get() {
-        AllocationContext result = instance.get();
-        instance.remove();
-        return result;
-    }
 }

src/main/java/io/seqera/debug/advice/AllocateMemAdvice.java

@@ -10,13 +10,15 @@
 public class AllocateMemAdvice {
 
     @Advice.OnMethodEnter
-    public static void before(@Advice.Argument(0) long size) {
-        AllocationContext.create("Unsafe.allocateMemory", size);
+    public static long before(@Advice.Argument(0) long size) {
+        // Return the size so it can be captured by @Enter in @OnMethodExit
+        return size;
     }
 
     @Advice.OnMethodExit
-    public static void after(@Advice.Return long address) {
-        LeaksManager.register(address, AllocationContext.get().withAddress(address));
+    public static void after(@Advice.Return long address, @Advice.Enter long size) {
+        AllocationContext c = AllocationContext.create("Unsafe.allocateMemory", size, address);
+        LeaksManager.register(address, c);
     }
 
 }

src/test/groovy/io/seqera/debug/AllocationContextTest.groovy

@@ -0,0 +1,25 @@
+package io.seqera.debug
+
+import java.time.Instant
+
+import spock.lang.Specification
+
+/**
+ *
+ * @author Paolo Di Tommaso <[email protected]>
+ */
+class AllocationContextTest extends Specification {
+
+    def 'should create alloc context' () {
+        when:
+        def c1 = AllocationContext.create('foo', 1, 2)
+
+        then:
+        c1.name == 'foo'
+        c1.size == 1
+        c1.address == 2
+        c1.createdAt.compareTo(Instant.now())<=0
+        c1.stackTrace.startsWith('  java.base/jdk.internal.reflect.NativeMethodAccessorImpl')
+    }
+
+}

src/test/groovy/io/seqera/debug/LeaksManagerTest.groovy

@@ -14,9 +14,9 @@ class LeaksManagerTest extends Specification {
     def 'should find all leaks older than specified duration' (){
         given:
         def all = [
-                new AllocationContext('foo', 1000, Instant.now().minusMillis(20_000), null),
-                new AllocationContext('foo', 1000, Instant.now().minusMillis(10_000), null),
-                new AllocationContext('foo', 1000, Instant.now().minusMillis(1_000), null),
+                new AllocationContext('foo', 1000, 1,  Instant.now().minusMillis(20_000), null),
+                new AllocationContext('foo', 1000, 2, Instant.now().minusMillis(10_000), null),
+                new AllocationContext('foo', 1000, 3, Instant.now().minusMillis(1_000), null),
         ]
 
         when: