Skip to content

Commit 849a199

Browse files
amarzialiamarziali
authored
Merge pull request #5038 from DataDog/andrea.marziali/rediscala-peer-info
add remote peer info for rediscala
2 parents aac4c0d + 38aeb61 commit 849a199

6 files changed

Lines changed: 141 additions & 13 deletions

File tree

‎dd-java-agent/instrumentation/rediscala-1.8.0/src/main/java/datadog/trace/instrumentation/rediscala/OnCompleteHandler.java‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,14 +3,23 @@
33
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope;
44
import static datadog.trace.instrumentation.rediscala.RediscalaClientDecorator.DECORATE;
55

6+
import akka.actor.ActorRef;
7+
import datadog.trace.bootstrap.ContextStore;
68
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
79
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
810
import scala.runtime.AbstractFunction1;
911
import scala.util.Try;
1012

1113
public final class OnCompleteHandler extends AbstractFunction1<Try<Object>, Void> {
1214

13-
public static final OnCompleteHandler INSTANCE = new OnCompleteHandler();
15+
private final ContextStore<ActorRef, RedisConnectionInfo> contextStore;
16+
private final ActorRef actorRef;
17+
18+
public OnCompleteHandler(
19+
ContextStore<ActorRef, RedisConnectionInfo> contextStore, ActorRef actorRef) {
20+
this.contextStore = contextStore;
21+
this.actorRef = actorRef;
22+
}
1423

1524
@Override
1625
public Void apply(final Try<Object> result) {
@@ -19,6 +28,9 @@ public Void apply(final Try<Object> result) {
1928
if (null != activeScope) {
2029
AgentSpan span = activeScope.span();
2130
try {
31+
if (actorRef != null) {
32+
DECORATE.onConnection(span, contextStore.get(actorRef));
33+
}
2234
if (result.isFailure()) {
2335
DECORATE.onError(span, result.failed().get());
2436
}

‎dd-java-agent/instrumentation/rediscala-1.8.0/src/main/java/datadog/trace/instrumentation/rediscala/RedisClientActorInstrumentation.java‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
package datadog.trace.instrumentation.rediscala;
2+
3+
import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.extendsClass;
4+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
5+
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
6+
7+
import akka.actor.ActorRef;
8+
import com.google.auto.service.AutoService;
9+
import datadog.trace.agent.tooling.Instrumenter;
10+
import datadog.trace.bootstrap.InstrumentationContext;
11+
import java.util.Collections;
12+
import java.util.Map;
13+
import net.bytebuddy.asm.Advice;
14+
import net.bytebuddy.description.type.TypeDescription;
15+
import net.bytebuddy.matcher.ElementMatcher;
16+
import redis.RedisClientActorLike;
17+
18+
@AutoService(Instrumenter.class)
19+
public class RedisClientActorInstrumentation extends Instrumenter.Tracing
20+
implements Instrumenter.ForTypeHierarchy {
21+
public RedisClientActorInstrumentation() {
22+
super("rediscala", "redis", "rediscala-connection");
23+
}
24+
25+
@Override
26+
public String hierarchyMarkerType() {
27+
return "redis.RedisClientActorLike";
28+
}
29+
30+
@Override
31+
public ElementMatcher<TypeDescription> hierarchyMatcher() {
32+
return extendsClass(named(hierarchyMarkerType()));
33+
}
34+
35+
@Override
36+
public String[] helperClassNames() {
37+
return new String[] {packageName + ".RedisConnectionInfo"};
38+
}
39+
40+
@Override
41+
public Map<String, String> contextStore() {
42+
return Collections.singletonMap("akka.actor.ActorRef", packageName + ".RedisConnectionInfo");
43+
}
44+
45+
@Override
46+
public void adviceTransformations(AdviceTransformation transformation) {
47+
transformation.applyAdvice(
48+
isMethod().and(named("onConnect")),
49+
RedisClientActorInstrumentation.class.getName() + "$ConnectionAdvice");
50+
}
51+
52+
public static class ConnectionAdvice {
53+
@Advice.OnMethodEnter(suppress = Throwable.class)
54+
public static void beforeOnConnect(@Advice.This final RedisClientActorLike thiz) {
55+
if (thiz.redisConnection() != null) {
56+
final Object tmpDbIndex = thiz.db().isDefined() ? thiz.db().get() : null;
57+
final int dbIndex = (tmpDbIndex instanceof Number) ? ((Number) tmpDbIndex).intValue() : 0;
58+
InstrumentationContext.get(ActorRef.class, RedisConnectionInfo.class)
59+
.put(
60+
thiz.redisConnection(), new RedisConnectionInfo(thiz.host(), thiz.port(), dbIndex));
61+
}
62+
}
63+
}
64+
}

‎dd-java-agent/instrumentation/rediscala-1.8.0/src/main/java/datadog/trace/instrumentation/rediscala/RedisConnectionInfo.java‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
package datadog.trace.instrumentation.rediscala;
2+
3+
public class RedisConnectionInfo {
4+
public final String host;
5+
public final int port;
6+
public final int dbIndex;
7+
8+
public RedisConnectionInfo(String host, int port, int dbIndex) {
9+
this.host = host;
10+
this.port = port;
11+
this.dbIndex = dbIndex;
12+
}
13+
}

‎dd-java-agent/instrumentation/rediscala-1.8.0/src/main/java/datadog/trace/instrumentation/rediscala/RediscalaClientDecorator.java‎

Lines changed: 14 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -6,11 +6,9 @@
66
import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes;
77
import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
88
import datadog.trace.bootstrap.instrumentation.decorator.DBTypeProcessingDatabaseClientDecorator;
9-
import redis.RedisCommand;
10-
import redis.protocol.RedisReply;
119

1210
public class RediscalaClientDecorator
13-
extends DBTypeProcessingDatabaseClientDecorator<RedisCommand<? extends RedisReply, ?>> {
11+
extends DBTypeProcessingDatabaseClientDecorator<RedisConnectionInfo> {
1412

1513
private static final CharSequence COMPONENT_NAME = UTF8BytesString.create("redis-command");
1614

@@ -47,20 +45,29 @@ protected String dbType() {
4745
}
4846

4947
@Override
50-
protected String dbUser(final RedisCommand<? extends RedisReply, ?> session) {
48+
protected String dbUser(final RedisConnectionInfo redisConnectionInfo) {
5149
return null;
5250
}
5351

5452
@Override
55-
protected String dbInstance(final RedisCommand<? extends RedisReply, ?> session) {
53+
protected String dbInstance(final RedisConnectionInfo redisConnectionInfo) {
5654
return null;
5755
}
5856

5957
@Override
60-
protected String dbHostname(RedisCommand<? extends RedisReply, ?> redisCommand) {
61-
return null;
58+
protected String dbHostname(final RedisConnectionInfo redisConnectionInfo) {
59+
return redisConnectionInfo.host;
6260
}
6361

6462
@Override
6563
protected void postProcessServiceAndOperationName(AgentSpan span, String dbType) {}
64+
65+
@Override
66+
public AgentSpan onConnection(final AgentSpan span, final RedisConnectionInfo connection) {
67+
if (connection != null) {
68+
setPeerPort(span, connection.port);
69+
span.setTag("db.redis.dbIndex", connection.dbIndex);
70+
}
71+
return super.onConnection(span, connection);
72+
}
6673
}

‎dd-java-agent/instrumentation/rediscala-1.8.0/src/main/java/datadog/trace/instrumentation/rediscala/RediscalaInstrumentation.java‎

Lines changed: 26 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11,14 +11,20 @@
1111
import static net.bytebuddy.matcher.ElementMatchers.returns;
1212
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
1313

14+
import akka.actor.ActorRef;
1415
import com.google.auto.service.AutoService;
1516
import datadog.trace.agent.tooling.Instrumenter;
1617
import datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers;
18+
import datadog.trace.bootstrap.ContextStore;
19+
import datadog.trace.bootstrap.InstrumentationContext;
1720
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
1821
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
22+
import java.util.Collections;
23+
import java.util.Map;
1924
import net.bytebuddy.asm.Advice;
2025
import net.bytebuddy.description.type.TypeDescription;
2126
import net.bytebuddy.matcher.ElementMatcher;
27+
import redis.ActorRequest;
2228
import redis.RedisCommand;
2329
import scala.concurrent.ExecutionContext;
2430
import scala.concurrent.Future;
@@ -51,10 +57,17 @@ public ElementMatcher<TypeDescription> hierarchyMatcher() {
5157
@Override
5258
public String[] helperClassNames() {
5359
return new String[] {
54-
packageName + ".OnCompleteHandler", packageName + ".RediscalaClientDecorator",
60+
packageName + ".OnCompleteHandler",
61+
packageName + ".RediscalaClientDecorator",
62+
packageName + ".RedisConnectionInfo"
5563
};
5664
}
5765

66+
@Override
67+
public Map<String, String> contextStore() {
68+
return Collections.singletonMap("akka.actor.ActorRef", packageName + ".RedisConnectionInfo");
69+
}
70+
5871
@Override
5972
public void adviceTransformations(AdviceTransformation transformation) {
6073
transformation.applyAdvice(
@@ -80,14 +93,24 @@ public static AgentScope onEnter(@Advice.Argument(0) final RedisCommand cmd) {
8093
public static void stopSpan(
8194
@Advice.Enter final AgentScope scope,
8295
@Advice.Thrown final Throwable throwable,
96+
@Advice.This final Object thiz,
8397
@Advice.FieldValue("executionContext") final ExecutionContext ctx,
8498
@Advice.Return(readOnly = false) final Future<Object> responseFuture) {
8599

86100
final AgentSpan span = scope.span();
87-
101+
final ContextStore<ActorRef, RedisConnectionInfo> contextStore =
102+
InstrumentationContext.get(ActorRef.class, RedisConnectionInfo.class);
103+
ActorRef connection = null;
104+
if (thiz instanceof ActorRequest) {
105+
connection = ((ActorRequest) thiz).redisConnection();
106+
}
88107
if (throwable == null) {
89-
responseFuture.onComplete(OnCompleteHandler.INSTANCE, ctx);
108+
responseFuture.onComplete(new OnCompleteHandler(contextStore, connection), ctx);
90109
} else {
110+
if (connection != null) {
111+
// try to get the info early
112+
DECORATE.onConnection(span, contextStore.get(connection));
113+
}
91114
DECORATE.onError(span, throwable);
92115
DECORATE.beforeFinish(span);
93116
span.finish();

‎dd-java-agent/instrumentation/rediscala-1.8.0/src/test/groovy/RediscalaClientTest.groovy‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
import static datadog.trace.api.config.TraceInstrumentationConfig.DB_CLIENT_HOST_SPLIT_BY_INSTANCE
2+
13
import akka.actor.ActorSystem
24
import datadog.trace.agent.test.naming.VersionedNamingTestBase
35
import datadog.trace.agent.test.utils.PortUtils
@@ -14,8 +16,6 @@ import scala.concurrent.Await
1416
import scala.concurrent.duration.Duration
1517
import spock.lang.Shared
1618

17-
import static datadog.trace.api.config.TraceInstrumentationConfig.DB_CLIENT_HOST_SPLIT_BY_INSTANCE
18-
1919
abstract class RediscalaClientTest extends VersionedNamingTestBase {
2020

2121
@Shared
@@ -92,6 +92,9 @@ abstract class RediscalaClientTest extends VersionedNamingTestBase {
9292
"$Tags.COMPONENT" "redis-command"
9393
"$Tags.SPAN_KIND" Tags.SPAN_KIND_CLIENT
9494
"$Tags.DB_TYPE" "redis"
95+
"$Tags.PEER_HOSTNAME" redisClient.host()
96+
"$Tags.PEER_PORT" redisClient.port()
97+
"db.redis.dbIndex" 0
9598
defaultTags()
9699
}
97100
}
@@ -124,6 +127,9 @@ abstract class RediscalaClientTest extends VersionedNamingTestBase {
124127
"$Tags.COMPONENT" "redis-command"
125128
"$Tags.SPAN_KIND" Tags.SPAN_KIND_CLIENT
126129
"$Tags.DB_TYPE" "redis"
130+
"$Tags.PEER_HOSTNAME" redisClient.host()
131+
"$Tags.PEER_PORT" redisClient.port()
132+
"db.redis.dbIndex" 0
127133
defaultTags()
128134
}
129135
}
@@ -138,6 +144,9 @@ abstract class RediscalaClientTest extends VersionedNamingTestBase {
138144
"$Tags.COMPONENT" "redis-command"
139145
"$Tags.SPAN_KIND" Tags.SPAN_KIND_CLIENT
140146
"$Tags.DB_TYPE" "redis"
147+
"$Tags.PEER_HOSTNAME" redisClient.host()
148+
"$Tags.PEER_PORT" redisClient.port()
149+
"db.redis.dbIndex" 0
141150
defaultTags()
142151
}
143152
}

0 commit comments

Comments
 (0)