Skip to content

Commit

Permalink
feat: better handle case when existing IP starts to support initial c…
Browse files Browse the repository at this point in the history 
…ontext (#282)

it could be already running when new IP is loaded leading to situation when initial context is not loaded
  • Loading branch information
@piotrwielgolaski-tomtom
piotrwielgolaski-tomtom authored Nov 29, 2024
1 parent ff1deaf commit 227bec3
Show file tree
Hide file tree
Showing 4 changed files with 48 additions and 26 deletions.
41 changes: 24 additions & 17 deletions james-agent/src/main/java/com/tomtom/james/informationpoint/advice/ContextAwareAdvice.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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()
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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");
}
Expand Down
24 changes: 18 additions & 6 deletions james-agent/src/main/java/com/tomtom/james/newagent/MethodExecutionContextHelper.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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) {
Expand All @@ -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)) {
Expand Down
7 changes: 5 additions & 2 deletions ...s-agent/src/test/groovy/com/tomtom/james/newagent/MethodExecutionContextHelperSpec.groovy
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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]
Expand All @@ -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)
}
}
2 changes: 1 addition & 1 deletion james-agent/src/test/groovy/com/tomtom/james/script/DisruptorAsyncScriptEngineSpec.groovy
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down

0 comments on commit 227bec3

Please sign in to comment.