feat: better handle case when existing IP starts to support initial context (#282)

it could be already running when new IP is loaded leading to situation when initial context is not loaded

Author: piotrwielgolaski-tomtom

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

@@ -69,7 +69,7 @@ public static void onEnter(String originTypeName,
                 return;
             }
 
-            final String key = MethodExecutionContextHelper.createContextKey();
+            final String key = MethodExecutionContextHelper.createContextKey(ip);
             LOG.trace(() -> "Initializing custom context setup for the call");
             final Object callContext = ScriptEngineSupplier.get().invokePrepareContext(
                     ip,
@@ -101,7 +101,7 @@ public static void onExit(
             Throwable thrown) {
         Instant eventTime = Instant.now();
         Duration executionTime = Duration.ofNanos(System.nanoTime() - _startTime);
-        boolean requireInitialContextCleanup = false;
+        AutoCloseable closeable = null;
 
         try {
             Optional<InformationPoint> optionalInformationPoint = InformationPointServiceSupplier.get()
@@ -110,46 +110,49 @@ public static void onExit(
                 LOG.trace(() -> "onExit: skipping because information point is gone");
                 return;
             }
-            final InformationPoint informationPoint = optionalInformationPoint.get();
-            final long successExecutionThreshold = informationPoint.getSuccessExecutionThreshold();
-            final double sampleRate = getSampleRate(thrown, informationPoint);
+            final InformationPoint ip = optionalInformationPoint.get();
+            final long successExecutionThreshold = ip.getSuccessExecutionThreshold();
+            final double sampleRate = getSampleRate(thrown, ip);
 
             LOG.trace(() -> "onExit: START "
                     + "[origin=" + origin
                     + ", informationPointClassName=" + informationPointClassName
                     + ", informationPointMethodName=" + informationPointMethodName
-                    + ", baseScript=" + (informationPoint.getBaseScript().isPresent())
-                    + ", script=" + (informationPoint.getScript().isPresent())
-                    + ", sampleRate=" + informationPoint.getSampleRate()
-                    + ", successSampleRate=" + informationPoint.getSuccessSampleRate()
-                    + ", errorSampleRate=" + informationPoint.getErrorSampleRate()
+                    + ", baseScript=" + (ip.getBaseScript().isPresent())
+                    + ", script=" + (ip.getScript().isPresent())
+                    + ", sampleRate=" + ip.getSampleRate()
+                    + ", successSampleRate=" + ip.getSuccessSampleRate()
+                    + ", errorSampleRate=" + ip.getErrorSampleRate()
                     + ", successExecutionThreshold=" + successExecutionThreshold
                     + ", instance=" + instance
                     + ", arguments=" + Arrays.asList(arguments)
                     + ", returned=" + returned
                     + ", thrown=" + thrown
                     + "]");
 
-            requireInitialContextCleanup = informationPoint.getRequiresInitialContext();
+            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())
+                    ? MethodExecutionContextHelper.getContextAsync(MethodExecutionContextHelper.getKeyForCurrentFrame(ip))
                     : CompletableFuture.completedFuture(null);
 
-            final String[] callStack = informationPoint.getRequiresCallStack() ? getCallStack() : EMPTY_CALL_STACK;
+            final String[] callStack = ip.getRequiresCallStack() ? getCallStack() : EMPTY_CALL_STACK;
             if (thrown == null) {
                 if (executionTime.toMillis() < successExecutionThreshold) {
                     LOG.trace(() -> "onExit: ExecutionTime skipped (executionTime=" + executionTime.toMillis() + ")");
                     return;
                 }
                 LOG.trace(() -> "onExit: Invoking success handler");
                 ScriptEngineSupplier.get().invokeSuccessHandler(
-                        informationPoint,
+                        ip,
                         origin,
                         createParameterList(origin, arguments),
                         instance,
@@ -163,7 +166,7 @@ public static void onExit(
             } else {
                 LOG.trace(() -> "onExit: Invoking error handler");
                 ScriptEngineSupplier.get().invokeErrorHandler(
-                        informationPoint,
+                        ip,
                         origin,
                         createParameterList(origin, arguments),
                         instance,
@@ -179,8 +182,12 @@ public static void onExit(
             LOG.error("Error executing onExit advice", t);
             throw t;
         } finally {
-            if (requireInitialContextCleanup) {
-                MethodExecutionContextHelper.removeContextKey();
+            try {
+                if (closeable != null) {
+                    closeable.close();
+                }
+            } catch (Exception e) {
+                LOG.error("Error executing onExit advice (finally)", e);
             }
             LOG.trace("onExit: END");
         }

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

@@ -16,6 +16,7 @@
 
 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;
 
@@ -53,18 +54,26 @@ public static void shutdown() {
         }
     }
 
-    public static String createContextKey() {
-        final String contextKey = UUID.randomUUID().toString();
+    public static String createContextKey(final InformationPoint ip) {
+        final String contextKey = ip.toString() + "-" + UUID.randomUUID();
         keysStack.get().push(contextKey);
         return contextKey;
     }
 
-    public static String getKeyForCurrentFrame() {
-        return keysStack.get().peek();
+    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 String removeContextKey() {
-        return keysStack.get().pop();
+    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) {
@@ -77,6 +86,9 @@ public static CompletableFuture<Object> storeContextAsync(final String key, fina
     }
 
     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)) {

james-agent/src/test/groovy/com/tomtom/james/newagent/MethodExecutionContextHelperSpec.groovy

@@ -16,6 +16,7 @@
 
 package com.tomtom.james.newagent
 
+import com.tomtom.james.common.api.informationpoint.InformationPoint
 import spock.lang.Specification
 
 import java.util.concurrent.TimeUnit
@@ -24,6 +25,8 @@ import static org.awaitility.Awaitility.await
 
 class MethodExecutionContextHelperSpec extends Specification {
 
+    static InformationPoint IP = InformationPoint.builder().withClassName("class").withMethodName("method").build()
+
     def "Should store context per thread"() {
         given:
         def keys = new String[2]
@@ -44,9 +47,9 @@ class MethodExecutionContextHelperSpec extends Specification {
     }
 
     static def runOnThread(String[] storage, int index) {
-        String key = MethodExecutionContextHelper.createContextKey()
+        String key = MethodExecutionContextHelper.createContextKey(IP)
         storage[index] = key
         MethodExecutionContextHelper.storeContextAsync(key, Integer.valueOf(100 + index))
-        MethodExecutionContextHelper.removeContextKey()
+        MethodExecutionContextHelper.removeContextKey(IP)
     }
 }

james-agent/src/test/groovy/com/tomtom/james/script/DisruptorAsyncScriptEngineSpec.groovy

@@ -124,7 +124,7 @@ class DisruptorAsyncScriptEngineSpec extends Specification {
 
         when:
         threadData.set("greetings from " + currentThread.getId())
-        def key = MethodExecutionContextHelper.createContextKey()
+        def key = MethodExecutionContextHelper.createContextKey(informationPoint)
         def context = scriptEngine.invokePrepareContext(informationPoint, origin, [param1, param2], instance, currentThread, key)
         MethodExecutionContextHelper.storeContextAsync(key, context)
         scriptEngine.invokeSuccessHandler(informationPoint, origin, [param1, param2], instance, currentThread, instant, duration, callStack, returnValue, MethodExecutionContextHelper.getContextAsync(key))