diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/AdviceUtils.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/AdviceUtils.java index 9fe033d7632..d536f859289 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/AdviceUtils.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/AdviceUtils.java @@ -41,6 +41,19 @@ public static ContextScope startTaskScope(State state) { return null; } + @Nullable + public static ContextScope startTpeTaskScope(State state) { + if (state != null) { + final State.TpeContinuation continuation = state.getAndResetTpeContinuation(); + if (continuation != null) { + final ContextScope scope = continuation.resume(); + continuation.stopTiming(); + return scope; + } + } + return null; + } + public static void endTaskScope(final ContextScope scope) { if (null != scope) { scope.close(); diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/State.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/State.java index 48cfe9078de..37c9b9e9d96 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/State.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/State.java @@ -40,6 +40,46 @@ public boolean captureAndSetContinuation(final Context context) { return false; } + /** + * Captures context for an unwrapped {@code ThreadPoolExecutor} submission. + * + *

The tagged continuation can only be consumed by {@code beforeExecute}; regular Runnable + * advice deliberately ignores it. This prevents a direct invocation, or another submission of the + * same Runnable, from stealing the queued submission's context. + */ + @Nullable + public TpeContinuation captureAndSetTpeContinuation(final Context context) { + while (true) { + ContextContinuation current = CONTINUATION.get(this); + if (current == null) { + if (!CONTINUATION.compareAndSet(this, null, CLAIMED)) { + continue; + } + try { + TpeContinuation continuation = new TpeContinuation(context.capture()); + CONTINUATION.lazySet(this, continuation); + return continuation; + } catch (Throwable error) { + CONTINUATION.compareAndSet(this, CLAIMED, null); + throw error; + } + } + if (current == CLAIMED || current instanceof TpeContinuation) { + return null; + } + // Generic Executor advice can run before ThreadPoolExecutor advice for the same call. + // Transfer + // that continuation instead of treating the duplicate instrumentation as task reuse. + if (current.context() != context) { + return null; + } + TpeContinuation continuation = new TpeContinuation(current); + if (CONTINUATION.compareAndSet(this, current, continuation)) { + return continuation; + } + } + } + public boolean setOrCancelContinuation(final ContextContinuation continuation) { if (CONTINUATION.compareAndSet(this, null, CLAIMED)) { // lazy write is guaranteed to be seen by getAndSet @@ -58,6 +98,28 @@ public void closeContinuation() { } } + @Nullable + public ContextContinuation getContinuation() { + ContextContinuation continuation = CONTINUATION.get(this); + return continuation == CLAIMED || continuation instanceof TpeContinuation ? null : continuation; + } + + @Nullable + public ContextContinuation getCancellableContinuation() { + ContextContinuation continuation = CONTINUATION.get(this); + return continuation == CLAIMED ? null : continuation; + } + + public void closeContinuation(ContextContinuation expected) { + if (expected != null && CONTINUATION.compareAndSet(this, expected, null)) { + if (expected instanceof TpeContinuation) { + ((TpeContinuation) expected).cancel(); + } else { + expected.release(); + } + } + } + public Context getContext() { ContextContinuation continuation = CONTINUATION.get(this); if (null == continuation || CLAIMED == continuation) { @@ -68,20 +130,54 @@ public Context getContext() { @Nullable public ContextContinuation getAndResetContinuation() { - ContextContinuation continuation = CONTINUATION.get(this); - if (null == continuation || CLAIMED == continuation) { - return null; + while (true) { + ContextContinuation continuation = CONTINUATION.get(this); + if (null == continuation + || CLAIMED == continuation + || continuation instanceof TpeContinuation) { + return null; + } + if (CONTINUATION.compareAndSet(this, continuation, null)) { + return continuation; + } + } + } + + @Nullable + public TpeContinuation getAndResetTpeContinuation() { + while (true) { + ContextContinuation continuation = CONTINUATION.get(this); + if (!(continuation instanceof TpeContinuation)) { + return null; + } + if (CONTINUATION.compareAndSet(this, continuation, null)) { + return (TpeContinuation) continuation; + } } - CONTINUATION.compareAndSet(this, continuation, null); - return continuation; + } + + @Nullable + public TpeContinuation getTpeContinuation() { + ContextContinuation continuation = CONTINUATION.get(this); + return continuation instanceof TpeContinuation ? (TpeContinuation) continuation : null; + } + + public void closeTpeContinuation(TpeContinuation expected) { + closeContinuation(expected); } public void setTiming(Timing timing) { - TIMING.lazySet(this, timing); + TpeContinuation continuation = getTpeContinuation(); + if (continuation != null) { + continuation.setTiming(timing); + } else { + TIMING.lazySet(this, timing); + } } public boolean isTimed() { - return TIMING.get(this) != null; + TpeContinuation continuation = getTpeContinuation(); + return TIMING.get(this) != null || (continuation != null && continuation.isTimed()); } public void stopTiming() { @@ -90,4 +186,55 @@ public void stopTiming() { QueueTimerHelper.stopQueuingTimer(timing); } } + + public static final class TpeContinuation implements ContextContinuation { + private final ContextContinuation delegate; + private volatile Timing timing; + + private TpeContinuation(ContextContinuation delegate) { + this.delegate = delegate; + } + + @Override + public ContextContinuation hold() { + delegate.hold(); + return this; + } + + @Override + public Context context() { + return delegate.context(); + } + + @Override + public datadog.context.ContextScope resume() { + return delegate.resume(); + } + + @Override + public void release() { + delegate.release(); + } + + private void setTiming(Timing timing) { + this.timing = timing; + } + + private boolean isTimed() { + return timing != null; + } + + public void stopTiming() { + Timing timing = this.timing; + this.timing = null; + if (timing != null) { + QueueTimerHelper.stopQueuingTimer(timing); + } + } + + private void cancel() { + release(); + stopTiming(); + } + } } diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/TPEHelper.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/TPEHelper.java index e95a6580849..b3354f98a6d 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/TPEHelper.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/TPEHelper.java @@ -1,14 +1,23 @@ package datadog.trace.bootstrap.instrumentation.java.concurrent; +import static datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils.shouldCapture; import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.ExcludeType.RUNNABLE; import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.exclude; +import datadog.context.Context; +import datadog.context.ContextContinuation; import datadog.context.ContextScope; import datadog.trace.api.GenericClassValue; import datadog.trace.api.InstrumenterConfig; import datadog.trace.api.Platform; import datadog.trace.bootstrap.ContextStore; import java.util.Set; +import java.util.concurrent.ArrayBlockingQueue; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.LinkedBlockingDeque; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.LinkedTransferQueue; +import java.util.concurrent.SynchronousQueue; import java.util.concurrent.ThreadPoolExecutor; /** @@ -63,17 +72,43 @@ public static boolean shouldPropagate(ThreadPoolExecutor executor) { && PROPAGATE.get(executor.getClass()); } - public static void capture(ContextStore contextStore, Runnable task) { + public static Runnable captureOrWrap( + ContextStore contextStore, + Runnable task, + Context context, + ThreadPoolExecutor executor) { + if (task != null && !exclude(RUNNABLE, task) && shouldCapture(context)) { + State state = contextStore.getOrCreate(task, State.FACTORY); + if (state.captureAndSetTpeContinuation(context) == null) { + return canWrapCollision(executor.getQueue()) ? Wrapper.wrap(task, context) : null; + } + } + return task; + } + + /** + * Retains the historical best-effort propagation for subclass overrides outside JDK admission. + */ + public static void captureLegacy(ContextStore contextStore, Runnable task) { if (task != null && !exclude(RUNNABLE, task)) { AdviceUtils.capture(contextStore, task); } } + private static boolean canWrapCollision(BlockingQueue queue) { + return queue instanceof ArrayBlockingQueue + || queue instanceof LinkedBlockingQueue + || queue instanceof LinkedBlockingDeque + || queue instanceof LinkedTransferQueue + || queue instanceof SynchronousQueue; + } + public static ContextScope startScope(ContextStore contextStore, Runnable task) { if (task == null || exclude(RUNNABLE, task)) { return null; } - return AdviceUtils.startTaskScope(contextStore, task); + State state = contextStore.get(task); + return AdviceUtils.startTpeTaskScope(state); } public static void setThreadLocalScope(ContextScope scope, Runnable task) { @@ -106,10 +141,53 @@ public static void endScope(ContextScope scope, Runnable task) { AdviceUtils.endTaskScope(scope); } - public static void cancelTask(ContextStore contextStore, Runnable task) { + public static final class RejectedTask { + private final Wrapper wrapper; + private final ContextContinuation continuation; + private final ContextScope scope; + + public RejectedTask(Wrapper wrapper, ContextContinuation continuation) { + this.wrapper = wrapper; + this.continuation = continuation; + this.scope = + wrapper != null + ? wrapper.activate() + : continuation == null ? null : continuation.resume(); + } + + public void close() { + try { + if (scope != null) { + scope.close(); + } + } finally { + if (wrapper != null) { + wrapper.cancel(); + } else if (continuation != null) { + continuation.release(); + } + } + } + } + + public static ContextContinuation prepareRejectedTask( + ContextStore contextStore, Runnable task) { if (task == null || exclude(RUNNABLE, task)) { - return; + return null; + } + State state = contextStore.get(task); + if (state == null) { + return null; + } + State.TpeContinuation tpeContinuation = state.getAndResetTpeContinuation(); + if (tpeContinuation != null) { + tpeContinuation.stopTiming(); + return tpeContinuation; + } + ContextContinuation continuation = state.getAndResetContinuation(); + if (continuation != null) { + state.stopTiming(); } - AdviceUtils.cancelTask(contextStore, task); + return continuation; } } diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/Wrapper.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/Wrapper.java index a26cc0c4535..22e1ab5cea4 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/Wrapper.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/Wrapper.java @@ -13,23 +13,41 @@ public class Wrapper implements Runnable, AutoCloseable { @SuppressWarnings({"unchecked", "rawtypes"}) public static Runnable wrap(T task) { - if (task instanceof Wrapper - || task instanceof RunnableFuture - || task == null - || exclude(RUNNABLE, task)) { + if (!isWrappable(task)) { return task; } ContextContinuation continuation = captureActiveSpan(); if (continuation.context() != Context.root()) { - if (task instanceof Comparable) { - return new ComparableRunnable(task, continuation); - } - return new Wrapper<>(task, continuation); + return newWrapper(task, continuation); } // don't wrap unless there is scope to propagate return task; } + @SuppressWarnings({"unchecked", "rawtypes"}) + public static Runnable wrap(T task, Context context) { + if (context == Context.root() || !isWrappable(task)) { + return task; + } + return newWrapper(task, context.capture()); + } + + private static boolean isWrappable(Runnable task) { + return task != null + && !(task instanceof Wrapper) + && !(task instanceof RunnableFuture) + && !exclude(RUNNABLE, task); + } + + @SuppressWarnings({"unchecked", "rawtypes"}) + private static Runnable newWrapper( + T task, ContextContinuation continuation) { + if (task instanceof Comparable) { + return new ComparableRunnable(task, continuation); + } + return new Wrapper<>(task, continuation); + } + public static Runnable unwrap(Runnable task) { return task instanceof Wrapper ? ((Wrapper) task).unwrap() : task; } @@ -59,7 +77,7 @@ public T unwrap() { return delegate; } - private ContextScope activate() { + public ContextScope activate() { return null == continuation ? null : continuation.resume(); } diff --git a/dd-java-agent/agent-bootstrap/src/test/java/datadog/trace/bootstrap/instrumentation/java/concurrent/StateTest.java b/dd-java-agent/agent-bootstrap/src/test/java/datadog/trace/bootstrap/instrumentation/java/concurrent/StateTest.java new file mode 100644 index 00000000000..a6a621d4dd0 --- /dev/null +++ b/dd-java-agent/agent-bootstrap/src/test/java/datadog/trace/bootstrap/instrumentation/java/concurrent/StateTest.java @@ -0,0 +1,110 @@ +package datadog.trace.bootstrap.instrumentation.java.concurrent; + +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import datadog.context.Context; +import datadog.context.ContextContinuation; +import datadog.trace.api.profiling.Timing; +import org.junit.jupiter.api.Test; + +class StateTest { + + @Test + void regularAdviceCannotConsumeTpeContinuation() { + State state = State.FACTORY.create(); + Context context = mock(Context.class); + ContextContinuation delegate = mock(ContextContinuation.class); + when(context.capture()).thenReturn(delegate); + + State.TpeContinuation continuation = state.captureAndSetTpeContinuation(context); + + assertNull(state.getAndResetContinuation()); + assertSame(continuation, state.getAndResetTpeContinuation()); + } + + @Test + void overlappingTpeCaptureDoesNotReplaceOwner() { + State state = State.FACTORY.create(); + Context first = mock(Context.class); + Context second = mock(Context.class); + ContextContinuation firstDelegate = mock(ContextContinuation.class); + when(first.capture()).thenReturn(firstDelegate); + + State.TpeContinuation continuation = state.captureAndSetTpeContinuation(first); + + assertNull(state.captureAndSetTpeContinuation(second)); + assertSame(continuation, state.getAndResetTpeContinuation()); + verify(second, never()).capture(); + } + + @Test + void transfersDuplicateGenericCaptureForSameSubmission() { + State state = State.FACTORY.create(); + Context context = mock(Context.class); + ContextContinuation delegate = mock(ContextContinuation.class); + when(context.capture()).thenReturn(delegate); + when(delegate.context()).thenReturn(context); + state.captureAndSetContinuation(context); + + State.TpeContinuation continuation = state.captureAndSetTpeContinuation(context); + + assertSame(continuation, state.getAndResetTpeContinuation()); + verify(context).capture(); + } + + @Test + void staleCancellationCannotCancelLaterSubmission() { + State state = State.FACTORY.create(); + Context first = mock(Context.class); + Context second = mock(Context.class); + ContextContinuation firstDelegate = mock(ContextContinuation.class); + ContextContinuation secondDelegate = mock(ContextContinuation.class); + when(first.capture()).thenReturn(firstDelegate); + when(second.capture()).thenReturn(secondDelegate); + + State.TpeContinuation stale = state.captureAndSetTpeContinuation(first); + assertSame(stale, state.getAndResetTpeContinuation()); + State.TpeContinuation current = state.captureAndSetTpeContinuation(second); + + state.closeTpeContinuation(stale); + + assertSame(current, state.getAndResetTpeContinuation()); + verify(firstDelegate, never()).release(); + verify(secondDelegate, never()).release(); + } + + @Test + void exactCancellationReleasesOwnedContinuation() { + State state = State.FACTORY.create(); + Context context = mock(Context.class); + ContextContinuation delegate = mock(ContextContinuation.class); + when(context.capture()).thenReturn(delegate); + State.TpeContinuation continuation = state.captureAndSetTpeContinuation(context); + + state.closeTpeContinuation(continuation); + + assertNull(state.getAndResetTpeContinuation()); + verify(delegate).release(); + } + + @Test + void taggedContinuationOwnsQueueTiming() { + State state = State.FACTORY.create(); + Context context = mock(Context.class); + ContextContinuation delegate = mock(ContextContinuation.class); + Timing timing = mock(Timing.class); + when(context.capture()).thenReturn(delegate); + State.TpeContinuation continuation = state.captureAndSetTpeContinuation(context); + + state.setTiming(timing); + + assertTrue(state.isTimed()); + assertSame(continuation, state.getAndResetTpeContinuation()); + } +} diff --git a/dd-java-agent/agent-bootstrap/src/test/java/datadog/trace/bootstrap/instrumentation/java/concurrent/TPEHelperTest.java b/dd-java-agent/agent-bootstrap/src/test/java/datadog/trace/bootstrap/instrumentation/java/concurrent/TPEHelperTest.java new file mode 100644 index 00000000000..d2bd193aa03 --- /dev/null +++ b/dd-java-agent/agent-bootstrap/src/test/java/datadog/trace/bootstrap/instrumentation/java/concurrent/TPEHelperTest.java @@ -0,0 +1,28 @@ +package datadog.trace.bootstrap.instrumentation.java.concurrent; + +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoInteractions; + +import datadog.context.Context; +import datadog.trace.bootstrap.ContextStore; +import java.util.concurrent.ThreadPoolExecutor; +import org.junit.jupiter.api.Test; + +class TPEHelperTest { + + @Test + void doesNotInspectQueueWhenThereIsNoContextToCapture() { + @SuppressWarnings("unchecked") + ContextStore contextStore = mock(ContextStore.class); + ThreadPoolExecutor executor = mock(ThreadPoolExecutor.class); + Runnable task = mock(Runnable.class); + + assertSame(task, TPEHelper.captureOrWrap(contextStore, task, Context.root(), executor)); + + verify(executor, never()).getQueue(); + verifyNoInteractions(contextStore); + } +} diff --git a/dd-java-agent/benchmark/build.gradle b/dd-java-agent/benchmark/build.gradle index 178eefd2def..d12ecd99338 100644 --- a/dd-java-agent/benchmark/build.gradle +++ b/dd-java-agent/benchmark/build.gradle @@ -5,6 +5,12 @@ plugins { apply from: "$rootDir/gradle/java.gradle" dependencies { + jmh(project(':internal-api')) { + transitive = false + } + jmh(project(':components:context')) { + transitive = false + } jmh project(':dd-trace-api') jmh libs.bytebuddyagent } @@ -38,8 +44,15 @@ jmh { jmhVersion = libs.versions.jmh.get() } +// Copy the agent to a fixed path because @Fork arguments must be compile-time constants. +def agentJarForBenchmarks = tasks.register('agentJarForBenchmarks', Copy) { + from project(':dd-java-agent').tasks.named('shadowJar') + into layout.buildDirectory.dir('agent') + rename { 'dd-java-agent.jar' } +} + tasks.named('jmh') { - dependsOn ':dd-java-agent:shadowJar' + dependsOn agentJarForBenchmarks } /* @@ -48,4 +61,3 @@ tasks.named('jmh') { (using https://github.com/brendangregg/FlameGraph) ./flamegraph.pl --color=java dd-java-agent/benchmark/build/reports/jmh/profiler-cleaned.txt > dd-java-agent/benchmark/build/reports/jmh/jmh-master.svg */ - diff --git a/dd-java-agent/benchmark/src/jmh/java/executorbench/RunnableSubmissionBenchmark.java b/dd-java-agent/benchmark/src/jmh/java/executorbench/RunnableSubmissionBenchmark.java new file mode 100644 index 00000000000..1502d945e3a --- /dev/null +++ b/dd-java-agent/benchmark/src/jmh/java/executorbench/RunnableSubmissionBenchmark.java @@ -0,0 +1,326 @@ +package executorbench; + +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; + +import datadog.trace.bootstrap.instrumentation.api.AgentScope; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import java.util.concurrent.ArrayBlockingQueue; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.Phaser; +import java.util.concurrent.SynchronousQueue; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Fork; +import org.openjdk.jmh.annotations.Level; +import org.openjdk.jmh.annotations.Measurement; +import org.openjdk.jmh.annotations.Mode; +import org.openjdk.jmh.annotations.OutputTimeUnit; +import org.openjdk.jmh.annotations.Param; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.Setup; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.annotations.TearDown; +import org.openjdk.jmh.annotations.Warmup; + +/** + * Measures context propagation through standard and subclassed {@link ThreadPoolExecutor}s. + * + *

The reusable, fresh, and root-context cases include a worker handoff. The overlapping case + * blocks the worker outside the measured interval so both submissions contend for the same task's + * ownership slot; its score covers two submissions. Allocation can be inspected with {@code -prof + * gc} for all cases except the overlapping benchmark, whose invocation-level fixture allocation is + * included in the profiler result. + * + *

Run from {@code dd-java-agent/benchmark} so the relative agent path resolves: + * + * 

{@code
+ * java -jar build/libs/benchmark-*-jmh.jar 'RunnableSubmissionBenchmark.*' -prof gc
+ * }
+ * + *

To compare commits with the same harness, replace {@code build/agent/dd-java-agent.jar} with + * each commit's agent before a run. Do not rebuild the JMH jar between arms. + */ +@BenchmarkMode(Mode.AverageTime) +@OutputTimeUnit(TimeUnit.MICROSECONDS) +@Warmup(iterations = 3, time = 1) +@Measurement(iterations = 5, time = 1) +@Fork( + value = 3, + jvmArgsAppend = { + "-javaagent:build/agent/dd-java-agent.jar", + "-Ddd.profiling.enabled=false", + "-Ddd.profiling.queueing.time.enabled=false", + "-Ddd.instrumentation.telemetry.enabled=false", + "-Ddd.remote_config.enabled=false", + "-Ddd.jmxfetch.enabled=false" + }) +public class RunnableSubmissionBenchmark { + + @Benchmark + public void activeReusableTask(ExecutorState state) throws InterruptedException { + state.submitActive(state.reusableTask); + } + + @Benchmark + public void activeFreshTask(ExecutorState state) throws InterruptedException { + state.submitActive(new CompletionTask(state, state.completion)); + } + + @Benchmark + public void rootReusableTask(ExecutorState state) throws InterruptedException { + state.submitRoot(state.reusableTask); + } + + @Benchmark + public void activeLambda(ExecutorState state) throws InterruptedException { + state.submitActiveLambda(); + } + + @Benchmark + public void overlappingReusableTask(CollisionState state) { + state.submitOverlapping(); + } + + @State(Scope.Thread) + public static class ExecutorState { + @Param({"base", "delegating", "nondelegating"}) + public String executorType; + + private ThreadPoolExecutor executor; + private Phaser completion; + private CompletionTask reusableTask; + private AgentSpan span; + private volatile boolean propagationFailed; + + @Setup(Level.Trial) + public void setup() { + executor = newExecutor(executorType); + executor.prestartAllCoreThreads(); + completion = new Phaser(1); + reusableTask = new CompletionTask(this, completion); + span = startSpan("benchmark", "runnable-submission"); + } + + @TearDown(Level.Trial) + public void tearDown() throws InterruptedException { + closeExecutor(executor); + span.finish(); + if (propagationFailed) { + throw new AssertionError("task did not receive the submitted span"); + } + } + + private void submitActive(CompletionTask task) throws InterruptedException { + int phase = completion.getPhase(); + task.expected = span; + try (AgentScope ignored = activateSpan(span)) { + executor.execute(task); + } + completion.awaitAdvanceInterruptibly(phase); + } + + private void submitRoot(CompletionTask task) throws InterruptedException { + int phase = completion.getPhase(); + task.expected = null; + executor.execute(task); + completion.awaitAdvanceInterruptibly(phase); + } + + private void submitActiveLambda() throws InterruptedException { + int phase = completion.getPhase(); + try (AgentScope ignored = activateSpan(span)) { + executor.execute( + () -> { + if (activeSpan() != span) { + propagationFailed = true; + } + completion.arrive(); + }); + } + completion.awaitAdvanceInterruptibly(phase); + } + } + + @State(Scope.Thread) + public static class CollisionState { + @Param({"base", "delegating"}) + public String executorType; + + private ThreadPoolExecutor executor; + private BlockingTask blocker; + private CollisionTask task; + private AgentSpan firstSpan; + private AgentSpan secondSpan; + + @Setup(Level.Invocation) + public void setup() throws InterruptedException { + executor = newExecutor(executorType); + executor.prestartAllCoreThreads(); + blocker = new BlockingTask(); + executor.execute(blocker); + blocker.started.await(); + firstSpan = startSpan("benchmark", "first-submission"); + secondSpan = startSpan("benchmark", "second-submission"); + task = new CollisionTask(firstSpan, secondSpan); + } + + @TearDown(Level.Invocation) + public void tearDown() throws InterruptedException { + blocker.release.countDown(); + task.finished.await(); + closeExecutor(executor); + firstSpan.finish(); + secondSpan.finish(); + if (task.propagationFailed) { + throw new AssertionError("overlapping submissions did not keep their own spans"); + } + } + + private void submitOverlapping() { + try (AgentScope ignored = activateSpan(firstSpan)) { + executor.execute(task); + } + try (AgentScope ignored = activateSpan(secondSpan)) { + executor.execute(task); + } + } + } + + private static ThreadPoolExecutor newExecutor(String executorType) { + switch (executorType) { + case "base": + return new ThreadPoolExecutor( + 1, 1, 0, TimeUnit.MILLISECONDS, new ArrayBlockingQueue(16)); + case "delegating": + return new DelegatingExecutor(); + case "nondelegating": + return new NonDelegatingExecutor(); + default: + throw new IllegalArgumentException("Unknown executor type: " + executorType); + } + } + + private static void closeExecutor(ThreadPoolExecutor executor) throws InterruptedException { + if (executor instanceof NonDelegatingExecutor) { + ((NonDelegatingExecutor) executor).closeWorker(); + } + executor.shutdownNow(); + if (!executor.awaitTermination(10, TimeUnit.SECONDS)) { + throw new AssertionError("executor did not terminate"); + } + } + + private static final class CompletionTask implements Runnable { + private final ExecutorState owner; + private final Phaser completion; + private volatile AgentSpan expected; + + private CompletionTask(ExecutorState owner, Phaser completion) { + this.owner = owner; + this.completion = completion; + } + + @Override + public void run() { + if (activeSpan() != expected) { + owner.propagationFailed = true; + } + completion.arrive(); + } + } + + private static final class CollisionTask implements Runnable { + private final AgentSpan firstExpected; + private final AgentSpan secondExpected; + private final CountDownLatch finished = new CountDownLatch(2); + private int execution; + private volatile boolean propagationFailed; + + private CollisionTask(AgentSpan firstExpected, AgentSpan secondExpected) { + this.firstExpected = firstExpected; + this.secondExpected = secondExpected; + } + + @Override + public void run() { + AgentSpan expected = execution++ == 0 ? firstExpected : secondExpected; + if (activeSpan() != expected) { + propagationFailed = true; + } + finished.countDown(); + } + } + + private static final class BlockingTask implements Runnable { + private final CountDownLatch started = new CountDownLatch(1); + private final CountDownLatch release = new CountDownLatch(1); + + @Override + public void run() { + started.countDown(); + try { + release.await(); + } catch (InterruptedException interrupted) { + Thread.currentThread().interrupt(); + } + } + } + + private static final class DelegatingExecutor extends ThreadPoolExecutor { + private DelegatingExecutor() { + super(1, 1, 0, TimeUnit.MILLISECONDS, new ArrayBlockingQueue(16)); + } + + @Override + public void execute(Runnable task) { + super.execute(task); + } + } + + private static final class NonDelegatingExecutor extends ThreadPoolExecutor { + private static final Runnable STOP = () -> {}; + + private final BlockingQueue tasks = new LinkedBlockingQueue<>(); + private final Thread worker; + + private NonDelegatingExecutor() { + super(0, 1, 0, TimeUnit.MILLISECONDS, new SynchronousQueue()); + worker = new Thread(this::runTasks, "nondelegating-benchmark-worker"); + worker.setDaemon(true); + worker.start(); + } + + @Override + public void execute(Runnable task) { + tasks.add(task); + } + + private void closeWorker() throws InterruptedException { + tasks.add(STOP); + worker.join(TimeUnit.SECONDS.toMillis(10)); + if (worker.isAlive()) { + throw new AssertionError("non-delegating worker did not terminate"); + } + } + + private void runTasks() { + try { + while (true) { + Runnable task = tasks.take(); + if (task == STOP) { + return; + } + task.run(); + } + } catch (InterruptedException interrupted) { + Thread.currentThread().interrupt(); + } + } + } +} diff --git a/dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-1.8/src/main/java/datadog/trace/instrumentation/java/concurrent/executor/RejectedExecutionHandlerInstrumentation.java b/dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-1.8/src/main/java/datadog/trace/instrumentation/java/concurrent/executor/RejectedExecutionHandlerInstrumentation.java index e8b7b9effd3..b7a5f50cb6f 100644 --- a/dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-1.8/src/main/java/datadog/trace/instrumentation/java/concurrent/executor/RejectedExecutionHandlerInstrumentation.java +++ b/dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-1.8/src/main/java/datadog/trace/instrumentation/java/concurrent/executor/RejectedExecutionHandlerInstrumentation.java @@ -4,17 +4,18 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.nameEndsWith; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.namedOneOf; -import static datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils.cancelTask; import static datadog.trace.instrumentation.java.concurrent.ConcurrentInstrumentationNames.EXECUTOR_INSTRUMENTATION_NAME; import static java.util.Arrays.asList; import static net.bytebuddy.matcher.ElementMatchers.isMethod; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; +import datadog.context.ContextContinuation; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.api.Config; import datadog.trace.api.InstrumenterConfig; import datadog.trace.bootstrap.InstrumentationContext; import datadog.trace.bootstrap.instrumentation.java.concurrent.State; +import datadog.trace.bootstrap.instrumentation.java.concurrent.TPEHelper; import datadog.trace.bootstrap.instrumentation.java.concurrent.Wrapper; import datadog.trace.bootstrap.instrumentation.jfr.backpressure.BackpressureProfiling; import java.util.concurrent.RunnableFuture; @@ -69,12 +70,17 @@ public void methodAdvice(MethodTransformer transformer) { public static final class Reject { // remove our wrapper before calling the handler (save wrapper, so we can cancel it later) @Advice.OnMethodEnter(suppress = Throwable.class) - public static Wrapper handle( + public static TPEHelper.RejectedTask handle( @Advice.This Object zis, @Advice.Argument(readOnly = false, value = 0) Runnable runnable) { Wrapper wrapper = null; + ContextContinuation continuation = null; if (runnable instanceof Wrapper) { wrapper = (Wrapper) runnable; runnable = wrapper.unwrap(); + } else { + continuation = + TPEHelper.prepareRejectedTask( + InstrumentationContext.get(Runnable.class, State.class), runnable); } if (Config.get().isProfilingBackPressureSamplingEnabled()) { // record this event before the handler executes, which will help @@ -82,28 +88,29 @@ public static Wrapper handle( // rejection policies which run on the caller (CallerRunsPolicy or user-provided) BackpressureProfiling.getInstance().process(zis.getClass(), runnable); } - return wrapper; + return wrapper == null && continuation == null + ? null + : new TPEHelper.RejectedTask(wrapper, continuation); } // must execute after in case the handler actually runs the runnable, // which is preferable to cancelling the continuation @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void reject( - @Advice.Enter Wrapper wrapper, @Advice.Argument(value = 0) Runnable runnable) { + @Advice.Enter TPEHelper.RejectedTask rejected, + @Advice.Argument(value = 0) Runnable runnable) { // not handling rejected work (which will often not manifest in an exception being thrown) // leads to unclosed continuations when executors get busy // note that this does not handle rejection mechanisms used in Scala, so those need to be // handled another way - if (null != wrapper) { - wrapper.cancel(); + if (null != rejected) { + rejected.close(); } else { if (runnable instanceof RunnableFuture) { - cancelTask( + datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils.cancelTask( InstrumentationContext.get(RunnableFuture.class, State.class), (RunnableFuture) runnable); } - // paranoid about double instrumentation until RunnableInstrumentation is removed - cancelTask(InstrumentationContext.get(Runnable.class, State.class), runnable); } } } diff --git a/dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-1.8/src/main/java/datadog/trace/instrumentation/java/concurrent/executor/ThreadPoolExecutorInstrumentation.java b/dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-1.8/src/main/java/datadog/trace/instrumentation/java/concurrent/executor/ThreadPoolExecutorInstrumentation.java index b7340c0f434..00ec82afaa2 100644 --- a/dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-1.8/src/main/java/datadog/trace/instrumentation/java/concurrent/executor/ThreadPoolExecutorInstrumentation.java +++ b/dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-1.8/src/main/java/datadog/trace/instrumentation/java/concurrent/executor/ThreadPoolExecutorInstrumentation.java @@ -3,6 +3,7 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.extendsClass; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.namedOneOf; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.currentContext; import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.ExcludeType.EXECUTOR; import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.ExcludeType.RUNNABLE; import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.ExcludeType.RUNNABLE_FUTURE; @@ -12,14 +13,19 @@ import static net.bytebuddy.matcher.ElementMatchers.not; import static net.bytebuddy.matcher.ElementMatchers.returns; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; +import static net.bytebuddy.matcher.ElementMatchers.takesArguments; +import datadog.context.ContextContinuation; import datadog.context.ContextScope; import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.bootstrap.ContextStore; import datadog.trace.bootstrap.InstrumentationContext; import datadog.trace.bootstrap.instrumentation.java.concurrent.QueueTimerHelper; import datadog.trace.bootstrap.instrumentation.java.concurrent.State; import datadog.trace.bootstrap.instrumentation.java.concurrent.TPEHelper; import datadog.trace.bootstrap.instrumentation.java.concurrent.Wrapper; +import java.util.List; +import java.util.ListIterator; import java.util.Queue; import java.util.concurrent.RunnableFuture; import java.util.concurrent.ThreadPoolExecutor; @@ -56,13 +62,10 @@ public final class ThreadPoolExecutorInstrumentation Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice { - // executors which do their own wrapping before calling super, - // leading to double wrapping, once at the child level and once - // in ThreadPoolExecutor - private static final ElementMatcher NO_WRAPPING_BEFORE_DELEGATION = - not( - isDeclaredBy( - namedOneOf("org.elasticsearch.common.util.concurrent.EsThreadPoolExecutor"))); + // This executor decorates tasks before delegating, so its override must not use the legacy + // compatibility path. Exact ownership is captured later by ThreadPoolExecutor.execute. + private static final ElementMatcher DECORATES_BEFORE_DELEGATION = + isDeclaredBy(namedOneOf("org.elasticsearch.common.util.concurrent.EsThreadPoolExecutor")); @Override public String hierarchyMarkerType() { @@ -80,9 +83,19 @@ public void methodAdvice(MethodTransformer transformer) { transformer.applyAdvice( named("execute") .and(isMethod()) - .and(NO_WRAPPING_BEFORE_DELEGATION) + .and(isDeclaredBy(named(ThreadPoolExecutor.class.getName()))) .and(takesArgument(0, named(Runnable.class.getName()))), getClass().getName() + "$Execute"); + transformer.applyAdvice( + // Preserve behavior for overrides that may bypass JDK admission. This path is + // intentionally best-effort; exact submission ownership begins in ThreadPoolExecutor. + named("execute") + .and(isMethod()) + .and(not(isDeclaredBy(named(ThreadPoolExecutor.class.getName())))) + .and(not(DECORATES_BEFORE_DELEGATION)) + .and(takesArgument(0, named(Runnable.class.getName()))) + .and(takesArguments(1)), + getClass().getName() + "$ExecuteOverride"); transformer.applyAdvice( named("beforeExecute") .and(isMethod()) @@ -94,8 +107,19 @@ public void methodAdvice(MethodTransformer transformer) { .and(takesArgument(0, named(Runnable.class.getName()))), getClass().getName() + "$AfterExecute"); transformer.applyAdvice( - named("remove").and(isMethod()).and(returns(Runnable.class)), + named("remove") + .and(isMethod()) + .and(isDeclaredBy(named(ThreadPoolExecutor.class.getName()))) + .and(takesArgument(0, named(Runnable.class.getName()))) + .and(returns(boolean.class)), getClass().getName() + "$Remove"); + transformer.applyAdvice( + named("shutdownNow") + .and(isMethod()) + .and(isDeclaredBy(named(ThreadPoolExecutor.class.getName()))) + .and(takesArguments(0)) + .and(returns(List.class)), + getClass().getName() + "$ShutdownNow"); } public static final class Execute { @@ -107,12 +131,21 @@ public static void capture( if (TPEHelper.useWrapping(task)) { task = Wrapper.wrap(task); } else { - TPEHelper.capture(InstrumentationContext.get(Runnable.class, State.class), task); - // queue time needs to be handled separately because there are RunnableFutures which are - // excluded as - // Runnables but it is not until now that they will be put on the executor's queue + Runnable captured = + TPEHelper.captureOrWrap( + InstrumentationContext.get(Runnable.class, State.class), + task, + currentContext(), + tpe); + if (captured == null) { + return; + } + task = captured; + // queue time needs to be handled separately because there are RunnableFutures which + // are excluded as Runnables but it is not until now that they will be put on the + // executor's queue if (!exclude(EXECUTOR, tpe)) { - if (!exclude(RUNNABLE, task)) { + if (!(task instanceof Wrapper) && !exclude(RUNNABLE, task)) { Queue queue = tpe.getQueue(); QueueTimerHelper.startQueuingTimer( InstrumentationContext.get(Runnable.class, State.class), @@ -135,13 +168,32 @@ public static void capture( } } + public static final class ExecuteOverride { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void capture( + @Advice.This final ThreadPoolExecutor tpe, + @Advice.Argument(readOnly = false, value = 0) Runnable task) { + if (TPEHelper.shouldPropagate(tpe)) { + if (TPEHelper.useWrapping(task)) { + task = Wrapper.wrap(task); + } else { + TPEHelper.captureLegacy(InstrumentationContext.get(Runnable.class, State.class), task); + } + } + } + } + public static final class BeforeExecute { @Advice.OnMethodEnter(suppress = Throwable.class) public static ContextScope beforeExecuteEnter( @Advice.This final ThreadPoolExecutor tpe, - @Advice.Argument(readOnly = false, value = 1) Runnable task) { + @Advice.Argument(readOnly = false, value = 1) Runnable task, + @Advice.Local("wrapper") Wrapper wrapper) { if (TPEHelper.shouldPropagate(tpe)) { if (TPEHelper.useWrapping(task)) { + if (task instanceof Wrapper) { + wrapper = (Wrapper) task; + } task = Wrapper.unwrap(task); } else { return TPEHelper.startScope( @@ -153,8 +205,16 @@ public static ContextScope beforeExecuteEnter( @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void beforeExecuteExit( - @Advice.Enter final ContextScope scope, @Advice.Argument(value = 1) Runnable task) { - if (scope != null) { + @Advice.Enter final ContextScope scope, + @Advice.Argument(value = 1) Runnable task, + @Advice.Thrown Throwable throwable, + @Advice.Local("wrapper") Wrapper wrapper) { + if (throwable != null) { + TPEHelper.endScope(scope, task); + if (wrapper != null) { + wrapper.cancel(); + } + } else if (scope != null) { TPEHelper.setThreadLocalScope(scope, task); } } @@ -185,19 +245,90 @@ public static void afterExecuteExit( } public static final class Remove { + @Advice.OnMethodEnter(skipOn = Advice.OnNonDefaultValue.class, suppress = Throwable.class) + public static boolean enter( + @Advice.This final ThreadPoolExecutor tpe, + @Advice.Argument(0) Runnable task, + @Advice.Local("owner") Runnable owner, + @Advice.Local("continuation") ContextContinuation continuation) { + if (!TPEHelper.shouldPropagate(tpe)) { + return false; + } + if (task instanceof Wrapper) { + owner = task; + return false; + } + if (task == null) { + return false; + } + ContextStore contextStore = + InstrumentationContext.get(Runnable.class, State.class); + for (Runnable queued : tpe.getQueue()) { + Runnable logical = queued instanceof Wrapper ? ((Wrapper) queued).unwrap() : queued; + if (task == logical || task.equals(logical)) { + if (queued instanceof Wrapper) { + if (tpe.getQueue().remove(queued)) { + ((Wrapper) queued).cancel(); + return true; + } + continue; + } + owner = queued; + State state = contextStore.get(queued); + continuation = state == null ? null : state.getCancellableContinuation(); + break; + } + } + return false; + } + @Advice.OnMethodExit(suppress = Throwable.class) public static void remove( @Advice.This final ThreadPoolExecutor tpe, - @Advice.Return(readOnly = false) Runnable removed) { - if (TPEHelper.shouldPropagate(tpe)) { - if (TPEHelper.useWrapping(removed)) { - if (removed instanceof Wrapper) { - Wrapper wrapper = ((Wrapper) removed); + @Advice.Enter boolean handled, + @Advice.Local("owner") Runnable owner, + @Advice.Local("continuation") ContextContinuation continuation, + @Advice.Return(readOnly = false) boolean removed) { + if (handled) { + removed = true; + return; + } + if (!TPEHelper.shouldPropagate(tpe)) { + return; + } + if (removed) { + if (owner instanceof Wrapper) { + Wrapper wrapper = ((Wrapper) owner); + wrapper.cancel(); + } else if (owner != null) { + State state = InstrumentationContext.get(Runnable.class, State.class).get(owner); + if (state != null) { + state.closeContinuation(continuation); + } + } + } + } + } + + public static final class ShutdownNow { + @Advice.OnMethodExit(suppress = Throwable.class) + public static void shutdown( + @Advice.This final ThreadPoolExecutor tpe, @Advice.Return List tasks) { + if (tasks != null && TPEHelper.shouldPropagate(tpe)) { + ContextStore contextStore = + InstrumentationContext.get(Runnable.class, State.class); + for (ListIterator iterator = tasks.listIterator(); iterator.hasNext(); ) { + Runnable task = iterator.next(); + if (task instanceof Wrapper) { + Wrapper wrapper = (Wrapper) task; wrapper.cancel(); - removed = wrapper.unwrap(); + iterator.set(wrapper.unwrap()); + } else { + State state = contextStore.get(task); + if (state != null) { + state.closeContinuation(state.getCancellableContinuation()); + } } - } else { - TPEHelper.cancelTask(InstrumentationContext.get(Runnable.class, State.class), removed); } } } diff --git a/dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-1.8/src/test/java/executor/ReusableRunnableSubmissionTest.java b/dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-1.8/src/test/java/executor/ReusableRunnableSubmissionTest.java new file mode 100644 index 00000000000..bb7ad910110 --- /dev/null +++ b/dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-1.8/src/test/java/executor/ReusableRunnableSubmissionTest.java @@ -0,0 +1,459 @@ +package executor; + +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import datadog.context.Context; +import datadog.context.ContextKey; +import datadog.context.ContextScope; +import datadog.trace.agent.test.AbstractInstrumentationTest; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.java.concurrent.Wrapper; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.ArrayBlockingQueue; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.PriorityBlockingQueue; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.jupiter.api.Test; + +class ReusableRunnableSubmissionTest extends AbstractInstrumentationTest { + + private static final Object TEST_CONTEXT = ContextKey.named("reusable-runnable-test"); + + @Test + void overlappingSubmissionsKeepTheirContextsAndCommonCaseIdentity() throws Exception { + BlockingTask blocker = new BlockingTask(); + RecordingTask shared = new RecordingTask(2); + ThreadPoolExecutor pool = newPool(3); + try { + pool.execute(blocker); + assertTrue(blocker.started.await(10, TimeUnit.SECONDS)); + submit(pool, shared, "first"); + submit(pool, shared, "second"); + + Object[] queued = pool.getQueue().toArray(); + assertSame(shared, queued[0]); + assertTrue(queued[1] instanceof Wrapper); + + blocker.release.countDown(); + assertTrue(shared.finished.await(10, TimeUnit.SECONDS)); + assertEquals(asList("first", "second"), shared.observed); + } finally { + blocker.release.countDown(); + pool.shutdownNow(); + } + } + + @Test + void directRunCannotStealQueuedSubmissionContext() throws Exception { + BlockingTask blocker = new BlockingTask(); + RecordingTask shared = new RecordingTask(2); + ThreadPoolExecutor pool = newPool(2); + try { + pool.execute(blocker); + assertTrue(blocker.started.await(10, TimeUnit.SECONDS)); + submit(pool, shared, "queued"); + + shared.run(); + blocker.release.countDown(); + + assertTrue(shared.finished.await(10, TimeUnit.SECONDS)); + assertEquals(asList(null, "queued"), shared.observed); + } finally { + blocker.release.countDown(); + pool.shutdownNow(); + } + } + + @Test + void callerRunsUsesRejectedCollisionContext() throws Exception { + BlockingTask blocker = new BlockingTask(); + RecordingTask shared = new RecordingTask(2); + ThreadPoolExecutor pool = + new ThreadPoolExecutor( + 1, + 1, + 0, + TimeUnit.SECONDS, + new ArrayBlockingQueue<>(1), + new ThreadPoolExecutor.CallerRunsPolicy()); + try { + pool.execute(blocker); + assertTrue(blocker.started.await(10, TimeUnit.SECONDS)); + submit(pool, shared, "queued"); + submit(pool, shared, "rejected"); + + blocker.release.countDown(); + assertTrue(shared.finished.await(10, TimeUnit.SECONDS)); + assertEquals(asList("rejected", "queued"), shared.observed); + } finally { + blocker.release.countDown(); + pool.shutdownNow(); + } + } + + @Test + void removeReleasesTheExactSubmissionSlot() throws Exception { + BlockingTask blocker = new BlockingTask(); + RecordingTask shared = new RecordingTask(1); + ThreadPoolExecutor pool = newPool(2); + try { + pool.execute(blocker); + assertTrue(blocker.started.await(10, TimeUnit.SECONDS)); + submit(pool, shared, "removed"); + assertTrue(pool.remove(shared)); + + submit(pool, shared, "replacement"); + assertSame(shared, pool.getQueue().peek()); + blocker.release.countDown(); + + assertTrue(shared.finished.await(10, TimeUnit.SECONDS)); + assertEquals(Collections.singletonList("replacement"), shared.observed); + } finally { + blocker.release.countDown(); + pool.shutdownNow(); + } + } + + @Test + void overlappingSubmissionsCanBeRemovedBySubmittedIdentity() throws Exception { + BlockingTask blocker = new BlockingTask(); + RecordingTask shared = new RecordingTask(2); + ThreadPoolExecutor pool = newPool(2); + try { + pool.execute(blocker); + assertTrue(blocker.started.await(10, TimeUnit.SECONDS)); + submit(pool, shared, "first"); + submit(pool, shared, "second"); + + assertTrue(pool.remove(shared)); + assertTrue(pool.remove(shared)); + assertTrue(pool.getQueue().isEmpty()); + } finally { + blocker.release.countDown(); + pool.shutdownNow(); + } + } + + @Test + void removeSelectsFirstLogicalOccurrenceAfterStateIsRecycled() throws Exception { + BlockingTask blocker = new BlockingTask(); + RecyclingTask shared = new RecyclingTask(); + ThreadPoolExecutor pool = newPool(3); + try { + pool.execute(blocker); + assertTrue(blocker.started.await(10, TimeUnit.SECONDS)); + submit(pool, shared, "running"); + submit(pool, shared, "wrapped"); + + blocker.release.countDown(); + assertTrue(shared.firstStarted.await(10, TimeUnit.SECONDS)); + submit(pool, shared, "recycled"); + + Object[] queued = pool.getQueue().toArray(); + assertTrue(queued[0] instanceof Wrapper); + assertSame(shared, queued[1]); + assertTrue(pool.remove(shared)); + assertSame(shared, pool.getQueue().peek()); + + shared.releaseFirst.countDown(); + assertTrue(shared.finished.await(10, TimeUnit.SECONDS)); + assertEquals(asList("running", "recycled"), shared.observed); + } finally { + blocker.release.countDown(); + shared.releaseFirst.countDown(); + pool.shutdownNow(); + } + } + + @Test + void priorityQueueCollisionPreservesApplicationBehavior() throws Exception { + BlockingTask blocker = new BlockingTask(); + PriorityTask shared = new PriorityTask(); + ThreadPoolExecutor pool = + new ThreadPoolExecutor(1, 1, 0, TimeUnit.SECONDS, new PriorityBlockingQueue()); + try { + pool.execute(blocker); + assertTrue(blocker.started.await(10, TimeUnit.SECONDS)); + submit(pool, shared, "first"); + submit(pool, shared, "not-propagated"); + + Object[] queued = pool.getQueue().toArray(); + assertSame(shared, queued[0]); + assertSame(shared, queued[1]); + blocker.release.countDown(); + + assertTrue(shared.finished.await(10, TimeUnit.SECONDS)); + assertEquals(asList("first", null), shared.observed); + } finally { + blocker.release.countDown(); + pool.shutdownNow(); + } + } + + @Test + void shutdownNowReturnsSubmittedIdentity() throws Exception { + BlockingTask blocker = new BlockingTask(); + RecordingTask shared = new RecordingTask(2); + ThreadPoolExecutor pool = newPool(2); + try { + pool.execute(blocker); + assertTrue(blocker.started.await(10, TimeUnit.SECONDS)); + submit(pool, shared, "first"); + submit(pool, shared, "second"); + + List returned = pool.shutdownNow(); + assertEquals(2, returned.size()); + assertSame(shared, returned.get(0)); + assertSame(shared, returned.get(1)); + } finally { + blocker.release.countDown(); + pool.shutdownNow(); + } + } + + @Test + void decoratingSubclassPropagatesAfterDecoration() throws Exception { + RecordingTask submitted = new RecordingTask(1); + DecoratingExecutor pool = new DecoratingExecutor(); + try { + submit(pool, submitted, "decorated"); + assertTrue(submitted.finished.await(10, TimeUnit.SECONDS)); + assertEquals(Collections.singletonList("decorated"), submitted.observed); + } finally { + pool.shutdownNow(); + } + } + + @Test + void delegatingSubclassKeepsPerSubmissionOwnership() throws Exception { + BlockingTask blocker = new BlockingTask(); + RecordingTask shared = new RecordingTask(2); + ThreadPoolExecutor pool = new DelegatingExecutor(); + try { + pool.execute(blocker); + assertTrue(blocker.started.await(10, TimeUnit.SECONDS)); + submit(pool, shared, "first"); + submit(pool, shared, "second"); + + Object[] queued = pool.getQueue().toArray(); + assertSame(shared, queued[0]); + assertTrue(queued[1] instanceof Wrapper); + + blocker.release.countDown(); + assertTrue(shared.finished.await(10, TimeUnit.SECONDS)); + assertEquals(asList("first", "second"), shared.observed); + } finally { + blocker.release.countDown(); + pool.shutdownNow(); + } + } + + @Test + void nonDelegatingSubclassKeepsNamedTaskPropagation() throws Exception { + RecordingTask submitted = new RecordingTask(1); + NonDelegatingExecutor pool = new NonDelegatingExecutor(); + try { + submit(pool, submitted, "custom"); + + assertTrue(submitted.finished.await(10, TimeUnit.SECONDS)); + assertEquals(Collections.singletonList("custom"), submitted.observed); + } finally { + pool.stopWorker(); + pool.shutdownNow(); + } + } + + @Test + void nonDelegatingSubclassKeepsLambdaPropagation() throws Exception { + List observed = Collections.synchronizedList(new ArrayList<>()); + CountDownLatch finished = new CountDownLatch(1); + Runnable submitted = + () -> { + observed.add(activeSpan()); + finished.countDown(); + }; + NonDelegatingExecutor pool = new NonDelegatingExecutor(); + AgentSpan parent = startSpan("test", "lambda-parent"); + try { + try (ContextScope ignored = activateSpan(parent)) { + pool.execute(submitted); + assertTrue(finished.await(10, TimeUnit.SECONDS)); + } + + assertEquals(Collections.singletonList(parent), observed); + } finally { + parent.finish(); + pool.stopWorker(); + pool.shutdownNow(); + } + } + + private static ThreadPoolExecutor newPool(int queueCapacity) { + return new ThreadPoolExecutor( + 1, 1, 0, TimeUnit.SECONDS, new ArrayBlockingQueue<>(queueCapacity)); + } + + private static void submit(ThreadPoolExecutor pool, Runnable task, String value) { + try (ContextScope ignored = Context.root().with(castContextKey(TEST_CONTEXT), value).attach()) { + pool.execute(task); + } + } + + private static List asList(String first, String second) { + List values = new ArrayList<>(2); + values.add(first); + values.add(second); + return values; + } + + @SuppressWarnings("unchecked") + private static T castContextKey(Object key) { + return (T) key; + } + + private static final class RecordingTask implements Runnable { + private final List observed = Collections.synchronizedList(new ArrayList<>()); + private final CountDownLatch finished; + + private RecordingTask(int executions) { + finished = new CountDownLatch(executions); + } + + @Override + public void run() { + observed.add(Context.current().get(castContextKey(TEST_CONTEXT))); + finished.countDown(); + } + } + + private static final class BlockingTask implements Runnable { + private final CountDownLatch started = new CountDownLatch(1); + private final CountDownLatch release = new CountDownLatch(1); + + @Override + public void run() { + started.countDown(); + try { + release.await(10, TimeUnit.SECONDS); + } catch (InterruptedException interrupted) { + Thread.currentThread().interrupt(); + } + } + } + + private static final class RecyclingTask implements Runnable { + private final List observed = Collections.synchronizedList(new ArrayList<>()); + private final AtomicInteger executions = new AtomicInteger(); + private final CountDownLatch firstStarted = new CountDownLatch(1); + private final CountDownLatch releaseFirst = new CountDownLatch(1); + private final CountDownLatch finished = new CountDownLatch(2); + + @Override + public void run() { + if (executions.getAndIncrement() == 0) { + firstStarted.countDown(); + try { + releaseFirst.await(10, TimeUnit.SECONDS); + } catch (InterruptedException interrupted) { + Thread.currentThread().interrupt(); + } + } + observed.add(Context.current().get(castContextKey(TEST_CONTEXT))); + finished.countDown(); + } + } + + private static final class PriorityTask implements Runnable, Comparable { + private final List observed = Collections.synchronizedList(new ArrayList<>()); + private final CountDownLatch finished = new CountDownLatch(2); + + @Override + public void run() { + observed.add(Context.current().get(castContextKey(TEST_CONTEXT))); + finished.countDown(); + } + + @Override + public int compareTo(PriorityTask ignored) { + return 0; + } + } + + private static final class DecoratingExecutor extends ThreadPoolExecutor { + private DecoratingExecutor() { + super(1, 1, 0, TimeUnit.SECONDS, new ArrayBlockingQueue<>(1)); + } + + @Override + public void execute(Runnable command) { + super.execute(new DelegatingTask(command)); + } + } + + private static final class DelegatingExecutor extends ThreadPoolExecutor { + private DelegatingExecutor() { + super(1, 1, 0, TimeUnit.SECONDS, new ArrayBlockingQueue<>(3)); + } + + @Override + public void execute(Runnable command) { + super.execute(command); + } + } + + private static final class NonDelegatingExecutor extends ThreadPoolExecutor { + private final BlockingQueue tasks = new LinkedBlockingQueue<>(); + private final Thread worker; + + private NonDelegatingExecutor() { + super(1, 1, 0, TimeUnit.SECONDS, new ArrayBlockingQueue<>(1)); + worker = new Thread(this::runTasks, "non-delegating-executor-test"); + worker.setDaemon(true); + worker.start(); + } + + @Override + public void execute(Runnable command) { + tasks.add(command); + } + + private void runTasks() { + try { + while (!Thread.currentThread().isInterrupted()) { + tasks.take().run(); + } + } catch (InterruptedException interrupted) { + Thread.currentThread().interrupt(); + } + } + + private void stopWorker() throws InterruptedException { + worker.interrupt(); + worker.join(TimeUnit.SECONDS.toMillis(10)); + } + } + + private static final class DelegatingTask implements Runnable { + private final Runnable delegate; + + private DelegatingTask(Runnable delegate) { + this.delegate = delegate; + } + + @Override + public void run() { + delegate.run(); + } + } +}