Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,9 @@
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.isAsyncPropagationEnabled;

import datadog.context.ContextContinuation;
import datadog.context.ContextScope;
import datadog.trace.bootstrap.ContextStore;
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;

/** Helper utils for Runnable/Callable instrumentation */
Expand All @@ -18,16 +19,16 @@ public class AdviceUtils {
* @param <T> task's type
* @return scope if scope was started, or null
*/
public static <T> AgentScope startTaskScope(
public static <T> ContextScope startTaskScope(
final ContextStore<T, State> contextStore, final T task) {
return startTaskScope(contextStore.get(task));
}

public static AgentScope startTaskScope(State state) {
public static ContextScope startTaskScope(State state) {
if (state != null) {
final AgentScope.Continuation continuation = state.getAndResetContinuation();
final ContextContinuation continuation = state.getAndResetContinuation();
if (continuation != null) {
final AgentScope scope = continuation.activate();
final ContextScope scope = continuation.resume();
// important - stop timing after the scope has been activated so the time in the queue can
// be attributed to the correct context without duplicating the propagated information
state.stopTiming();
Expand All @@ -37,7 +38,7 @@ public static AgentScope startTaskScope(State state) {
return null;
}

public static void endTaskScope(final AgentScope scope) {
public static void endTaskScope(final ContextScope scope) {
if (null != scope) {
scope.close();
}
Expand Down
Original file line number Diff line number Diff line change
@@ -1,11 +1,11 @@
package datadog.trace.bootstrap.instrumentation.java.concurrent;

import datadog.trace.bootstrap.instrumentation.api.AgentScope;
import datadog.context.ContextContinuation;

public final class ComparableRunnable<T extends Runnable & Comparable<T>> extends Wrapper<T>
implements Comparable<ComparableRunnable<T>> {

public ComparableRunnable(T delegate, AgentScope.Continuation continuation) {
public ComparableRunnable(T delegate, ContextContinuation continuation) {
super(delegate, continuation);
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,9 @@
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.isAsyncPropagationEnabled;
import static datadog.trace.bootstrap.instrumentation.java.concurrent.ContinuationClaim.CLAIMED;

import datadog.context.ContextContinuation;
import datadog.context.ContextScope;
import datadog.trace.bootstrap.ContextStore;
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
import java.util.concurrent.atomic.AtomicReferenceFieldUpdater;
import org.slf4j.Logger;
Expand All @@ -22,12 +23,12 @@ public final class ConcurrentState {

public static ContextStore.Factory<ConcurrentState> FACTORY = ConcurrentState::new;

private volatile AgentScope.Continuation continuation = null;
private volatile ContextContinuation continuation = null;

private static final AtomicReferenceFieldUpdater<ConcurrentState, AgentScope.Continuation>
private static final AtomicReferenceFieldUpdater<ConcurrentState, ContextContinuation>
CONTINUATION =
AtomicReferenceFieldUpdater.newUpdater(
ConcurrentState.class, AgentScope.Continuation.class, "continuation");
ConcurrentState.class, ContextContinuation.class, "continuation");

private ConcurrentState() {}

Expand All @@ -44,7 +45,7 @@ public static <K> ConcurrentState captureContinuation(
return state;
}

public static <K> AgentScope activateAndContinueContinuation(
public static <K> ContextScope activateAndContinueContinuation(
ContextStore<K, ConcurrentState> contextStore, K key) {
final ConcurrentState state = contextStore.get(key);
if (state == null) {
Expand All @@ -54,7 +55,10 @@ public static <K> AgentScope activateAndContinueContinuation(
}

public static <K> void closeScope(
ContextStore<K, ConcurrentState> contextStore, K key, AgentScope scope, Throwable throwable) {
ContextStore<K, ConcurrentState> contextStore,
K key,
ContextScope scope,
Throwable throwable) {
final ConcurrentState state = contextStore.get(key);
if (scope != null) {
scope.close();
Expand Down Expand Up @@ -88,27 +92,27 @@ private boolean captureAndSetContinuation(final AgentSpan span) {
return false;
}

private AgentScope activateAndContinueContinuation() {
final AgentScope.Continuation continuation = CONTINUATION.get(this);
private ContextScope activateAndContinueContinuation() {
final ContextContinuation continuation = CONTINUATION.get(this);
if (continuation != null && continuation != CLAIMED) {
return continuation.activate();
return continuation.resume();
}
return null;
}

private void cancelContinuation() {
final AgentScope.Continuation continuation = CONTINUATION.get(this);
final ContextContinuation continuation = CONTINUATION.get(this);
if (continuation != null && continuation != CLAIMED) {
continuation.cancel();
continuation.release();
}
}

private void cancelAndClearContinuation() {
final AgentScope.Continuation continuation = CONTINUATION.get(this);
final ContextContinuation continuation = CONTINUATION.get(this);
if (continuation != null && continuation != CLAIMED) {
// We should never be able to reuse this state
CONTINUATION.compareAndSet(this, continuation, CLAIMED);
continuation.cancel();
continuation.release();
}
}
}
Original file line number Diff line number Diff line change
@@ -1,29 +1,30 @@
package datadog.trace.bootstrap.instrumentation.java.concurrent;

import datadog.trace.bootstrap.instrumentation.api.AgentScope;
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
import datadog.context.Context;
import datadog.context.ContextContinuation;
import datadog.context.ContextScope;

final class ContinuationClaim implements AgentScope.Continuation {
final class ContinuationClaim implements ContextContinuation {

public static final ContinuationClaim CLAIMED = new ContinuationClaim();

@Override
public AgentScope.Continuation hold() {
public ContextContinuation hold() {
throw new IllegalStateException();
}

@Override
public AgentScope activate() {
public ContextScope resume() {
throw new IllegalStateException();
}

@Override
public AgentSpan span() {
public Context context() {
throw new IllegalStateException();
}

@Override
public void cancel() {
public void release() {
throw new IllegalStateException();
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -3,21 +3,21 @@
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.captureSpan;
import static datadog.trace.bootstrap.instrumentation.java.concurrent.ContinuationClaim.CLAIMED;

import datadog.context.ContextContinuation;
import datadog.trace.api.profiling.Timing;
import datadog.trace.bootstrap.ContextStore;
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
import java.util.concurrent.atomic.AtomicReferenceFieldUpdater;

public final class State {

public static ContextStore.Factory<State> FACTORY = State::new;

private static final AtomicReferenceFieldUpdater<State, AgentScope.Continuation> CONTINUATION =
private static final AtomicReferenceFieldUpdater<State, ContextContinuation> CONTINUATION =
AtomicReferenceFieldUpdater.newUpdater(
State.class, AgentScope.Continuation.class, "continuation");
State.class, ContextContinuation.class, "continuation");

private volatile AgentScope.Continuation continuation = null;
private volatile ContextContinuation continuation = null;

private static final AtomicReferenceFieldUpdater<State, Timing> TIMING =
AtomicReferenceFieldUpdater.newUpdater(State.class, Timing.class, "timing");
Expand All @@ -40,34 +40,34 @@ public boolean captureAndSetContinuation(final AgentSpan span) {
return false;
}

public boolean setOrCancelContinuation(final AgentScope.Continuation continuation) {
public boolean setOrCancelContinuation(final ContextContinuation continuation) {
if (CONTINUATION.compareAndSet(this, null, CLAIMED)) {
// lazy write is guaranteed to be seen by getAndSet
CONTINUATION.lazySet(this, continuation);
return true;
} else {
continuation.cancel();
continuation.release();
return false;
}
}

public void closeContinuation() {
AgentScope.Continuation continuation = getAndResetContinuation();
ContextContinuation continuation = getAndResetContinuation();
if (null != continuation) {
continuation.cancel();
continuation.release();
}
}

public AgentSpan getSpan() {
AgentScope.Continuation continuation = CONTINUATION.get(this);
ContextContinuation continuation = CONTINUATION.get(this);
if (null != continuation) {
return continuation.span();
return AgentSpan.fromContext(continuation.context());
}
return null;
}

public AgentScope.Continuation getAndResetContinuation() {
AgentScope.Continuation continuation = CONTINUATION.get(this);
public ContextContinuation getAndResetContinuation() {
ContextContinuation continuation = CONTINUATION.get(this);
if (null == continuation || CLAIMED == continuation) {
return null;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3,18 +3,18 @@
import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.ExcludeType.RUNNABLE;
import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.exclude;

import datadog.context.ContextScope;
import datadog.trace.api.GenericClassValue;
import datadog.trace.api.InstrumenterConfig;
import datadog.trace.bootstrap.ContextStore;
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
import java.util.Set;
import java.util.concurrent.ThreadPoolExecutor;

/**
* This class is a helper for the ThreadPoolExecutorInstrumentation. The instrumentation has two
* modes, where the legacy mode uses wrapping of the Runnable and the new mode uses the State
* context store field in the actual Runnable. The ThreadLocal below is needed to transport the
* AgentScope between two methods when using the context store approach. More details can be found
* ContextScope between two methods when using the context store approach. More details can be found
* in the ThreadPoolExecutorInstrumentation.
*/
public final class TPEHelper {
Expand All @@ -24,7 +24,7 @@ public final class TPEHelper {
// A ThreadPoolExecutor with one of these types will newer be propagated/wrapped
private static final Set<String> excludedClasses;
// A ThreadLocal to store the Scope between beforeExecute and afterExecute if wrapping is not used
private static final ThreadLocal<AgentScope> threadLocalScope;
private static final ThreadLocal<ContextScope> threadLocalScope;

private static final ClassValue<Boolean> WRAP =
GenericClassValue.of(
Expand Down Expand Up @@ -78,37 +78,37 @@ public static void capture(ContextStore<Runnable, State> contextStore, Runnable
}
}

public static AgentScope startScope(ContextStore<Runnable, State> contextStore, Runnable task) {
public static ContextScope startScope(ContextStore<Runnable, State> contextStore, Runnable task) {
if (task == null || exclude(RUNNABLE, task)) {
return null;
}
return AdviceUtils.startTaskScope(contextStore, task);
}

public static void setThreadLocalScope(AgentScope scope, Runnable task) {
public static void setThreadLocalScope(ContextScope scope, Runnable task) {
if (scope == null || task == null || exclude(RUNNABLE, task)) {
return;
}
AgentScope current = threadLocalScope.get();
ContextScope current = threadLocalScope.get();
if (current != null) {
current.close();
}
threadLocalScope.set(scope);
}

public static AgentScope getAndClearThreadLocalScope(Runnable task) {
public static ContextScope getAndClearThreadLocalScope(Runnable task) {
if (task == null || exclude(RUNNABLE, task)) {
return null;
}
AgentScope scope = threadLocalScope.get();
ContextScope scope = threadLocalScope.get();
// Intentionally use `.set(null)` instead of `.remove()` for performance reasons.
// For details see: https://github.com/DataDog/dd-trace-java/pull/9856#discussion_r2527729963
// noinspection ThreadLocalSetWithNull
threadLocalScope.set(null);
return scope;
}

public static void endScope(AgentScope scope, Runnable task) {
public static void endScope(ContextScope scope, Runnable task) {
if (task == null || exclude(RUNNABLE, task)) {
return;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,8 @@
import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.ExcludeType.RUNNABLE;
import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.exclude;

import datadog.context.ContextContinuation;
import datadog.context.ContextScope;
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
import java.util.concurrent.RunnableFuture;

public class Wrapper<T extends Runnable> implements Runnable, AutoCloseable {
Expand All @@ -19,7 +19,7 @@ public static <T extends Runnable> Runnable wrap(T task) {
|| exclude(RUNNABLE, task)) {
return task;
}
AgentScope.Continuation continuation = captureActiveSpan();
ContextContinuation continuation = captureActiveSpan();
if (continuation != noopContinuation()) {
if (task instanceof Comparable) {
return new ComparableRunnable(task, continuation);
Expand All @@ -35,9 +35,9 @@ public static Runnable unwrap(Runnable task) {
}

protected final T delegate;
private final AgentScope.Continuation continuation;
private final ContextContinuation continuation;

public Wrapper(T delegate, AgentScope.Continuation continuation) {
public Wrapper(T delegate, ContextContinuation continuation) {
this.delegate = delegate;
this.continuation = continuation;
}
Expand All @@ -51,16 +51,16 @@ public void run() {

public void cancel() {
if (null != continuation) {
continuation.cancel();
continuation.release();
}
}

public T unwrap() {
return delegate;
}

private AgentScope activate() {
return null == continuation ? null : continuation.activate();
private ContextScope activate() {
return null == continuation ? null : continuation.resume();
}

@Override
Expand Down
Loading