Skip to content

Commit 2d96885

Browse files
amarzialiamarziali
committed
Migrate netty server
1 parent 1a7c76c commit 2d96885

8 files changed

Lines changed: 124 additions & 6 deletions

File tree

‎dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/AttributeKeys.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,9 @@ public final class AttributeKeys {
2727
CONNECT_PARENT_CONTINUATION_ATTRIBUTE_KEY =
2828
attributeKey("datadog.connect.parent.continuation");
2929

30+
public static final AttributeKey<Context> PARENT_CONTEXT_ATTRIBUTE_KEY =
31+
attributeKey("datadog.server.parent-context");
32+
3033
public static final AttributeKey<HttpHeaders> REQUEST_HEADERS_ATTRIBUTE_KEY =
3134
attributeKey("datadog.server.request.headers");
3235

‎dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/NettyChannelPipelineInstrumentation.java‎

Lines changed: 30 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package datadog.trace.instrumentation.netty40;
22

3+
import static datadog.trace.agent.tooling.InstrumenterModule.TargetSystem.CONTEXT_TRACKING;
34
import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface;
45
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
56
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.namedOneOf;
@@ -13,12 +14,14 @@
1314
import com.google.auto.service.AutoService;
1415
import datadog.trace.agent.tooling.Instrumenter;
1516
import datadog.trace.agent.tooling.InstrumenterModule;
17+
import datadog.trace.agent.tooling.annotation.AppliesOn;
1618
import datadog.trace.api.InstrumenterConfig;
1719
import datadog.trace.bootstrap.CallDepthThreadLocalMap;
1820
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
1921
import datadog.trace.instrumentation.netty40.client.HttpClientRequestTracingHandler;
2022
import datadog.trace.instrumentation.netty40.client.HttpClientResponseTracingHandler;
2123
import datadog.trace.instrumentation.netty40.client.HttpClientTracingHandler;
24+
import datadog.trace.instrumentation.netty40.server.HttpServerContextTrackingHandler;
2225
import datadog.trace.instrumentation.netty40.server.HttpServerRequestTracingHandler;
2326
import datadog.trace.instrumentation.netty40.server.HttpServerResponseTracingHandler;
2427
import datadog.trace.instrumentation.netty40.server.HttpServerTracingHandler;
@@ -77,6 +80,7 @@ public String[] helperClassNames() {
7780
packageName + ".server.NettyHttpServerDecorator$NettyBlockResponseFunction",
7881
packageName + ".server.BlockingResponseHandler",
7982
packageName + ".server.BlockingResponseHandler$IgnoreAllWritesHandler",
83+
packageName + ".server.HttpServerContextTrackingHandler",
8084
packageName + ".server.HttpServerRequestTracingHandler",
8185
packageName + ".server.HttpServerResponseTracingHandler",
8286
packageName + ".server.HttpServerTracingHandler",
@@ -90,21 +94,45 @@ public String[] helperClassNames() {
9094

9195
@Override
9296
public void methodAdvice(MethodTransformer transformer) {
93-
transformer.applyAdvice(
97+
transformer.applyAdvices(
9498
isMethod()
9599
.and(namedOneOf("addFirst", "addLast"))
96100
.and(takesArgument(2, named("io.netty.channel.ChannelHandler"))),
101+
NettyChannelPipelineInstrumentation.class.getName() + "$ContextTrackingAddHandlerAdvice",
97102
NettyChannelPipelineInstrumentation.class.getName() + "$AddHandlerAdvice");
98-
transformer.applyAdvice(
103+
transformer.applyAdvices(
99104
isMethod()
100105
.and(namedOneOf("addBefore", "addAfter"))
101106
.and(takesArgument(3, named("io.netty.channel.ChannelHandler"))),
107+
NettyChannelPipelineInstrumentation.class.getName() + "$ContextTrackingAddHandlerAdvice",
102108
NettyChannelPipelineInstrumentation.class.getName() + "$AddHandlerAdvice");
103109
transformer.applyAdvice(
104110
isMethod().and(named("connect")).and(returns(named("io.netty.channel.ChannelFuture"))),
105111
NettyChannelPipelineInstrumentation.class.getName() + "$ConnectAdvice");
106112
}
107113

114+
@AppliesOn(CONTEXT_TRACKING)
115+
public static class ContextTrackingAddHandlerAdvice {
116+
// No OnMethodEnter — avoids double-incrementing CallDepthThreadLocalMap,
117+
// which would cause AddHandlerAdvice.OnMethodExit to see depth > 0 and skip.
118+
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
119+
public static void addContextTrackingHandler(
120+
@Advice.This final ChannelPipeline pipeline,
121+
@Advice.Argument(value = 2, optional = true) final Object handler2,
122+
@Advice.Argument(value = 3, optional = true) final ChannelHandler handler3) {
123+
ChannelHandler handler =
124+
handler2 instanceof ChannelHandler ? (ChannelHandler) handler2 : handler3;
125+
try {
126+
if (handler instanceof HttpServerCodec || handler instanceof HttpRequestDecoder) {
127+
NettyPipelineHelper.addHandlerAfter(
128+
pipeline, handler, HttpServerContextTrackingHandler.INSTANCE);
129+
}
130+
} catch (final IllegalArgumentException e) {
131+
// Prevented adding duplicate handlers.
132+
}
133+
}
134+
}
135+
108136
/**
109137
* When certain handlers are added to the pipeline, we want to add our corresponding tracing
110138
* handlers. If those handlers are later removed, we may want to remove our handlers. That is not

‎dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/HttpServerContextTrackingHandler.java‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
package datadog.trace.instrumentation.netty40.server;
2+
3+
import static datadog.trace.instrumentation.netty40.AttributeKeys.PARENT_CONTEXT_ATTRIBUTE_KEY;
4+
import static datadog.trace.instrumentation.netty40.server.NettyHttpServerDecorator.DECORATE;
5+
6+
import datadog.context.Context;
7+
import io.netty.channel.ChannelHandler;
8+
import io.netty.channel.ChannelHandlerContext;
9+
import io.netty.channel.ChannelInboundHandlerAdapter;
10+
import io.netty.handler.codec.http.HttpRequest;
11+
12+
@ChannelHandler.Sharable
13+
public class HttpServerContextTrackingHandler extends ChannelInboundHandlerAdapter {
14+
public static final HttpServerContextTrackingHandler INSTANCE =
15+
new HttpServerContextTrackingHandler();
16+
17+
@Override
18+
public void channelRead(final ChannelHandlerContext ctx, final Object msg) {
19+
if (msg instanceof HttpRequest) {
20+
final Context parentContext = DECORATE.extract(((HttpRequest) msg).headers());
21+
ctx.channel().attr(PARENT_CONTEXT_ATTRIBUTE_KEY).set(parentContext);
22+
}
23+
ctx.fireChannelRead(msg);
24+
}
25+
}

‎dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/HttpServerRequestTracingHandler.java‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import static datadog.trace.instrumentation.netty40.AttributeKeys.ANALYZED_RESPONSE_KEY;
55
import static datadog.trace.instrumentation.netty40.AttributeKeys.BLOCKED_RESPONSE_KEY;
66
import static datadog.trace.instrumentation.netty40.AttributeKeys.CONTEXT_ATTRIBUTE_KEY;
7+
import static datadog.trace.instrumentation.netty40.AttributeKeys.PARENT_CONTEXT_ATTRIBUTE_KEY;
78
import static datadog.trace.instrumentation.netty40.AttributeKeys.REQUEST_HEADERS_ATTRIBUTE_KEY;
89
import static datadog.trace.instrumentation.netty40.server.NettyHttpServerDecorator.DECORATE;
910

@@ -40,7 +41,9 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) {
4041

4142
final HttpRequest request = (HttpRequest) msg;
4243
final HttpHeaders headers = request.headers();
43-
final Context parentContext = DECORATE.extract(headers);
44+
final Context storedParentContext = channel.attr(PARENT_CONTEXT_ATTRIBUTE_KEY).getAndRemove();
45+
final Context parentContext =
46+
storedParentContext != null ? storedParentContext : DECORATE.extract(headers);
4447
final Context context = DECORATE.startSpan(headers, parentContext);
4548

4649
try (final ContextScope ignored = context.attach()) {

‎dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/NettyChannelPipelineInstrumentation.java‎

Lines changed: 30 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package datadog.trace.instrumentation.netty41;
22

3+
import static datadog.trace.agent.tooling.InstrumenterModule.TargetSystem.CONTEXT_TRACKING;
34
import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface;
45
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
56
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.namedOneOf;
@@ -13,12 +14,14 @@
1314
import com.google.auto.service.AutoService;
1415
import datadog.trace.agent.tooling.Instrumenter;
1516
import datadog.trace.agent.tooling.InstrumenterModule;
17+
import datadog.trace.agent.tooling.annotation.AppliesOn;
1618
import datadog.trace.api.InstrumenterConfig;
1719
import datadog.trace.bootstrap.CallDepthThreadLocalMap;
1820
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
1921
import datadog.trace.instrumentation.netty41.client.HttpClientRequestTracingHandler;
2022
import datadog.trace.instrumentation.netty41.client.HttpClientResponseTracingHandler;
2123
import datadog.trace.instrumentation.netty41.client.HttpClientTracingHandler;
24+
import datadog.trace.instrumentation.netty41.server.HttpServerContextTrackingHandler;
2225
import datadog.trace.instrumentation.netty41.server.HttpServerRequestTracingHandler;
2326
import datadog.trace.instrumentation.netty41.server.HttpServerResponseTracingHandler;
2427
import datadog.trace.instrumentation.netty41.server.HttpServerTracingHandler;
@@ -79,6 +82,7 @@ public String[] helperClassNames() {
7982
packageName + ".server.NettyHttpServerDecorator$NettyBlockResponseFunction",
8083
packageName + ".server.BlockingResponseHandler",
8184
packageName + ".server.BlockingResponseHandler$IgnoreAllWritesHandler",
85+
packageName + ".server.HttpServerContextTrackingHandler",
8286
packageName + ".server.HttpServerRequestTracingHandler",
8387
packageName + ".server.HttpServerResponseTracingHandler",
8488
packageName + ".server.HttpServerTracingHandler",
@@ -93,21 +97,45 @@ public String[] helperClassNames() {
9397

9498
@Override
9599
public void methodAdvice(MethodTransformer transformer) {
96-
transformer.applyAdvice(
100+
transformer.applyAdvices(
97101
isMethod()
98102
.and(namedOneOf("addFirst", "addLast"))
99103
.and(takesArgument(2, named("io.netty.channel.ChannelHandler"))),
104+
NettyChannelPipelineInstrumentation.class.getName() + "$ContextTrackingAddHandlerAdvice",
100105
NettyChannelPipelineInstrumentation.class.getName() + "$AddHandlerAdvice");
101-
transformer.applyAdvice(
106+
transformer.applyAdvices(
102107
isMethod()
103108
.and(namedOneOf("addBefore", "addAfter"))
104109
.and(takesArgument(3, named("io.netty.channel.ChannelHandler"))),
110+
NettyChannelPipelineInstrumentation.class.getName() + "$ContextTrackingAddHandlerAdvice",
105111
NettyChannelPipelineInstrumentation.class.getName() + "$AddHandlerAdvice");
106112
transformer.applyAdvice(
107113
isMethod().and(named("connect")).and(returns(named("io.netty.channel.ChannelFuture"))),
108114
NettyChannelPipelineInstrumentation.class.getName() + "$ConnectAdvice");
109115
}
110116

117+
@AppliesOn(CONTEXT_TRACKING)
118+
public static class ContextTrackingAddHandlerAdvice {
119+
// No OnMethodEnter — avoids double-incrementing CallDepthThreadLocalMap,
120+
// which would cause AddHandlerAdvice.OnMethodExit to see depth > 0 and skip.
121+
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
122+
public static void addContextTrackingHandler(
123+
@Advice.This final ChannelPipeline pipeline,
124+
@Advice.Argument(value = 2, optional = true) final Object handler2,
125+
@Advice.Argument(value = 3, optional = true) final ChannelHandler handler3) {
126+
ChannelHandler handler =
127+
handler2 instanceof ChannelHandler ? (ChannelHandler) handler2 : handler3;
128+
try {
129+
if (handler instanceof HttpServerCodec || handler instanceof HttpRequestDecoder) {
130+
NettyPipelineHelper.addHandlerAfter(
131+
pipeline, handler, HttpServerContextTrackingHandler.INSTANCE);
132+
}
133+
} catch (final IllegalArgumentException e) {
134+
// Prevented adding duplicate handlers.
135+
}
136+
}
137+
}
138+
111139
/**
112140
* When certain handlers are added to the pipeline, we want to add our corresponding tracing
113141
* handlers. If those handlers are later removed, we may want to remove our handlers. That is not

‎dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/HttpServerContextTrackingHandler.java‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
package datadog.trace.instrumentation.netty41.server;
2+
3+
import static datadog.trace.instrumentation.netty41.AttributeKeys.PARENT_CONTEXT_ATTRIBUTE_KEY;
4+
import static datadog.trace.instrumentation.netty41.server.NettyHttpServerDecorator.DECORATE;
5+
6+
import datadog.context.Context;
7+
import io.netty.channel.ChannelHandler;
8+
import io.netty.channel.ChannelHandlerContext;
9+
import io.netty.channel.ChannelInboundHandlerAdapter;
10+
import io.netty.handler.codec.http.HttpRequest;
11+
12+
@ChannelHandler.Sharable
13+
public class HttpServerContextTrackingHandler extends ChannelInboundHandlerAdapter {
14+
public static final HttpServerContextTrackingHandler INSTANCE =
15+
new HttpServerContextTrackingHandler();
16+
17+
@Override
18+
public void channelRead(final ChannelHandlerContext ctx, final Object msg) {
19+
if (msg instanceof HttpRequest) {
20+
final Context parentContext = DECORATE.extract(((HttpRequest) msg).headers());
21+
ctx.channel().attr(PARENT_CONTEXT_ATTRIBUTE_KEY).set(parentContext);
22+
}
23+
ctx.fireChannelRead(msg);
24+
}
25+
}

‎dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/HttpServerRequestTracingHandler.java‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import static datadog.trace.instrumentation.netty41.AttributeKeys.ANALYZED_RESPONSE_KEY;
55
import static datadog.trace.instrumentation.netty41.AttributeKeys.BLOCKED_RESPONSE_KEY;
66
import static datadog.trace.instrumentation.netty41.AttributeKeys.CONTEXT_ATTRIBUTE_KEY;
7+
import static datadog.trace.instrumentation.netty41.AttributeKeys.PARENT_CONTEXT_ATTRIBUTE_KEY;
78
import static datadog.trace.instrumentation.netty41.AttributeKeys.REQUEST_HEADERS_ATTRIBUTE_KEY;
89
import static datadog.trace.instrumentation.netty41.server.NettyHttpServerDecorator.DECORATE;
910

@@ -39,7 +40,9 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) {
3940

4041
final HttpRequest request = (HttpRequest) msg;
4142
final HttpHeaders headers = request.headers();
42-
final Context parentContext = DECORATE.extract(headers);
43+
final Context storedParentContext = channel.attr(PARENT_CONTEXT_ATTRIBUTE_KEY).getAndRemove();
44+
final Context parentContext =
45+
storedParentContext != null ? storedParentContext : DECORATE.extract(headers);
4346
final Context context = DECORATE.startSpan(headers, parentContext);
4447

4548
try (final ContextScope ignored = context.attach()) {

‎dd-java-agent/instrumentation/netty/netty-common/src/main/java/datadog/trace/instrumentation/netty41/AttributeKeys.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,9 @@ public final class AttributeKeys {
2727
CONNECT_PARENT_CONTINUATION_ATTRIBUTE_KEY =
2828
attributeKey("datadog.connect.parent.continuation");
2929

30+
public static final AttributeKey<Context> PARENT_CONTEXT_ATTRIBUTE_KEY =
31+
attributeKey("datadog.server.parent-context");
32+
3033
public static final AttributeKey<HttpHeaders> REQUEST_HEADERS_ATTRIBUTE_KEY =
3134
attributeKey("datadog.server.request.headers");
3235

0 commit comments

Comments
 (0)