Skip to content

Commit 2ea8ef9

Browse files
lrwhlrwh
committed
Add Java-WebSocket instrumentation
1 parent b1b6910 commit 2ea8ef9

16 files changed

Lines changed: 1076 additions & 0 deletions

‎dd-java-agent/instrumentation/websocket/java-websocket/build.gradle‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
muzzle {
2+
pass {
3+
name = "Java-WebSocket"
4+
group = "org.java-websocket"
5+
module = "Java-WebSocket"
6+
versions = "[1.4.1,)"
7+
}
8+
}
9+
10+
apply from: "$rootDir/gradle/java.gradle"
11+
12+
13+
dependencies {
14+
compileOnly group: 'org.java-websocket', name: 'Java-WebSocket', version: '1.5.3'
15+
16+
testImplementation group: 'org.java-websocket', name: 'Java-WebSocket', version: '1.5.3'
17+
}

‎dd-java-agent/instrumentation/websocket/java-websocket/src/main/java/datadog/trace/instrumentation/websocket/org/OrgWebsocketModule.java‎

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
package datadog.trace.instrumentation.websocket.org;
2+
3+
import com.google.auto.service.AutoService;
4+
import datadog.trace.agent.tooling.Instrumenter;
5+
import datadog.trace.agent.tooling.InstrumenterModule;
6+
import datadog.trace.api.InstrumenterConfig;
7+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
8+
import java.util.Arrays;
9+
import java.util.HashMap;
10+
import java.util.List;
11+
import java.util.Map;
12+
13+
@AutoService(InstrumenterModule.class)
14+
public class OrgWebsocketModule extends InstrumenterModule.Tracing {
15+
16+
public OrgWebsocketModule() {
17+
this("java-websocket", "websocket");
18+
}
19+
20+
protected OrgWebsocketModule(String instrumentationName, String... additionalNames) {
21+
super(instrumentationName, additionalNames);
22+
}
23+
24+
@Override
25+
public Map<String, String> contextStore() {
26+
Map<String, String> contextStores = new HashMap<>();
27+
contextStores.put("org.java_websocket.WebSocket", WebsocketAgentSpanContext.class.getName());
28+
contextStores.put("org.java_websocket.client.WebSocketClient", AgentSpan.class.getName());
29+
return contextStores;
30+
}
31+
32+
@Override
33+
public String[] helperClassNames() {
34+
return new String[] {
35+
packageName + ".WebSocketDecorator",
36+
packageName + ".WebSocketClientDecorator",
37+
packageName + ".WebSocketServerDecorator",
38+
packageName + ".WebsocketExtractAdapter",
39+
packageName + ".TraceDraft_6455",
40+
packageName + ".WebsocketHeaderInjector",
41+
packageName + ".WebsocketHeaderExtract",
42+
packageName + ".WebsocketAgentSpanContext"
43+
};
44+
}
45+
46+
@Override
47+
protected boolean defaultEnabled() {
48+
return InstrumenterConfig.get().isWebsocketTracingEnabled();
49+
}
50+
51+
@Override
52+
public String muzzleDirective() {
53+
return "Java-WebSocket";
54+
}
55+
56+
@Override
57+
public List<Instrumenter> typeInstrumentations() {
58+
return Arrays.asList(
59+
new WebsocketClientInstrumentation(),
60+
new WebsocketServerInstrumentation(),
61+
new WebSocketSendInstrumentation());
62+
}
63+
}

‎dd-java-agent/instrumentation/websocket/java-websocket/src/main/java/datadog/trace/instrumentation/websocket/org/TraceDraft_6455.java‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
package datadog.trace.instrumentation.websocket.org;
2+
3+
import java.util.ArrayList;
4+
import java.util.List;
5+
import org.java_websocket.drafts.Draft;
6+
import org.java_websocket.drafts.Draft_6455;
7+
import org.java_websocket.extensions.IExtension;
8+
import org.java_websocket.protocols.IProtocol;
9+
10+
public class TraceDraft_6455 extends Draft_6455 {
11+
12+
public TraceDraft_6455() {
13+
super();
14+
}
15+
16+
public TraceDraft_6455(List<IExtension> inputExtensions) {
17+
super(inputExtensions);
18+
}
19+
20+
public TraceDraft_6455(List<IExtension> inputExtensions, List<IProtocol> inputProtocols) {
21+
super(inputExtensions, inputProtocols);
22+
}
23+
24+
public TraceDraft_6455(
25+
List<IExtension> inputExtensions, List<IProtocol> inputProtocols, int maxFrameSize) {
26+
super(inputExtensions, inputProtocols, maxFrameSize);
27+
}
28+
29+
public static TraceDraft_6455 fromDraft6455(Draft_6455 original) {
30+
List<IExtension> extensions = new ArrayList<>();
31+
for (IExtension ext : original.getKnownExtensions()) {
32+
extensions.add(ext);
33+
}
34+
35+
List<IProtocol> protocols = new ArrayList<>();
36+
for (IProtocol proto : original.getKnownProtocols()) {
37+
protocols.add(proto);
38+
}
39+
40+
return new TraceDraft_6455(extensions, protocols, original.getMaxFrameSize());
41+
}
42+
43+
@Override
44+
public Draft copyInstance() {
45+
ArrayList<IExtension> newExtensions = new ArrayList<>();
46+
for (IExtension knownExtension : this.getKnownExtensions()) {
47+
newExtensions.add(knownExtension.copyInstance());
48+
}
49+
50+
ArrayList<IProtocol> newProtocols = new ArrayList<>();
51+
52+
for (IProtocol knownProtocol : this.getKnownProtocols()) {
53+
newProtocols.add(knownProtocol.copyInstance());
54+
}
55+
56+
return new TraceDraft_6455(newExtensions, newProtocols, getMaxFrameSize());
57+
}
58+
59+
@Override
60+
public boolean equals(Object o) {
61+
return this == o || super.equals(o);
62+
}
63+
64+
@Override
65+
public int hashCode() {
66+
return super.hashCode();
67+
}
68+
}

‎dd-java-agent/instrumentation/websocket/java-websocket/src/main/java/datadog/trace/instrumentation/websocket/org/WebSocketClientDecorator.java‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
package datadog.trace.instrumentation.websocket.org;
2+
3+
import static datadog.context.propagation.Propagators.defaultPropagator;
4+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
5+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
6+
import static datadog.trace.bootstrap.instrumentation.api.Tags.SPAN_KIND_CLIENT;
7+
import static datadog.trace.instrumentation.websocket.org.WebsocketHeaderInjector.SETTER;
8+
9+
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
10+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
11+
import datadog.trace.bootstrap.instrumentation.api.Tags;
12+
import org.java_websocket.client.WebSocketClient;
13+
14+
public class WebSocketClientDecorator extends WebSocketDecorator {
15+
public static final WebSocketClientDecorator CLIENT_DECORATE = new WebSocketClientDecorator();
16+
17+
@Override
18+
protected CharSequence spanKind() {
19+
return SPAN_KIND_CLIENT;
20+
}
21+
22+
private static final String OPERATION_NAME = "websocket.handshake";
23+
24+
public AgentScope startHandshakeSpan(WebSocketClient client) {
25+
String uri = client.getURI().toString();
26+
AgentSpan span = startSpan("websocket.handshake", OPERATION_NAME);
27+
span.setTag(Tags.HTTP_URL, uri);
28+
span.setTag(Tags.HTTP_METHOD, "GET");
29+
AgentScope scope = activateSpan(span);
30+
defaultPropagator().inject(scope.context(), client, SETTER);
31+
afterStart(span);
32+
return scope;
33+
}
34+
35+
public void onHandshakeSuccess(AgentSpan span, int httpStatus) {
36+
span.setTag(Tags.HTTP_STATUS, httpStatus);
37+
span.setTag("websocket.handshake.success", true);
38+
}
39+
}

‎dd-java-agent/instrumentation/websocket/java-websocket/src/main/java/datadog/trace/instrumentation/websocket/org/WebSocketDecorator.java‎

Lines changed: 107 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,107 @@
1+
package datadog.trace.instrumentation.websocket.org;
2+
3+
import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.extractContextAndGetSpanContext;
4+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
5+
import static datadog.trace.bootstrap.instrumentation.websocket.HandlersExtractor.MESSAGE_TYPE_BINARY;
6+
import static datadog.trace.bootstrap.instrumentation.websocket.HandlersExtractor.MESSAGE_TYPE_TEXT;
7+
import static datadog.trace.instrumentation.websocket.org.WebsocketExtractAdapter.GETTER;
8+
9+
import datadog.trace.api.DDSpanTypes;
10+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
11+
import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
12+
import datadog.trace.bootstrap.instrumentation.api.Tags;
13+
import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
14+
import datadog.trace.bootstrap.instrumentation.decorator.BaseDecorator;
15+
import java.net.InetSocketAddress;
16+
import java.nio.ByteBuffer;
17+
import org.java_websocket.WebSocket;
18+
import org.java_websocket.handshake.Handshakedata;
19+
20+
public abstract class WebSocketDecorator extends BaseDecorator {
21+
public static final CharSequence WEBSOCKET = UTF8BytesString.create("org-java-websocket");
22+
public static final CharSequence WEBSOCKET_OPEN = UTF8BytesString.create("websocket.open");
23+
public static final CharSequence WEBSOCKET_RECEIVE = UTF8BytesString.create("websocket.receive");
24+
public static final CharSequence WEBSOCKET_SEND = UTF8BytesString.create("websocket.send");
25+
public static final CharSequence WEBSOCKET_CLOSE = UTF8BytesString.create("websocket.close");
26+
private static final String[] INSTRUMENTATION_NAMES = {WEBSOCKET.toString()};
27+
private static final String REMOTE_ADDRESS = "remote.address";
28+
private static final String MESSAGE_TYPE = "message.type";
29+
private static final String MESSAGE_SIZE = "message.size";
30+
private static final String MESSAGE_REMOTE = "message.remote";
31+
private static final String MESSAGE_CODE = "message.code";
32+
private static final String MESSAGE_REASON = "message.reason";
33+
34+
@Override
35+
protected String[] instrumentationNames() {
36+
return INSTRUMENTATION_NAMES;
37+
}
38+
39+
@Override
40+
protected CharSequence spanType() {
41+
return DDSpanTypes.WEBSOCKET;
42+
}
43+
44+
@Override
45+
protected CharSequence component() {
46+
return WEBSOCKET;
47+
}
48+
49+
protected abstract CharSequence spanKind();
50+
51+
public AgentSpan open(WebSocket conn, Handshakedata handshake) {
52+
AgentSpanContext parentContext = extractContextAndGetSpanContext(handshake, GETTER);
53+
AgentSpan span = startSpan(instrumentationNames()[0], WEBSOCKET_OPEN, parentContext);
54+
InetSocketAddress remoteAddress = conn.getRemoteSocketAddress();
55+
if (remoteAddress != null) {
56+
span.setTag(REMOTE_ADDRESS, remoteAddress.toString());
57+
}
58+
afterStart(span);
59+
return span;
60+
}
61+
62+
public AgentSpan onMessage(Object message, AgentSpanContext spanContext) {
63+
AgentSpan span = startSpan(instrumentationNames()[0], WEBSOCKET_RECEIVE, spanContext);
64+
CharSequence messageType;
65+
int msgSize;
66+
if (message instanceof ByteBuffer) {
67+
messageType = MESSAGE_TYPE_BINARY;
68+
msgSize = ((ByteBuffer) message).remaining();
69+
} else {
70+
messageType = MESSAGE_TYPE_TEXT;
71+
msgSize = message == null ? 0 : ((CharSequence) message).length();
72+
}
73+
span.setTag(MESSAGE_SIZE, msgSize);
74+
span.setTag(MESSAGE_TYPE, messageType);
75+
afterStart(span);
76+
return span;
77+
}
78+
79+
public AgentSpan send(WebsocketAgentSpanContext spanContext) {
80+
AgentSpan span;
81+
if (spanContext == null) {
82+
span = startSpan(instrumentationNames()[0], WEBSOCKET_SEND);
83+
} else {
84+
span = startSpan(instrumentationNames()[0], WEBSOCKET_SEND, spanContext.getSpanContext());
85+
}
86+
afterStart(span);
87+
if (spanContext != null) {
88+
span.setTag(Tags.SPAN_KIND, spanContext.getSpanKind());
89+
}
90+
return span;
91+
}
92+
93+
public AgentSpan onClose(int code, String reason, boolean remote, AgentSpanContext spanContext) {
94+
AgentSpan span = startSpan(instrumentationNames()[0], WEBSOCKET_CLOSE, spanContext);
95+
span.setTag(MESSAGE_REASON, reason);
96+
span.setTag(MESSAGE_CODE, code);
97+
span.setTag(MESSAGE_REMOTE, remote);
98+
afterStart(span);
99+
return span;
100+
}
101+
102+
@Override
103+
public AgentSpan afterStart(AgentSpan span) {
104+
span.setTag(Tags.SPAN_KIND, spanKind());
105+
return super.afterStart(span);
106+
}
107+
}

‎dd-java-agent/instrumentation/websocket/java-websocket/src/main/java/datadog/trace/instrumentation/websocket/org/WebSocketSendInstrumentation.java‎

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,81 @@
1+
package datadog.trace.instrumentation.websocket.org;
2+
3+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
4+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
5+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopSpan;
6+
import static datadog.trace.instrumentation.websocket.org.WebSocketServerDecorator.SERVER_DECORATOR;
7+
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
8+
import static net.bytebuddy.matcher.ElementMatchers.isPrivate;
9+
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
10+
import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
11+
12+
import datadog.trace.agent.tooling.Instrumenter;
13+
import datadog.trace.bootstrap.InstrumentationContext;
14+
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
15+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
16+
import java.util.Collection;
17+
import net.bytebuddy.asm.Advice;
18+
import org.java_websocket.WebSocket;
19+
import org.java_websocket.WebSocketImpl;
20+
import org.java_websocket.framing.CloseFrame;
21+
import org.java_websocket.framing.Framedata;
22+
import org.java_websocket.framing.PingFrame;
23+
import org.java_websocket.framing.PongFrame;
24+
25+
public class WebSocketSendInstrumentation
26+
implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
27+
28+
public WebSocketSendInstrumentation() {}
29+
30+
@Override
31+
public String instrumentedType() {
32+
return "org.java_websocket.WebSocketImpl";
33+
}
34+
35+
@Override
36+
public void methodAdvice(MethodTransformer transformer) {
37+
transformer.applyAdvice(
38+
isMethod()
39+
.and(isPrivate())
40+
.and(named("send"))
41+
.and(takesArguments(1))
42+
.and(takesArgument(0, Collection.class)),
43+
getClass().getName() + "$OnSendAdvice");
44+
}
45+
46+
public static class OnSendAdvice {
47+
@Advice.OnMethodEnter(suppress = Throwable.class)
48+
public static AgentScope onEnter(
49+
@Advice.This WebSocketImpl impl, @Advice.Argument(value = 0) Collection<Framedata> frames) {
50+
WebsocketAgentSpanContext context =
51+
InstrumentationContext.get(WebSocket.class, WebsocketAgentSpanContext.class).get(impl);
52+
if (context == null) {
53+
// close after send
54+
final AgentSpan wsSpan = SERVER_DECORATOR.send(null);
55+
return activateSpan(wsSpan);
56+
}
57+
// ignore ping/pong/close
58+
for (Framedata frame : frames) {
59+
if (frame instanceof PingFrame
60+
|| frame instanceof PongFrame
61+
|| frame instanceof CloseFrame) {
62+
return activateSpan(noopSpan());
63+
}
64+
}
65+
final AgentSpan wsSpan = SERVER_DECORATOR.send(context);
66+
return activateSpan(wsSpan);
67+
}
68+
69+
@Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
70+
public static void onExit(
71+
@Advice.Enter final AgentScope scope, @Advice.Thrown final Throwable throwable) {
72+
if (scope == null) {
73+
return;
74+
}
75+
SERVER_DECORATOR.onError(scope.span(), throwable);
76+
SERVER_DECORATOR.beforeFinish(scope.span());
77+
scope.close();
78+
scope.span().finish();
79+
}
80+
}
81+
}

‎dd-java-agent/instrumentation/websocket/java-websocket/src/main/java/datadog/trace/instrumentation/websocket/org/WebSocketServerDecorator.java‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
package datadog.trace.instrumentation.websocket.org;
2+
3+
import static datadog.trace.bootstrap.instrumentation.api.Tags.SPAN_KIND_SERVER;
4+
5+
public class WebSocketServerDecorator extends WebSocketDecorator {
6+
public static final WebSocketServerDecorator SERVER_DECORATOR = new WebSocketServerDecorator();
7+
8+
@Override
9+
protected CharSequence spanKind() {
10+
return SPAN_KIND_SERVER;
11+
}
12+
}

0 commit comments

Comments
 (0)