44import datadog .trace .api .gateway .Flow ;
55import datadog .trace .api .internal .TraceSegment ;
66import datadog .trace .bootstrap .blocking .BlockingActionHelper ;
7+ import datadog .trace .instrumentation .netty41 .ServerRequestContext ;
78import io .netty .channel .ChannelHandler ;
89import io .netty .channel .ChannelHandlerContext ;
910import io .netty .channel .ChannelInboundHandlerAdapter ;
1516import io .netty .handler .codec .http .HttpRequest ;
1617import io .netty .handler .codec .http .HttpResponseStatus ;
1718import io .netty .handler .codec .http .HttpUtil ;
19+ import io .netty .handler .codec .http .HttpVersion ;
1820import io .netty .util .ReferenceCountUtil ;
1921import java .util .Map ;
2022import org .slf4j .Logger ;
2123import org .slf4j .LoggerFactory ;
2224
2325public class BlockingResponseHandler extends ChannelInboundHandlerAdapter {
2426 public static final Logger log = LoggerFactory .getLogger (BlockingResponseHandler .class );
27+ static final String BEFORE_BLOCKING_HANDLER_NAME = "before_blocking_handler" ;
28+ static final String HANDLER_NAME = "blocking_handler" ;
29+ private static final String IGNORE_ALL_WRITES_HANDLER = "ignore_all_writes_handler" ;
30+ private static final String MISSING_RESPONSE_TRACING_HANDLER_MESSAGE =
31+ "Unable to block because HttpServerResponseTracingHandler was not found on the pipeline" ;
2532 private static volatile boolean HAS_WARNED ;
2633
2734 private final TraceSegment segment ;
2835 private final int statusCode ;
2936 private final BlockingContentType bct ;
3037 private final Map <String , String > extraHeaders ;
3138 private final String securityResponseId ;
39+ private final ServerRequestContext serverContext ;
3240
33- private boolean hasBlockedAlready ;
41+ // Current callers are event-loop confined; volatile preserves visibility if
42+ // commitBlockingResponse is invoked from another thread.
43+ private volatile boolean hasBlockedAlready ;
3444
3545 public BlockingResponseHandler (
3646 TraceSegment segment ,
3747 int statusCode ,
3848 BlockingContentType bct ,
3949 Map <String , String > extraHeaders ,
40- String securityResponseId ) {
50+ String securityResponseId ,
51+ ServerRequestContext serverContext ) {
4152 this .segment = segment ;
4253 this .statusCode = statusCode ;
4354 this .bct = bct ;
4455 this .extraHeaders = extraHeaders ;
4556 this .securityResponseId = securityResponseId ;
57+ this .serverContext = serverContext ;
4658 }
4759
48- public BlockingResponseHandler (TraceSegment segment , Flow .Action .RequestBlockingAction rba ) {
49- this (
50- segment ,
51- rba .getStatusCode (),
52- rba .getBlockingContentType (),
53- rba .getExtraHeaders (),
54- rba .getSecurityResponseId ());
60+ public BlockingResponseHandler (
61+ TraceSegment segment ,
62+ Flow .Action .RequestBlockingAction rba ,
63+ ServerRequestContext serverContext ) {
64+ this .segment = segment ;
65+ this .statusCode = rba .getStatusCode ();
66+ this .bct = rba .getBlockingContentType ();
67+ this .extraHeaders = rba .getExtraHeaders ();
68+ this .securityResponseId = rba .getSecurityResponseId ();
69+ this .serverContext = serverContext ;
5570 }
5671
5772 @ Override
@@ -66,82 +81,162 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) {
6681 return ;
6782 }
6883
84+ HttpRequest request = (HttpRequest ) msg ;
85+ if (!commitBlockingResponse (ctx , request .protocolVersion (), request .headers ().get ("accept" ))) {
86+ // Do not let a failed block intercept later requests on this keep-alive connection.
87+ if (ctx .pipeline ().get (BEFORE_BLOCKING_HANDLER_NAME ) != null ) {
88+ ctx .pipeline ().remove (BEFORE_BLOCKING_HANDLER_NAME );
89+ }
90+ ctx .pipeline ().remove (this );
91+ ctx .fireChannelRead (msg );
92+ return ;
93+ }
94+
95+ ReferenceCountUtil .release (msg );
96+ }
97+
98+ boolean commitBlockingResponse (
99+ final ChannelHandlerContext ctx ,
100+ final HttpVersion protocolVersion ,
101+ final String acceptHeader ) {
69102 ChannelHandlerContext ctxForDownstream =
70103 ctx .pipeline ().context (HttpServerResponseTracingHandler .class );
71104 if (ctxForDownstream == null ) {
72105 ctxForDownstream = ctx .pipeline ().context (HttpServerTracingHandler .class );
73106 }
74-
75107 if (ctxForDownstream == null ) {
76- if (HAS_WARNED ) {
77- log .debug (
78- "Unable to block because HttpServerResponseTracingHandler was not found on the pipeline" );
79- } else {
80- log .warn (
81- "Unable to block because HttpServerResponseTracingHandler was not found on the pipeline" );
82- HAS_WARNED = true ;
83- }
84- ctx .fireChannelRead (msg );
85- return ;
108+ logMissingResponseTracingHandler ();
109+ return false ;
86110 }
87-
88- HttpRequest request = (HttpRequest ) msg ;
89-
90- int httpCode = BlockingActionHelper .getHttpCode (statusCode );
91- HttpResponseStatus httpResponseStatus = HttpResponseStatus .valueOf (httpCode );
92- FullHttpResponse response =
93- new DefaultFullHttpResponse (request .protocolVersion (), httpResponseStatus );
94-
95- HttpHeaders headers = response .headers ();
96- headers .set ("Connection" , "close" );
97-
98- for (Map .Entry <String , String > h : this .extraHeaders .entrySet ()) {
99- headers .set (h .getKey (), h .getValue ());
111+ this .hasBlockedAlready = true ;
112+ ServerRequestContext .markRequestBlocked (ctx .channel ());
113+
114+ PendingBlockResponse pendingBlockResponse =
115+ new PendingBlockResponse (
116+ segment ,
117+ statusCode ,
118+ bct ,
119+ extraHeaders ,
120+ securityResponseId ,
121+ protocolVersion ,
122+ acceptHeader );
123+
124+ if (serverContext != null
125+ && ServerRequestContext .nextResponse (ctx .channel ()) != serverContext ) {
126+ serverContext .deferBlockResponse (pendingBlockResponse );
127+ return true ;
100128 }
101129
102- if (bct != BlockingContentType .NONE ) {
103- String acceptHeader = request .headers ().get ("accept" );
104- BlockingActionHelper .TemplateType type =
105- BlockingActionHelper .determineTemplateType (bct , acceptHeader );
106- headers .set ("Content-type" , BlockingActionHelper .getContentType (type ));
130+ writeBlockResponse (ctxForDownstream , pendingBlockResponse );
131+ return true ;
132+ }
107133
108- byte [] template = BlockingActionHelper .getTemplate (type , this .securityResponseId );
109- HttpUtil .setContentLength (response , template .length );
110- response .content ().writeBytes (template );
134+ private static void logMissingResponseTracingHandler () {
135+ if (HAS_WARNED ) {
136+ log .debug (MISSING_RESPONSE_TRACING_HANDLER_MESSAGE );
137+ } else {
138+ log .warn (MISSING_RESPONSE_TRACING_HANDLER_MESSAGE );
139+ HAS_WARNED = true ;
111140 }
141+ }
112142
113- this .hasBlockedAlready = true ;
114-
115- ReferenceCountUtil .release (msg );
143+ static boolean maybeWriteDeferredBlockResponse (
144+ ChannelHandlerContext ctx , ServerRequestContext serverContext ) {
145+ if (serverContext == null ) {
146+ return false ;
147+ }
148+ Object deferredBlockResponse = serverContext .deferredBlockResponse ();
149+ if (!(deferredBlockResponse instanceof PendingBlockResponse )) {
150+ return false ;
151+ }
152+ serverContext .deferBlockResponse (null );
153+ writeBlockResponse (ctx , (PendingBlockResponse ) deferredBlockResponse );
154+ return true ;
155+ }
116156
157+ private static void writeBlockResponse (
158+ ChannelHandlerContext ctxForDownstream , PendingBlockResponse pendingBlockResponse ) {
117159 // write starts in the handler before the one associated with ctx
118160 // so add one that will be skipped (but that will prevent any writes later coming from later
119161 // handlers).
120162 // We do not want to start from the end of the
121163 // pipeline because there is an increased risk of hitting duplex handlers that
122164 // expect to have seen a request before processing the response
123- ctxForDownstream =
124- ctxForDownstream
125- .pipeline ()
126- .addAfter (
127- ctxForDownstream .name (),
128- "ignore_all_writes_handler" ,
129- IgnoreAllWritesHandler .INSTANCE )
130- .context ("ignore_all_writes_handler" );
131-
132- segment .effectivelyBlocked ();
133-
134- ctxForDownstream
135- .writeAndFlush (response )
165+ if (ctxForDownstream .pipeline ().get (IGNORE_ALL_WRITES_HANDLER ) == null ) {
166+ ctxForDownstream
167+ .pipeline ()
168+ .addAfter (
169+ ctxForDownstream .name (), IGNORE_ALL_WRITES_HANDLER , IgnoreAllWritesHandler .INSTANCE );
170+ }
171+ ChannelHandlerContext writeContext =
172+ ctxForDownstream .pipeline ().context (IGNORE_ALL_WRITES_HANDLER );
173+
174+ writeContext
175+ .writeAndFlush (pendingBlockResponse .toResponse ())
136176 .addListener (
137177 fut -> {
138178 if (!fut .isSuccess ()) {
139179 log .warn ("Write of blocking response failed" , fut .cause ());
140180 }
141- ctx .channel ().close ();
181+ writeContext .channel ().close ();
142182 });
143183 }
144184
185+ private static class PendingBlockResponse {
186+ private final TraceSegment segment ;
187+ private final int statusCode ;
188+ private final BlockingContentType bct ;
189+ private final Map <String , String > extraHeaders ;
190+ private final String securityResponseId ;
191+ private final HttpVersion protocolVersion ;
192+ private final String acceptHeader ;
193+
194+ // Prevent the generation of BlockingResponseHandler$1 by making this constructor
195+ // package-private to allow the BlockingResponseHandler to call this. This module emits Java 8
196+ // bytecode, so it cannot use Java 11 nested access.
197+ PendingBlockResponse (
198+ TraceSegment segment ,
199+ int statusCode ,
200+ BlockingContentType bct ,
201+ Map <String , String > extraHeaders ,
202+ String securityResponseId ,
203+ HttpVersion protocolVersion ,
204+ String acceptHeader ) {
205+ this .segment = segment ;
206+ this .statusCode = statusCode ;
207+ this .bct = bct ;
208+ this .extraHeaders = extraHeaders ;
209+ this .securityResponseId = securityResponseId ;
210+ this .protocolVersion = protocolVersion ;
211+ this .acceptHeader = acceptHeader ;
212+ }
213+
214+ FullHttpResponse toResponse () {
215+ int httpCode = BlockingActionHelper .getHttpCode (statusCode );
216+ HttpResponseStatus httpResponseStatus = HttpResponseStatus .valueOf (httpCode );
217+ FullHttpResponse response = new DefaultFullHttpResponse (protocolVersion , httpResponseStatus );
218+
219+ HttpHeaders headers = response .headers ();
220+ headers .set ("Connection" , "close" );
221+
222+ for (Map .Entry <String , String > h : extraHeaders .entrySet ()) {
223+ headers .set (h .getKey (), h .getValue ());
224+ }
225+
226+ if (bct != BlockingContentType .NONE ) {
227+ BlockingActionHelper .TemplateType type =
228+ BlockingActionHelper .determineTemplateType (bct , acceptHeader );
229+ headers .set ("Content-type" , BlockingActionHelper .getContentType (type ));
230+
231+ byte [] template = BlockingActionHelper .getTemplate (type , securityResponseId );
232+ HttpUtil .setContentLength (response , template .length );
233+ response .content ().writeBytes (template );
234+ }
235+ segment .effectivelyBlocked ();
236+ return response ;
237+ }
238+ }
239+
145240 @ ChannelHandler .Sharable
146241 public static class IgnoreAllWritesHandler extends ChannelOutboundHandlerAdapter {
147242 public static final IgnoreAllWritesHandler INSTANCE = new IgnoreAllWritesHandler ();
0 commit comments