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 @@ -22,10 +22,9 @@

public class RedisAPICallAdvice {
@Advice.OnMethodEnter(suppress = Throwable.class)
public static boolean beforeCall(
public static AgentScope beforeCall(
@Advice.Origin final Method currentMethod,
@Advice.This final RedisAPI self,
@Advice.Local("callScope") AgentScope scope,
@Advice.Argument(
value = 0,
readOnly = false,
Expand Down Expand Up @@ -62,7 +61,7 @@ public static boolean beforeCall(
// either), so this seems to be the only way to communicate that we have already wrapped
// the handler. :(
if (CallDepthThreadLocalMap.incrementCallDepth(RedisAPI.class) > 0) {
return true;
return null;
}

// TODO what is the recreated for every read about in the @Advice.Origin javadoc?
Expand Down Expand Up @@ -100,7 +99,7 @@ public static boolean beforeCall(
}

if (null == handler || handler instanceof ResponseHandlerWrapper) {
return true;
return null;
}

final AgentSpan parentSpan = activeSpan();
Expand All @@ -112,7 +111,7 @@ public static boolean beforeCall(
The potential racy condition when the handler may be added to an already finished task is handled
by RedisAPIImplSendAdvice.
*/
scope = activateSpan(clientSpan);
AgentScope scope = activateSpan(clientSpan);
ResponseHandlerWrapper respHandler =
new ResponseHandlerWrapper(handler, clientSpan, parentContinuation);
handler = respHandler;
Expand All @@ -139,23 +138,24 @@ public static boolean beforeCall(
Store the response handler in the context so that it can be retrieved in RedisAPIImplSendAdvice
*/
InstrumentationContext.get(RedisAPI.class, ResponseHandlerWrapper.class).put(self, respHandler);
return true;
return scope;
}

@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
public static void afterCall(
@Advice.Thrown final Throwable throwable,
@Advice.This final RedisAPI self,
@Advice.Local("callScope") AgentScope scope,
@Advice.Enter final boolean decrement) {
if (decrement) {
CallDepthThreadLocalMap.decrementCallDepth(RedisAPI.class);
@Advice.Enter final AgentScope scope) {
CallDepthThreadLocalMap.decrementCallDepth(RedisAPI.class);

if (null == scope) {
return;
}

scope.close();

// Clean the response handler from the context
InstrumentationContext.get(RedisAPI.class, ResponseHandlerWrapper.class).put(self, null);
InstrumentationContext.get(RedisAPI.class, ResponseHandlerWrapper.class).remove(self);
}

// Only apply this advice for versions that we instrument 3.9.x
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@
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.captureSpan;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopScope;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopSpan;
import static datadog.trace.instrumentation.vertx_redis_client.VertxRedisClientDecorator.DECORATE;
import static datadog.trace.instrumentation.vertx_redis_client.VertxRedisClientDecorator.REDIS_COMMAND;
Expand Down Expand Up @@ -32,6 +31,10 @@ public static AgentScope beforeSend(
@Advice.Argument(value = 0, readOnly = false) Request request,
@Advice.Local("ddParentContinuation") ContextContinuation parentContinuation)
throws Throwable {
// If we had already wrapped the innermost handler in the RedisAPI call, then we should
// not wrap it again here. See comment in RedisAPICallAdvice
boolean nested = CallDepthThreadLocalMap.incrementCallDepth(RedisAPI.class) > 0;

ContextStore<Request, Boolean> ctxt = InstrumentationContext.get(Request.class, Boolean.class);
Boolean handled = ctxt.get(request);
if (null != handled && handled) {
Expand All @@ -44,6 +47,12 @@ public static AgentScope beforeSend(
}
ctxt.put(request, Boolean.TRUE);

// Mark the request handled even when nested, so a later async re-send of the same
// Request (e.g. via a pooled connection) isn't mistaken for a brand-new command.
if (nested) {
return null;
}

AgentSpan parentSpan = activeSpan();

if (parentSpan != null && REDIS_COMMAND.equals(parentSpan.getOperationName())) {
Expand All @@ -53,12 +62,6 @@ public static AgentScope beforeSend(

parentContinuation = null == parentSpan ? captureSpan(noopSpan()) : captureSpan(parentSpan);

// If we had already wrapped the innermost handler in the RedisAPI call, then we should
// not wrap it again here. See comment in RedisAPICallAdvice
if (CallDepthThreadLocalMap.incrementCallDepth(RedisAPI.class) > 0) {
return noopScope();
}

final AgentSpan clientSpan =
DECORATE.startAndDecorateSpan(
request.command(), InstrumentationContext.get(Command.class, UTF8BytesString.class));
Expand All @@ -72,20 +75,19 @@ public static void afterSend(
@Advice.Local("ddParentContinuation") ContextContinuation parentContinuation,
@Advice.Enter final AgentScope clientScope,
@Advice.This final Object thiz) {
CallDepthThreadLocalMap.decrementCallDepth(RedisAPI.class);
if (thiz instanceof RedisConnection) {
final SocketAddress socketAddress =
InstrumentationContext.get(RedisConnection.class, SocketAddress.class)
.get((RedisConnection) thiz);
final AgentSpan span = clientScope != null ? clientScope.span() : activeSpan();

if (socketAddress != null && span != null) {
final AgentSpan spanWithConnection = clientScope == noopScope() ? activeSpan() : span;
DECORATE.onConnection(spanWithConnection, socketAddress);
DECORATE.setPeerPort(spanWithConnection, socketAddress.port());
// Verify the activeSpan() fallback is actually a REDIS_COMMAND span
if (socketAddress != null && span != null && REDIS_COMMAND.equals(span.getOperationName())) {
DECORATE.onConnection(span, socketAddress);
DECORATE.setPeerPort(span, socketAddress.port());
}
}
if (clientScope != null) {
CallDepthThreadLocalMap.decrementCallDepth(RedisAPI.class);
Promise<Response> promise = Promise.promise();
responseFuture.onComplete(
new ResponseHandler(promise, clientScope.span(), parentContinuation));
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@
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.captureSpan;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopScope;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopSpan;
import static datadog.trace.instrumentation.vertx_redis_client.VertxRedisClientDecorator.DECORATE;
import static datadog.trace.instrumentation.vertx_redis_client.VertxRedisClientDecorator.REDIS_COMMAND;
Expand Down Expand Up @@ -33,6 +32,10 @@ public static AgentScope beforeSend(
@Advice.Argument(value = 0, readOnly = false) Request request,
@Advice.Argument(value = 1, readOnly = false) Handler<AsyncResult<Response>> handler)
throws Throwable {
// If we had already wrapped the innermost handler in the RedisAPI call, then we should
// not wrap it again here. See comment in RedisAPICallAdvice
boolean nested = CallDepthThreadLocalMap.incrementCallDepth(RedisAPI.class) > 0;

if (null == handler || handler instanceof ResponseHandlerWrapper) {
return null;
}
Expand All @@ -49,10 +52,10 @@ public static AgentScope beforeSend(
}
ctxt.put(request, Boolean.TRUE);

// If we had already wrapped the innermost handler in the RedisAPI call, then we should
// not wrap it again here. See comment in RedisAPICallAdvice
if (CallDepthThreadLocalMap.incrementCallDepth(RedisAPI.class) > 0) {
return noopScope();
// Mark the request handled even when nested, so a later async re-send of the same
// Request (e.g. via a pooled connection) isn't mistaken for a brand-new command.
if (nested) {
return null;
}

AgentSpan parentSpan = activeSpan();
Expand All @@ -74,18 +77,19 @@ public static AgentScope beforeSend(
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
public static void afterSend(
@Advice.Enter final AgentScope clientScope, @Advice.This final Object thiz) {
CallDepthThreadLocalMap.decrementCallDepth(RedisAPI.class);
if (thiz instanceof RedisConnection) {
final SocketAddress socketAddress =
InstrumentationContext.get(RedisConnection.class, SocketAddress.class)
.get((RedisConnection) thiz);
final AgentSpan span = clientScope != null ? clientScope.span() : activeSpan();
if (socketAddress != null && span != null) {
// Verify the activeSpan() fallback is actually a REDIS_COMMAND span
if (socketAddress != null && span != null && REDIS_COMMAND.equals(span.getOperationName())) {
DECORATE.onConnection(span, socketAddress);
DECORATE.setPeerPort(span, socketAddress.port());
}
}
if (null != clientScope) {
CallDepthThreadLocalMap.decrementCallDepth(RedisAPI.class);
clientScope.close();
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ import datadog.trace.agent.test.asserts.ListWriterAssert
import datadog.trace.agent.test.asserts.TraceAssert
import datadog.trace.agent.test.naming.VersionedNamingTestBase
import datadog.trace.api.DDSpanTypes
import datadog.trace.bootstrap.CallDepthThreadLocalMap
import datadog.trace.bootstrap.instrumentation.api.Tags
import datadog.trace.core.DDSpan
import io.vertx.core.AsyncResult
Expand All @@ -12,6 +13,7 @@ import io.vertx.core.Vertx
import io.vertx.core.VertxOptions
import io.vertx.redis.client.Command
import io.vertx.redis.client.Redis
import io.vertx.redis.client.RedisAPI
import io.vertx.redis.client.Request
import io.vertx.redis.client.Response
import org.testcontainers.containers.wait.strategy.Wait
Expand Down Expand Up @@ -75,6 +77,10 @@ abstract class VertxRedisTestBase extends VersionedNamingTestBase {
cleanUpTraces()
}

def cleanup() {
assert CallDepthThreadLocalMap.getCallDepth(RedisAPI) == 0
}
Comment thread
Copilot marked this conversation as resolved.

void cleanUpTraces() {
def cleanupSpan = runUnderTrace("cleanup") {
redis.send(Request.cmd(Command.FLUSHALL), Promise.promise())
Expand Down
Loading