Skip to content

Commit 78c6f88

Browse files
tylerbensontylerbenson
authored
Merge pull request DataDog#535 from DataDog/tyler/rabbit
Add instrumentation for RabbitMQ’s AMQP library
2 parents 80f160b + b693335 commit 78c6f88

18 files changed

Lines changed: 946 additions & 11 deletions

File tree

‎.circleci/config.yml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,8 @@ jobs:
5252
- image: *default_container
5353
# This is used by spymemcached instrumentation tests
5454
- image: memcached
55+
# This is used by rabbitmq instrumentation tests
56+
- image: rabbitmq
5557

5658
steps:
5759
- checkout

‎dd-java-agent/instrumentation/rabbitmq-amqp-2.7/rabbitmq-amqp-2.7.gradle‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
muzzle {
2+
pass {
3+
group = "com.rabbitmq"
4+
module = 'amqp-client'
5+
versions = "[2.7.0,)"
6+
assertInverse = true
7+
}
8+
}
9+
10+
apply from: "${rootDir}/gradle/java.gradle"
11+
12+
apply plugin: 'org.unbroken-dome.test-sets'
13+
14+
testSets {
15+
latestDepTest {
16+
dirName = 'test'
17+
}
18+
}
19+
20+
dependencies {
21+
compileOnly group: 'com.rabbitmq', name: 'amqp-client', version: '2.7.0'
22+
23+
compile project(':dd-trace-ot')
24+
compile project(':dd-java-agent:agent-tooling')
25+
26+
compile deps.bytebuddy
27+
compile deps.opentracing
28+
compile deps.autoservice
29+
annotationProcessor deps.autoservice
30+
implementation deps.autoservice
31+
32+
testCompile project(':dd-java-agent:testing')
33+
34+
testCompile group: 'com.rabbitmq', name: 'amqp-client', version: '2.7.0'
35+
testCompile group: 'org.springframework.amqp', name: 'spring-rabbit', version: '1.1.0.RELEASE'
36+
37+
testCompile deps.testcontainers
38+
39+
latestDepTestCompile group: 'com.rabbitmq', name: 'amqp-client', version: '+'
40+
latestDepTestCompile group: 'org.springframework.amqp', name: 'spring-rabbit', version: '+'
41+
}
42+
43+
configurations.testRuntime {
44+
resolutionStrategy {
45+
force group: 'com.rabbitmq', name: 'amqp-client', version: '2.7.0'
46+
}
47+
}

‎dd-java-agent/instrumentation/rabbitmq-amqp-2.7/src/main/java/datadog/trace/instrumentation/rabbitmq/amqp/RabbitChannelInstrumentation.java‎

Lines changed: 284 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,284 @@
1+
package datadog.trace.instrumentation.rabbitmq.amqp;
2+
3+
import static datadog.trace.agent.tooling.ByteBuddyElementMatchers.safeHasSuperType;
4+
import static io.opentracing.log.Fields.ERROR_OBJECT;
5+
import static net.bytebuddy.matcher.ElementMatchers.canThrow;
6+
import static net.bytebuddy.matcher.ElementMatchers.isGetter;
7+
import static net.bytebuddy.matcher.ElementMatchers.isInterface;
8+
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
9+
import static net.bytebuddy.matcher.ElementMatchers.isPublic;
10+
import static net.bytebuddy.matcher.ElementMatchers.isSetter;
11+
import static net.bytebuddy.matcher.ElementMatchers.nameEndsWith;
12+
import static net.bytebuddy.matcher.ElementMatchers.named;
13+
import static net.bytebuddy.matcher.ElementMatchers.not;
14+
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
15+
import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
16+
17+
import com.google.auto.service.AutoService;
18+
import com.rabbitmq.client.AMQP;
19+
import com.rabbitmq.client.Channel;
20+
import com.rabbitmq.client.Connection;
21+
import com.rabbitmq.client.Consumer;
22+
import com.rabbitmq.client.GetResponse;
23+
import com.rabbitmq.client.MessageProperties;
24+
import datadog.trace.agent.tooling.Instrumenter;
25+
import datadog.trace.api.DDSpanTypes;
26+
import datadog.trace.api.DDTags;
27+
import datadog.trace.bootstrap.CallDepthThreadLocalMap;
28+
import io.opentracing.Scope;
29+
import io.opentracing.Span;
30+
import io.opentracing.SpanContext;
31+
import io.opentracing.noop.NoopSpan;
32+
import io.opentracing.propagation.Format;
33+
import io.opentracing.tag.Tags;
34+
import io.opentracing.util.GlobalTracer;
35+
import java.io.IOException;
36+
import java.util.Collections;
37+
import java.util.HashMap;
38+
import java.util.LinkedHashMap;
39+
import java.util.Map;
40+
import java.util.concurrent.TimeUnit;
41+
import net.bytebuddy.asm.Advice;
42+
import net.bytebuddy.description.type.TypeDescription;
43+
import net.bytebuddy.matcher.ElementMatcher;
44+
45+
@AutoService(Instrumenter.class)
46+
public class RabbitChannelInstrumentation extends Instrumenter.Default {
47+
48+
public RabbitChannelInstrumentation() {
49+
super("amqp", "rabbitmq");
50+
}
51+
52+
@Override
53+
public ElementMatcher<TypeDescription> typeMatcher() {
54+
return not(isInterface()).and(safeHasSuperType(named("com.rabbitmq.client.Channel")));
55+
}
56+
57+
@Override
58+
public String[] helperClassNames() {
59+
return new String[] {
60+
packageName + ".TextMapInjectAdapter",
61+
packageName + ".TextMapExtractAdapter",
62+
packageName + ".TracedDelegatingConsumer",
63+
};
64+
}
65+
66+
@Override
67+
public Map<? extends ElementMatcher, String> transformers() {
68+
// We want the advice applied in a specific order, so use an ordered map.
69+
final Map<ElementMatcher, String> transformers = new LinkedHashMap<>();
70+
transformers.put(
71+
isMethod()
72+
.and(
73+
not(
74+
isGetter()
75+
.or(isSetter())
76+
.or(nameEndsWith("Listener"))
77+
.or(nameEndsWith("Listeners"))
78+
.or(named("processAsync"))
79+
.or(named("open"))
80+
.or(named("close"))
81+
.or(named("abort"))
82+
.or(named("basicGet"))))
83+
.and(isPublic())
84+
.and(canThrow(IOException.class).or(canThrow(InterruptedException.class))),
85+
ChannelMethodAdvice.class.getName());
86+
transformers.put(
87+
isMethod().and(named("basicPublish")).and(takesArguments(6)),
88+
ChannelPublishAdvice.class.getName());
89+
transformers.put(
90+
isMethod().and(named("basicGet")).and(takesArgument(0, String.class)),
91+
ChannelGetAdvice.class.getName());
92+
transformers.put(
93+
isMethod()
94+
.and(named("basicConsume"))
95+
.and(takesArgument(0, String.class))
96+
.and(takesArgument(6, named("com.rabbitmq.client.Consumer"))),
97+
ChannelConsumeAdvice.class.getName());
98+
return transformers;
99+
}
100+
101+
public static class ChannelMethodAdvice {
102+
@Advice.OnMethodEnter
103+
public static Scope startSpan(
104+
@Advice.This final Channel channel, @Advice.Origin("Channel.#m") final String method) {
105+
final int callDepth = CallDepthThreadLocalMap.incrementCallDepth(Channel.class);
106+
if (callDepth > 0) {
107+
return null;
108+
}
109+
110+
final Connection connection = channel.getConnection();
111+
112+
return GlobalTracer.get()
113+
.buildSpan("amqp.command")
114+
.withTag(DDTags.SERVICE_NAME, "rabbitmq")
115+
.withTag(DDTags.RESOURCE_NAME, method)
116+
.withTag(DDTags.SPAN_TYPE, DDSpanTypes.MESSAGE_CLIENT)
117+
.withTag(Tags.SPAN_KIND.getKey(), Tags.SPAN_KIND_CLIENT)
118+
.withTag(Tags.COMPONENT.getKey(), "rabbitmq-amqp")
119+
.withTag(Tags.PEER_HOSTNAME.getKey(), connection.getAddress().getHostName())
120+
.withTag(Tags.PEER_PORT.getKey(), connection.getPort())
121+
.startActive(true);
122+
}
123+
124+
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
125+
public static void stopSpan(
126+
@Advice.Enter final Scope scope, @Advice.Thrown final Throwable throwable) {
127+
if (scope != null) {
128+
if (throwable != null) {
129+
final Span span = scope.span();
130+
Tags.ERROR.set(span, true);
131+
span.log(Collections.singletonMap(ERROR_OBJECT, throwable));
132+
}
133+
scope.close();
134+
CallDepthThreadLocalMap.reset(Channel.class);
135+
}
136+
}
137+
}
138+
139+
public static class ChannelPublishAdvice {
140+
@Advice.OnMethodEnter(suppress = Throwable.class)
141+
public static void setResourceNameAddHeaders(
142+
@Advice.Argument(0) final String exchange,
143+
@Advice.Argument(1) final String routingKey,
144+
@Advice.Argument(value = 4, readOnly = false) AMQP.BasicProperties props,
145+
@Advice.Argument(5) final byte[] body) {
146+
final Span span = GlobalTracer.get().activeSpan();
147+
148+
if (span != null) {
149+
final String exchangeName = exchange == null || exchange.isEmpty() ? "<default>" : exchange;
150+
final String routing =
151+
routingKey == null || routingKey.isEmpty()
152+
? "<all>"
153+
: routingKey.startsWith("amq.gen-") ? "<generated>" : routingKey;
154+
span.setTag(DDTags.RESOURCE_NAME, "basic.publish " + exchangeName + " -> " + routing);
155+
span.setTag(DDTags.SPAN_TYPE, DDSpanTypes.MESSAGE_PRODUCER);
156+
span.setTag(Tags.SPAN_KIND.getKey(), Tags.SPAN_KIND_PRODUCER);
157+
span.setTag("amqp.exchange", exchange);
158+
span.setTag("amqp.routing_key", routingKey);
159+
span.setTag("message.size", body == null ? 0 : body.length);
160+
161+
// This is the internal behavior when props are null. We're just doing it earlier now.
162+
if (props == null) {
163+
props = MessageProperties.MINIMAL_BASIC;
164+
}
165+
span.setTag("amqp.delivery_mode", props.getDeliveryMode());
166+
167+
// We need to copy the BasicProperties and provide a header map we can modify
168+
Map<String, Object> headers = props.getHeaders();
169+
headers = (headers == null) ? new HashMap<String, Object>() : new HashMap<>(headers);
170+
171+
GlobalTracer.get()
172+
.inject(span.context(), Format.Builtin.TEXT_MAP, new TextMapInjectAdapter(headers));
173+
174+
props =
175+
new AMQP.BasicProperties(
176+
props.getContentType(),
177+
props.getContentEncoding(),
178+
headers,
179+
props.getDeliveryMode(),
180+
props.getPriority(),
181+
props.getCorrelationId(),
182+
props.getReplyTo(),
183+
props.getExpiration(),
184+
props.getMessageId(),
185+
props.getTimestamp(),
186+
props.getType(),
187+
props.getUserId(),
188+
props.getAppId(),
189+
props.getClusterId());
190+
}
191+
}
192+
}
193+
194+
public static class ChannelGetAdvice {
195+
@Advice.OnMethodEnter
196+
public static long takeTimestamp(
197+
@Advice.Local("placeholderScope") Scope scope, @Advice.Local("callDepth") int callDepth) {
198+
callDepth = CallDepthThreadLocalMap.incrementCallDepth(Channel.class);
199+
// Don't want RabbitCommandInstrumentation to mess up our actual parent span.
200+
scope = GlobalTracer.get().scopeManager().activate(NoopSpan.INSTANCE, true);
201+
return System.currentTimeMillis();
202+
}
203+
204+
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
205+
public static void extractAndStartSpan(
206+
@Advice.This final Channel channel,
207+
@Advice.Argument(0) final String queue,
208+
@Advice.Enter final long startTime,
209+
@Advice.Local("placeholderScope") final Scope scope,
210+
@Advice.Local("callDepth") final int callDepth,
211+
@Advice.Return final GetResponse response,
212+
@Advice.Thrown final Throwable throwable) {
213+
214+
if (scope.span() instanceof NoopSpan) {
215+
scope.close();
216+
}
217+
218+
if (callDepth > 0) {
219+
return;
220+
}
221+
SpanContext parentContext = null;
222+
223+
if (response != null && response.getProps() != null) {
224+
final Map<String, Object> headers = response.getProps().getHeaders();
225+
226+
parentContext =
227+
headers == null
228+
? null
229+
: GlobalTracer.get()
230+
.extract(Format.Builtin.TEXT_MAP, new TextMapExtractAdapter(headers));
231+
}
232+
233+
if (parentContext == null) {
234+
final Span parent = GlobalTracer.get().activeSpan();
235+
if (parent != null) {
236+
parentContext = parent.context();
237+
}
238+
}
239+
240+
final Connection connection = channel.getConnection();
241+
242+
final Integer length = response == null ? null : response.getBody().length;
243+
244+
final String queueName = queue.startsWith("amq.gen-") ? "<generated>" : queue;
245+
246+
final Span span =
247+
GlobalTracer.get()
248+
.buildSpan("amqp.command")
249+
.withStartTimestamp(TimeUnit.MILLISECONDS.toMicros(startTime))
250+
.asChildOf(parentContext)
251+
.withTag(DDTags.SERVICE_NAME, "rabbitmq")
252+
.withTag(DDTags.RESOURCE_NAME, "basic.get " + queueName)
253+
.withTag(DDTags.SPAN_TYPE, DDSpanTypes.MESSAGE_CONSUMER)
254+
.withTag(Tags.SPAN_KIND.getKey(), Tags.SPAN_KIND_CONSUMER)
255+
.withTag(Tags.COMPONENT.getKey(), "rabbitmq-amqp")
256+
.withTag("amqp.command", "basic.get")
257+
.withTag("amqp.queue", queue)
258+
.withTag("message.size", length)
259+
.withTag(Tags.PEER_HOSTNAME.getKey(), connection.getAddress().getHostName())
260+
.withTag(Tags.PEER_PORT.getKey(), connection.getPort())
261+
.start();
262+
263+
if (throwable != null) {
264+
Tags.ERROR.set(span, true);
265+
span.log(Collections.singletonMap(ERROR_OBJECT, throwable));
266+
}
267+
268+
span.finish();
269+
CallDepthThreadLocalMap.reset(Channel.class);
270+
}
271+
}
272+
273+
public static class ChannelConsumeAdvice {
274+
@Advice.OnMethodEnter(suppress = Throwable.class)
275+
public static void wrapConsumer(
276+
@Advice.Argument(0) final String queue,
277+
@Advice.Argument(value = 6, readOnly = false) Consumer consumer) {
278+
// We have to save off the queue name here because it isn't available to the consumer later.
279+
if (consumer != null && !(consumer instanceof TracedDelegatingConsumer)) {
280+
consumer = new TracedDelegatingConsumer(queue, consumer);
281+
}
282+
}
283+
}
284+
}

0 commit comments

Comments
 (0)