feat: simplify logic of handling initial context (#283)

- reduce thread locals
- drop usage of extra executors

james-agent/src/main/java/com/tomtom/james/agent/Agent.java

@@ -24,7 +24,6 @@
 import com.tomtom.james.configuration.AgentConfigurationFactory;
 import com.tomtom.james.configuration.ConfigurationInitializationException;
 import com.tomtom.james.newagent.JVMAgentCleaner;
-import com.tomtom.james.newagent.MethodExecutionContextHelper;
 import com.tomtom.james.publisher.EventPublisherFactory;
 import com.tomtom.james.script.ScriptEngineFactory;
 import com.tomtom.james.store.informationpoints.io.InformationPointStore;
@@ -63,7 +62,7 @@ private static void setupAgent(Instrumentation instrumentation) {
             //InformationPointService informationPointService = new InformationPointServiceImpl(store, instrumentation);
             //controllersManager.initializeControllers(informationPointService, engine, publisher);
 
-            JVMAgentCleaner.init(controllersManager, engine, publisher, MethodExecutionContextHelper::shutdown);
+            JVMAgentCleaner.init(controllersManager, engine, publisher);
             if (configuration.isShutdownHookEnabled()) {
                 ShutdownHook shutdownHook = new ShutdownHook(configuration, JVMAgentCleaner::close);
                 Runtime.getRuntime().addShutdownHook(shutdownHook);

james-agent/src/main/java/com/tomtom/james/informationpoint/advice/ContextAwareAdvice.java

@@ -29,7 +29,6 @@
 import com.tomtom.james.common.api.informationpoint.InformationPoint;
 import com.tomtom.james.common.api.script.RuntimeInformationPointParameter;
 import com.tomtom.james.common.log.Logger;
-import com.tomtom.james.newagent.MethodExecutionContextHelper;
 import org.apache.logging.log4j.util.StackLocatorUtil;
 
 /*
@@ -44,7 +43,8 @@ private ContextAwareAdvice() {
     }
 
     @SuppressWarnings("unused")
-    public static void onEnter(String originTypeName,
+    public static void onEnter(ExecutionContext context,
+                               String originTypeName,
                                String originMethodName,
                                Method origin,
                                Object instance,
@@ -68,18 +68,15 @@ public static void onEnter(String originTypeName,
                 LOG.trace(() -> "onEnter: noInitialContextSupportRequired - skipping");
                 return;
             }
-
-            final String key = MethodExecutionContextHelper.createContextKey(ip);
             LOG.trace(() -> "Initializing custom context setup for the call");
             final Object callContext = ScriptEngineSupplier.get().invokePrepareContext(
                     ip,
                     origin,
                     createParameterList(origin, arguments),
                     instance,
                     Thread.currentThread(),
-                    key);
-
-            MethodExecutionContextHelper.storeContextAsync(key, callContext);
+                    context.toString());
+            context.setInitialContext(callContext);
 
         } catch (Throwable t) {
             LOG.error("Error executing onEnter advice", t);
@@ -91,7 +88,7 @@ public static void onEnter(String originTypeName,
 
     @SuppressWarnings("unused")
     public static void onExit(
-            long _startTime,
+            ExecutionContext context,
             String informationPointClassName,
             String informationPointMethodName,
             Method origin,
@@ -100,8 +97,7 @@ public static void onExit(
             Object returned,
             Throwable thrown) {
         Instant eventTime = Instant.now();
-        Duration executionTime = Duration.ofNanos(System.nanoTime() - _startTime);
-        AutoCloseable closeable = null;
+        Duration executionTime = context.getElapsedTime();
 
         try {
             Optional<InformationPoint> optionalInformationPoint = InformationPointServiceSupplier.get()
@@ -130,20 +126,11 @@ public static void onExit(
                     + ", thrown=" + thrown
                     + "]");
 
-            boolean requireInitialContextCleanup = ip.getRequiresInitialContext();
-            if (requireInitialContextCleanup) {
-                closeable = () -> MethodExecutionContextHelper.removeContextKey(ip);
-            }
-
             if ((sampleRate < 100) && (sampleRate < ThreadLocalRandom.current().nextDouble() * 100)) {
                 LOG.trace(() -> "onExit: Sample skipped (sampleRate=" + sampleRate + ")");
                 return;
             }
 
-            final CompletableFuture<Object> initialContextAsyncProvider = requireInitialContextCleanup
-                    ? MethodExecutionContextHelper.getContextAsync(MethodExecutionContextHelper.getKeyForCurrentFrame(ip))
-                    : CompletableFuture.completedFuture(null);
-
             final String[] callStack = ip.getRequiresCallStack() ? getCallStack() : EMPTY_CALL_STACK;
             if (thrown == null) {
                 if (executionTime.toMillis() < successExecutionThreshold) {
@@ -161,7 +148,7 @@ public static void onExit(
                         executionTime,
                         callStack,
                         returned,
-                        initialContextAsyncProvider
+                        CompletableFuture.completedFuture(context.getInitialContext())
                 );
             } else {
                 LOG.trace(() -> "onExit: Invoking error handler");
@@ -175,20 +162,13 @@ public static void onExit(
                         executionTime,
                         callStack,
                         thrown,
-                        initialContextAsyncProvider
+                        CompletableFuture.completedFuture(context.getInitialContext())
                 );
             }
         } catch (Throwable t) {
             LOG.error("Error executing onExit advice", t);
             throw t;
         } finally {
-            try {
-                if (closeable != null) {
-                    closeable.close();
-                }
-            } catch (Exception e) {
-                LOG.error("Error executing onExit advice (finally)", e);
-            }
             LOG.trace("onExit: END");
         }
     }
@@ -228,12 +208,6 @@ public static List<RuntimeInformationPointParameter> createParameterList(Method
         return result;
     }
 
-    @SuppressWarnings("unused")
-    public static void onEnter(String originTypeName,
-                               String originMethodName) {
-        onEnter(originTypeName, originMethodName, null, null, null);
-    }
-
 
     @SuppressWarnings("unused")
     public static void onExit(

james-agent/src/main/java/com/tomtom/james/informationpoint/advice/ExecutionContext.java

@@ -0,0 +1,24 @@
+package com.tomtom.james.informationpoint.advice;
+
+import java.time.Duration;
+
+public class ExecutionContext {
+    private final long started;
+    private Object initialContext;
+
+    public ExecutionContext() {
+        this.started = System.nanoTime();
+    }
+
+    public Duration getElapsedTime() {
+        return Duration.ofNanos(System.nanoTime() - started);
+    }
+
+    public Object getInitialContext() {
+        return initialContext;
+    }
+
+    public void setInitialContext(final Object initialContext) {
+        this.initialContext = initialContext;
+    }
+}

james-agent/src/main/java/com/tomtom/james/newagent/JVMAgent.java

@@ -173,7 +173,7 @@ private static void setupAgent(Instrumentation inst) {
                     removeInformationPointQueue
                     );
             LOG.trace("initialize controllers time=" + stopwatch.elapsed());
-            JVMAgentCleaner.init(controllersManager, engine, publisher, MethodExecutionContextHelper::shutdown);
+            JVMAgentCleaner.init(controllersManager, engine, publisher);
             if (configuration.isShutdownHookEnabled()) {
                 ShutdownHook shutdownHook = new ShutdownHook(configuration, JVMAgentCleaner::close);
                 Runtime.getRuntime().addShutdownHook(shutdownHook);

james-agent/src/main/java/com/tomtom/james/newagent/MethodExecutionContextHelper.java

@@ -16,93 +16,27 @@
 
 package com.tomtom.james.newagent;
 
-import com.tomtom.james.common.api.informationpoint.InformationPoint;
-import com.tomtom.james.common.log.Logger;
-import com.tomtom.james.util.MoreExecutors;
-
+import com.tomtom.james.informationpoint.advice.ExecutionContext;
 import java.util.ArrayDeque;
-import java.util.UUID;
-import java.util.WeakHashMap;
-import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.ExecutorService;
-import java.util.concurrent.TimeUnit;
 
+/**
+ * Helper class for managing method execution context - used by GroovyJames to inject method entry and exit hooks.
+ */
 public class MethodExecutionContextHelper {
-    private static final Logger LOG = Logger.getLogger(MethodExecutionContextHelper.class);
-
-    public static ThreadLocal<ArrayDeque<String>> keysStack = ThreadLocal.withInitial(() -> new ArrayDeque<>(8));
-    private static WeakHashMap<String, Object> contextStore = new WeakHashMap<>();
 
-    private static ExecutorService contextStoreAccessExecutor = MoreExecutors
-            .createNamedDaemonExecutorService("james-context-access-%d", 1);
-
-    private static ExecutorService contextCallbackExecutor = MoreExecutors
-            .createNamedDaemonExecutorService("james-context-callback-%d", 5);
-
-    public static void shutdown() {
-        try {
-            contextStoreAccessExecutor.shutdown();
-            contextStoreAccessExecutor.awaitTermination(5, TimeUnit.SECONDS);
-        } catch (InterruptedException e) {
-            LOG.warn("Executor contextStoreAccessExecutor shutdown interrupted " + e);
-        }
-        try {
-            contextCallbackExecutor.shutdown();
-            contextCallbackExecutor.awaitTermination(5, TimeUnit.SECONDS);
-        } catch (InterruptedException e) {
-            LOG.warn("Executor contextCallbackExecutor shutdown interrupted " + e);
-        }
-    }
+    private static ThreadLocal<ArrayDeque<ExecutionContext>> contextStack = ThreadLocal.withInitial(() -> new ArrayDeque<>(8));
 
-    public static String createContextKey(final InformationPoint ip) {
-        final String contextKey = ip.toString() + "-" + UUID.randomUUID();
-        keysStack.get().push(contextKey);
-        return contextKey;
+    public static ExecutionContext executionStarted() {
+        final ExecutionContext context = new ExecutionContext();
+        contextStack.get().push(context);
+        return context;
     }
 
-    public static String getKeyForCurrentFrame(final InformationPoint informationPoint) {
-        final String key = keysStack.get().peek();
-        if (key != null && key.startsWith(informationPoint.toString())) {
-            return key;
-        }
-        return null;
+    public static ExecutionContext getExecutionContext() {
+        return contextStack.get().peek();
     }
 
-    public static String removeContextKey(final InformationPoint informationPoint) {
-        String key = getKeyForCurrentFrame(informationPoint);
-        if (key != null) {
-            return keysStack.get().pop();
-        }
-        return null;
-    }
-
-    public static CompletableFuture<Object> storeContextAsync(final String key, final Object value) {
-        final CompletableFuture result = new CompletableFuture();
-        contextStoreAccessExecutor.submit(() -> {
-            contextStore.put(key, value);
-            CompletableFuture.supplyAsync(() -> result.complete(value), contextCallbackExecutor);
-        });
-        return result;
-    }
-
-    public static CompletableFuture<Object> getContextAsync(final String key) {
-        if (key == null) {
-            return CompletableFuture.completedFuture(null);
-        }
-        final CompletableFuture result = new CompletableFuture();
-        contextStoreAccessExecutor.submit(() -> {
-            if (contextStore.containsKey(key)) {
-                final Object context = contextStore.get(key);
-                CompletableFuture.supplyAsync(() -> result.complete(context), contextCallbackExecutor);
-                return;
-            }
-
-            CompletableFuture.supplyAsync(() -> {
-                final String msg = String.format("Key '%s' not found in context container", key);
-                result.completeExceptionally(new IllegalArgumentException(msg));
-                return null;
-            }, contextCallbackExecutor);
-        });
-        return result;
+    public static ExecutionContext executionFinished() {
+        return contextStack.get().pop();
     }
 }

james-agent/src/main/java/com/tomtom/james/newagent/james/GroovyJames.java

@@ -17,18 +17,9 @@ public GroovyJames(Queue<JamesObjective> objectives, long sleepTime) {
         this.setName(getClass().getSimpleName());
     }
 
-    // TODO check double if that is all chars that we need to escape
-    private String escapeScriptString(String script) {
-        return script.replace("\\", "\\\\")
-                .replace("\"", "\\\"")
-                .replace("\r", "\\r")
-                .replace("\n", "\\n");
-    }
-
     protected void insertBefore(CtMethod method, ExtendedInformationPoint informationPoint) throws CannotCompileException {
         StringBuilder s = new StringBuilder("");
-        s.append(" com.tomtom.james.newagent.MethodExecutionTimeHelper.executionStarted();\n");
-        s.append(" com.tomtom.james.informationpoint.advice.ContextAwareAdvice.onEnter(");
+        s.append(" com.tomtom.james.informationpoint.advice.ContextAwareAdvice.onEnter( com.tomtom.james.newagent.MethodExecutionContextHelper.executionStarted(), \n");
         s.append("\"" + informationPoint.getClassName() + "\", ");
         s.append("\"" + informationPoint.getMethodName() + "\", ");
         s.append(informationPoint.getMethodBodyClassName() + ".class.getDeclaredMethod(\"" + informationPoint.getMethodName() + "\",$sig), "); // method
@@ -42,9 +33,8 @@ protected void insertBefore(CtMethod method, ExtendedInformationPoint informatio
     }
 
     protected void insertAfter(CtMethod method, ExtendedInformationPoint informationPoint) throws CannotCompileException {
-        String script = escapeScriptString(informationPoint.getScript().get());
         StringBuilder s = new StringBuilder();
-        s.append(" com.tomtom.james.informationpoint.advice.ContextAwareAdvice.onExit( com.tomtom.james.newagent.MethodExecutionTimeHelper.getStartTime(), \n");
+        s.append(" com.tomtom.james.informationpoint.advice.ContextAwareAdvice.onExit( com.tomtom.james.newagent.MethodExecutionContextHelper.getExecutionContext(), \n");
         s.append("\"" + informationPoint.getClassName() + "\", ");
         s.append("\"" + informationPoint.getMethodName() + "\", ");
         s.append(informationPoint.getMethodBodyClassName() + ".class.getDeclaredMethod(\"" + informationPoint.getMethodName() + "\",$sig), "); // method
@@ -64,9 +54,8 @@ protected void insertAfter(CtMethod method, ExtendedInformationPoint information
     }
 
     protected void addCatch(ClassPool pool, CtMethod method, ExtendedInformationPoint informationPoint) throws CannotCompileException, NotFoundException {
-        String script = escapeScriptString(informationPoint.getScript().get());
         StringBuilder s = new StringBuilder();
-        s.append(" com.tomtom.james.informationpoint.advice.ContextAwareAdvice.onExit( com.tomtom.james.newagent.MethodExecutionTimeHelper.getStartTime(), ");
+        s.append(" com.tomtom.james.informationpoint.advice.ContextAwareAdvice.onExit( com.tomtom.james.newagent.MethodExecutionContextHelper.getExecutionContext(), ");
         s.append("\"" + informationPoint.getClassName() + "\", ");
         s.append("\"" + informationPoint.getMethodName() + "\", ");
         s.append(informationPoint.getMethodBodyClassName() + ".class.getDeclaredMethod(\"" + informationPoint.getMethodName() + "\",$sig), "); // method
@@ -87,7 +76,7 @@ protected void addCatch(ClassPool pool, CtMethod method, ExtendedInformationPoin
 
         // finally block
         StringBuilder f = new StringBuilder("");
-        f.append(" com.tomtom.james.newagent.MethodExecutionTimeHelper.executionFinished(); \n");
+        f.append(" com.tomtom.james.newagent.MethodExecutionContextHelper.executionFinished(); \n");
         method.insertAfter(f.toString(), true);
     }
 }