Resolve TODOs in executor instrumentations (#3760)

This commit is contained in:
Mateusz Rzeszutek 2021-08-04 12:49:53 +02:00 committed by GitHub
parent 1c58c132a3
commit 21b6492c17
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 14 additions and 25 deletions

View File

@ -184,17 +184,20 @@ public class JavaExecutorInstrumentation extends AbstractExecutorInstrumentation
Collection<Callable<?>> wrappedTasks = new ArrayList<>(tasks.size());
Context context = Java8BytecodeBridge.currentContext();
for (Callable<?> task : tasks) {
// TODO: missing shouldPropagateContext() check
if (task != null) {
if (ExecutorAdviceHelper.shouldPropagateContext(context, task)) {
Callable<?> newTask = CallableWrapper.wrapIfNeeded(task);
wrappedTasks.add(newTask);
ContextStore<Callable<?>, PropagatedContext> contextStore =
InstrumentationContext.get(Callable.class, PropagatedContext.class);
ExecutorAdviceHelper.attachContextToTask(context, contextStore, newTask);
} else {
// note that task may be null here
wrappedTasks.add(task);
}
}
tasks = wrappedTasks;
// TODO: why are we returning tasks and not just states?
// returning tasks and not propagatedContexts to avoid allocating another list just for an
// edge case (exception)
return tasks;
}
@ -211,23 +214,13 @@ public class JavaExecutorInstrumentation extends AbstractExecutorInstrumentation
any parent spans in case of an error.
(according to ExecutorService docs and AbstractExecutorService code)
*/
if (null != throwable) {
if (throwable != null) {
for (Callable<?> task : wrappedTasks) {
if (task != null) {
ContextStore<Callable<?>, PropagatedContext> contextStore =
InstrumentationContext.get(Callable.class, PropagatedContext.class);
PropagatedContext propagatedContext = contextStore.get(task);
if (propagatedContext != null) {
/*
Note: this may potentially clear somebody else's parent span if we didn't set it
up in setupState because it was already present before us. This should be safe but
may lead to non-attributed async work in some very rare cases.
Alternative is to not clear parent span here if we did not set it up in setupState
but this may potentially lead to memory leaks if callers do not properly handle
exceptions.
*/
propagatedContext.clear();
}
ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable);
}
}
}

View File

@ -53,14 +53,12 @@ public class GuavaListenableFutureInstrumentation implements TypeInstrumentation
@Advice.OnMethodEnter(suppress = Throwable.class)
public static PropagatedContext addListenerEnter(
@Advice.Argument(value = 0, readOnly = false) Runnable task) {
Runnable newTask = RunnableWrapper.wrapIfNeeded(task);
Context context = Java8BytecodeBridge.currentContext();
// TODO: shouldn't we check task? not newTask?
if (ExecutorAdviceHelper.shouldPropagateContext(context, newTask)) {
task = newTask;
if (ExecutorAdviceHelper.shouldPropagateContext(context, task)) {
task = RunnableWrapper.wrapIfNeeded(task);
ContextStore<Runnable, PropagatedContext> contextStore =
InstrumentationContext.get(Runnable.class, PropagatedContext.class);
return ExecutorAdviceHelper.attachContextToTask(context, contextStore, newTask);
return ExecutorAdviceHelper.attachContextToTask(context, contextStore, task);
}
return null;
}

View File

@ -42,14 +42,12 @@ public class JettyQueuedThreadPoolInstrumentation implements TypeInstrumentation
@Advice.OnMethodEnter(suppress = Throwable.class)
public static PropagatedContext enterJobSubmit(
@Advice.Argument(value = 0, readOnly = false) Runnable task) {
Runnable newTask = RunnableWrapper.wrapIfNeeded(task);
Context context = Java8BytecodeBridge.currentContext();
// TODO: shouldn't we check task? not newTask?
if (ExecutorAdviceHelper.shouldPropagateContext(context, newTask)) {
task = newTask;
if (ExecutorAdviceHelper.shouldPropagateContext(context, task)) {
task = RunnableWrapper.wrapIfNeeded(task);
ContextStore<Runnable, PropagatedContext> contextStore =
InstrumentationContext.get(Runnable.class, PropagatedContext.class);
return ExecutorAdviceHelper.attachContextToTask(context, contextStore, newTask);
return ExecutorAdviceHelper.attachContextToTask(context, contextStore, task);
}
return null;
}