Skip to content

Commit 05b9e03

Browse files
piotrwielgolaski-tomtompiotrwielgolaski-tomtom
authored
feat: simplify logic of handling initial context (#283)
- reduce thread locals - drop usage of extra executors
1 parent 227bec3 commit 05b9e03

File tree

10 files changed

+75
-258
lines changed

10 files changed

+75
-258
lines changed

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

+1-2
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,6 @@
2424
import com.tomtom.james.configuration.AgentConfigurationFactory;
2525
import com.tomtom.james.configuration.ConfigurationInitializationException;
2626
import com.tomtom.james.newagent.JVMAgentCleaner;
27-
import com.tomtom.james.newagent.MethodExecutionContextHelper;
2827
import com.tomtom.james.publisher.EventPublisherFactory;
2928
import com.tomtom.james.script.ScriptEngineFactory;
3029
import com.tomtom.james.store.informationpoints.io.InformationPointStore;
@@ -63,7 +62,7 @@ private static void setupAgent(Instrumentation instrumentation) {
6362
//InformationPointService informationPointService = new InformationPointServiceImpl(store, instrumentation);
6463
//controllersManager.initializeControllers(informationPointService, engine, publisher);
6564

66-
JVMAgentCleaner.init(controllersManager, engine, publisher, MethodExecutionContextHelper::shutdown);
65+
JVMAgentCleaner.init(controllersManager, engine, publisher);
6766
if (configuration.isShutdownHookEnabled()) {
6867
ShutdownHook shutdownHook = new ShutdownHook(configuration, JVMAgentCleaner::close);
6968
Runtime.getRuntime().addShutdownHook(shutdownHook);

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

+8-34
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,6 @@
2929
import com.tomtom.james.common.api.informationpoint.InformationPoint;
3030
import com.tomtom.james.common.api.script.RuntimeInformationPointParameter;
3131
import com.tomtom.james.common.log.Logger;
32-
import com.tomtom.james.newagent.MethodExecutionContextHelper;
3332
import org.apache.logging.log4j.util.StackLocatorUtil;
3433

3534
/*
@@ -44,7 +43,8 @@ private ContextAwareAdvice() {
4443
}
4544

4645
@SuppressWarnings("unused")
47-
public static void onEnter(String originTypeName,
46+
public static void onEnter(ExecutionContext context,
47+
String originTypeName,
4848
String originMethodName,
4949
Method origin,
5050
Object instance,
@@ -68,18 +68,15 @@ public static void onEnter(String originTypeName,
6868
LOG.trace(() -> "onEnter: noInitialContextSupportRequired - skipping");
6969
return;
7070
}
71-
72-
final String key = MethodExecutionContextHelper.createContextKey(ip);
7371
LOG.trace(() -> "Initializing custom context setup for the call");
7472
final Object callContext = ScriptEngineSupplier.get().invokePrepareContext(
7573
ip,
7674
origin,
7775
createParameterList(origin, arguments),
7876
instance,
7977
Thread.currentThread(),
80-
key);
81-
82-
MethodExecutionContextHelper.storeContextAsync(key, callContext);
78+
context.toString());
79+
context.setInitialContext(callContext);
8380

8481
} catch (Throwable t) {
8582
LOG.error("Error executing onEnter advice", t);
@@ -91,7 +88,7 @@ public static void onEnter(String originTypeName,
9188

9289
@SuppressWarnings("unused")
9390
public static void onExit(
94-
long _startTime,
91+
ExecutionContext context,
9592
String informationPointClassName,
9693
String informationPointMethodName,
9794
Method origin,
@@ -100,8 +97,7 @@ public static void onExit(
10097
Object returned,
10198
Throwable thrown) {
10299
Instant eventTime = Instant.now();
103-
Duration executionTime = Duration.ofNanos(System.nanoTime() - _startTime);
104-
AutoCloseable closeable = null;
100+
Duration executionTime = context.getElapsedTime();
105101

106102
try {
107103
Optional<InformationPoint> optionalInformationPoint = InformationPointServiceSupplier.get()
@@ -130,20 +126,11 @@ public static void onExit(
130126
+ ", thrown=" + thrown
131127
+ "]");
132128

133-
boolean requireInitialContextCleanup = ip.getRequiresInitialContext();
134-
if (requireInitialContextCleanup) {
135-
closeable = () -> MethodExecutionContextHelper.removeContextKey(ip);
136-
}
137-
138129
if ((sampleRate < 100) && (sampleRate < ThreadLocalRandom.current().nextDouble() * 100)) {
139130
LOG.trace(() -> "onExit: Sample skipped (sampleRate=" + sampleRate + ")");
140131
return;
141132
}
142133

143-
final CompletableFuture<Object> initialContextAsyncProvider = requireInitialContextCleanup
144-
? MethodExecutionContextHelper.getContextAsync(MethodExecutionContextHelper.getKeyForCurrentFrame(ip))
145-
: CompletableFuture.completedFuture(null);
146-
147134
final String[] callStack = ip.getRequiresCallStack() ? getCallStack() : EMPTY_CALL_STACK;
148135
if (thrown == null) {
149136
if (executionTime.toMillis() < successExecutionThreshold) {
@@ -161,7 +148,7 @@ public static void onExit(
161148
executionTime,
162149
callStack,
163150
returned,
164-
initialContextAsyncProvider
151+
CompletableFuture.completedFuture(context.getInitialContext())
165152
);
166153
} else {
167154
LOG.trace(() -> "onExit: Invoking error handler");
@@ -175,20 +162,13 @@ public static void onExit(
175162
executionTime,
176163
callStack,
177164
thrown,
178-
initialContextAsyncProvider
165+
CompletableFuture.completedFuture(context.getInitialContext())
179166
);
180167
}
181168
} catch (Throwable t) {
182169
LOG.error("Error executing onExit advice", t);
183170
throw t;
184171
} finally {
185-
try {
186-
if (closeable != null) {
187-
closeable.close();
188-
}
189-
} catch (Exception e) {
190-
LOG.error("Error executing onExit advice (finally)", e);
191-
}
192172
LOG.trace("onExit: END");
193173
}
194174
}
@@ -228,12 +208,6 @@ public static List<RuntimeInformationPointParameter> createParameterList(Method
228208
return result;
229209
}
230210

231-
@SuppressWarnings("unused")
232-
public static void onEnter(String originTypeName,
233-
String originMethodName) {
234-
onEnter(originTypeName, originMethodName, null, null, null);
235-
}
236-
237211

238212
@SuppressWarnings("unused")
239213
public static void onExit(

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

+24
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
package com.tomtom.james.informationpoint.advice;
2+
3+
import java.time.Duration;
4+
5+
public class ExecutionContext {
6+
private final long started;
7+
private Object initialContext;
8+
9+
public ExecutionContext() {
10+
this.started = System.nanoTime();
11+
}
12+
13+
public Duration getElapsedTime() {
14+
return Duration.ofNanos(System.nanoTime() - started);
15+
}
16+
17+
public Object getInitialContext() {
18+
return initialContext;
19+
}
20+
21+
public void setInitialContext(final Object initialContext) {
22+
this.initialContext = initialContext;
23+
}
24+
}

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -173,7 +173,7 @@ private static void setupAgent(Instrumentation inst) {
173173
removeInformationPointQueue
174174
);
175175
LOG.trace("initialize controllers time=" + stopwatch.elapsed());
176-
JVMAgentCleaner.init(controllersManager, engine, publisher, MethodExecutionContextHelper::shutdown);
176+
JVMAgentCleaner.init(controllersManager, engine, publisher);
177177
if (configuration.isShutdownHookEnabled()) {
178178
ShutdownHook shutdownHook = new ShutdownHook(configuration, JVMAgentCleaner::close);
179179
Runtime.getRuntime().addShutdownHook(shutdownHook);

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

+13-79
Original file line numberDiff line numberDiff line change
@@ -16,93 +16,27 @@
1616

1717
package com.tomtom.james.newagent;
1818

19-
import com.tomtom.james.common.api.informationpoint.InformationPoint;
20-
import com.tomtom.james.common.log.Logger;
21-
import com.tomtom.james.util.MoreExecutors;
22-
19+
import com.tomtom.james.informationpoint.advice.ExecutionContext;
2320
import java.util.ArrayDeque;
24-
import java.util.UUID;
25-
import java.util.WeakHashMap;
26-
import java.util.concurrent.CompletableFuture;
27-
import java.util.concurrent.ExecutorService;
28-
import java.util.concurrent.TimeUnit;
2921

22+
/**
23+
* Helper class for managing method execution context - used by GroovyJames to inject method entry and exit hooks.
24+
*/
3025
public class MethodExecutionContextHelper {
31-
private static final Logger LOG = Logger.getLogger(MethodExecutionContextHelper.class);
32-
33-
public static ThreadLocal<ArrayDeque<String>> keysStack = ThreadLocal.withInitial(() -> new ArrayDeque<>(8));
34-
private static WeakHashMap<String, Object> contextStore = new WeakHashMap<>();
3526

36-
private static ExecutorService contextStoreAccessExecutor = MoreExecutors
37-
.createNamedDaemonExecutorService("james-context-access-%d", 1);
38-
39-
private static ExecutorService contextCallbackExecutor = MoreExecutors
40-
.createNamedDaemonExecutorService("james-context-callback-%d", 5);
41-
42-
public static void shutdown() {
43-
try {
44-
contextStoreAccessExecutor.shutdown();
45-
contextStoreAccessExecutor.awaitTermination(5, TimeUnit.SECONDS);
46-
} catch (InterruptedException e) {
47-
LOG.warn("Executor contextStoreAccessExecutor shutdown interrupted " + e);
48-
}
49-
try {
50-
contextCallbackExecutor.shutdown();
51-
contextCallbackExecutor.awaitTermination(5, TimeUnit.SECONDS);
52-
} catch (InterruptedException e) {
53-
LOG.warn("Executor contextCallbackExecutor shutdown interrupted " + e);
54-
}
55-
}
27+
private static ThreadLocal<ArrayDeque<ExecutionContext>> contextStack = ThreadLocal.withInitial(() -> new ArrayDeque<>(8));
5628

57-
public static String createContextKey(final InformationPoint ip) {
58-
final String contextKey = ip.toString() + "-" + UUID.randomUUID();
59-
keysStack.get().push(contextKey);
60-
return contextKey;
29+
public static ExecutionContext executionStarted() {
30+
final ExecutionContext context = new ExecutionContext();
31+
contextStack.get().push(context);
32+
return context;
6133
}
6234

63-
public static String getKeyForCurrentFrame(final InformationPoint informationPoint) {
64-
final String key = keysStack.get().peek();
65-
if (key != null && key.startsWith(informationPoint.toString())) {
66-
return key;
67-
}
68-
return null;
35+
public static ExecutionContext getExecutionContext() {
36+
return contextStack.get().peek();
6937
}
7038

71-
public static String removeContextKey(final InformationPoint informationPoint) {
72-
String key = getKeyForCurrentFrame(informationPoint);
73-
if (key != null) {
74-
return keysStack.get().pop();
75-
}
76-
return null;
77-
}
78-
79-
public static CompletableFuture<Object> storeContextAsync(final String key, final Object value) {
80-
final CompletableFuture result = new CompletableFuture();
81-
contextStoreAccessExecutor.submit(() -> {
82-
contextStore.put(key, value);
83-
CompletableFuture.supplyAsync(() -> result.complete(value), contextCallbackExecutor);
84-
});
85-
return result;
86-
}
87-
88-
public static CompletableFuture<Object> getContextAsync(final String key) {
89-
if (key == null) {
90-
return CompletableFuture.completedFuture(null);
91-
}
92-
final CompletableFuture result = new CompletableFuture();
93-
contextStoreAccessExecutor.submit(() -> {
94-
if (contextStore.containsKey(key)) {
95-
final Object context = contextStore.get(key);
96-
CompletableFuture.supplyAsync(() -> result.complete(context), contextCallbackExecutor);
97-
return;
98-
}
99-
100-
CompletableFuture.supplyAsync(() -> {
101-
final String msg = String.format("Key '%s' not found in context container", key);
102-
result.completeExceptionally(new IllegalArgumentException(msg));
103-
return null;
104-
}, contextCallbackExecutor);
105-
});
106-
return result;
39+
public static ExecutionContext executionFinished() {
40+
return contextStack.get().pop();
10741
}
10842
}

‎james-agent/src/main/java/com/tomtom/james/newagent/MethodExecutionTimeHelper.java

-47
This file was deleted.

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

+4-15
Original file line numberDiff line numberDiff line change
@@ -17,18 +17,9 @@ public GroovyJames(Queue<JamesObjective> objectives, long sleepTime) {
1717
this.setName(getClass().getSimpleName());
1818
}
1919

20-
// TODO check double if that is all chars that we need to escape
21-
private String escapeScriptString(String script) {
22-
return script.replace("\\", "\\\\")
23-
.replace("\"", "\\\"")
24-
.replace("\r", "\\r")
25-
.replace("\n", "\\n");
26-
}
27-
2820
protected void insertBefore(CtMethod method, ExtendedInformationPoint informationPoint) throws CannotCompileException {
2921
StringBuilder s = new StringBuilder("");
30-
s.append(" com.tomtom.james.newagent.MethodExecutionTimeHelper.executionStarted();\n");
31-
s.append(" com.tomtom.james.informationpoint.advice.ContextAwareAdvice.onEnter(");
22+
s.append(" com.tomtom.james.informationpoint.advice.ContextAwareAdvice.onEnter( com.tomtom.james.newagent.MethodExecutionContextHelper.executionStarted(), \n");
3223
s.append("\"" + informationPoint.getClassName() + "\", ");
3324
s.append("\"" + informationPoint.getMethodName() + "\", ");
3425
s.append(informationPoint.getMethodBodyClassName() + ".class.getDeclaredMethod(\"" + informationPoint.getMethodName() + "\",$sig), "); // method
@@ -42,9 +33,8 @@ protected void insertBefore(CtMethod method, ExtendedInformationPoint informatio
4233
}
4334

4435
protected void insertAfter(CtMethod method, ExtendedInformationPoint informationPoint) throws CannotCompileException {
45-
String script = escapeScriptString(informationPoint.getScript().get());
4636
StringBuilder s = new StringBuilder();
47-
s.append(" com.tomtom.james.informationpoint.advice.ContextAwareAdvice.onExit( com.tomtom.james.newagent.MethodExecutionTimeHelper.getStartTime(), \n");
37+
s.append(" com.tomtom.james.informationpoint.advice.ContextAwareAdvice.onExit( com.tomtom.james.newagent.MethodExecutionContextHelper.getExecutionContext(), \n");
4838
s.append("\"" + informationPoint.getClassName() + "\", ");
4939
s.append("\"" + informationPoint.getMethodName() + "\", ");
5040
s.append(informationPoint.getMethodBodyClassName() + ".class.getDeclaredMethod(\"" + informationPoint.getMethodName() + "\",$sig), "); // method
@@ -64,9 +54,8 @@ protected void insertAfter(CtMethod method, ExtendedInformationPoint information
6454
}
6555

6656
protected void addCatch(ClassPool pool, CtMethod method, ExtendedInformationPoint informationPoint) throws CannotCompileException, NotFoundException {
67-
String script = escapeScriptString(informationPoint.getScript().get());
6857
StringBuilder s = new StringBuilder();
69-
s.append(" com.tomtom.james.informationpoint.advice.ContextAwareAdvice.onExit( com.tomtom.james.newagent.MethodExecutionTimeHelper.getStartTime(), ");
58+
s.append(" com.tomtom.james.informationpoint.advice.ContextAwareAdvice.onExit( com.tomtom.james.newagent.MethodExecutionContextHelper.getExecutionContext(), ");
7059
s.append("\"" + informationPoint.getClassName() + "\", ");
7160
s.append("\"" + informationPoint.getMethodName() + "\", ");
7261
s.append(informationPoint.getMethodBodyClassName() + ".class.getDeclaredMethod(\"" + informationPoint.getMethodName() + "\",$sig), "); // method
@@ -87,7 +76,7 @@ protected void addCatch(ClassPool pool, CtMethod method, ExtendedInformationPoin
8776

8877
// finally block
8978
StringBuilder f = new StringBuilder("");
90-
f.append(" com.tomtom.james.newagent.MethodExecutionTimeHelper.executionFinished(); \n");
79+
f.append(" com.tomtom.james.newagent.MethodExecutionContextHelper.executionFinished(); \n");
9180
method.insertAfter(f.toString(), true);
9281
}
9382
}

0 commit comments

Comments
 (0)