Skip to content

Commit f1b9d8a

Browse files
sezen-datadogsezen-datadog
committed
smola advice implemented
Signed-off-by: sezen.leblay <[email protected]>
1 parent 489b3e1 commit f1b9d8a

File tree

24 files changed

+687
-65
lines changed

24 files changed

+687
-65
lines changed

‎dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java

Lines changed: 7 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -387,12 +387,6 @@ public AgentSpan onResponse(final AgentSpan span, final RESPONSE response) {
387387
final int status = status(response);
388388
onResponseStatus(span, status);
389389

390-
BiFunction<RequestContext, OutputStream, Void> start =
391-
tracer().getUniversalCallbackProvider().getCallback(EVENTS.responseBodyStart());
392-
393-
final OutputStream responseBody = responseBody(response);
394-
start.apply(span.getRequestContext(), responseBody);
395-
396390
AgentPropagation.ContextVisitor<RESPONSE> getter = responseGetter();
397391
if (getter != null) {
398392
ResponseHeaderTagClassifier tagger =
@@ -403,7 +397,7 @@ public AgentSpan onResponse(final AgentSpan span, final RESPONSE response) {
403397
}
404398

405399
if (!isAppSecOnResponseSeparate()) {
406-
callIGCallbackResponseAndHeaders(span, response, status, responseBody);
400+
callIGCallbackResponseAndHeaders(span, response, status);
407401
}
408402
}
409403
return span;
@@ -467,8 +461,7 @@ private Flow<Void> callIGCallbackRequestHeaders(AgentSpan span, REQUEST_CARRIER
467461
IGKeyClassifier.create(
468462
requestContext,
469463
cbp.getCallback(EVENTS.requestHeader()),
470-
cbp.getCallback(EVENTS.requestHeaderDone()),
471-
null);
464+
cbp.getCallback(EVENTS.requestHeaderDone()));
472465
if (null != igKeyClassifier) {
473466
getter.forEachKey(carrier, igKeyClassifier);
474467
return igKeyClassifier.done();
@@ -497,27 +490,21 @@ private Flow<Void> callIGCallbackRequestSessionId(final AgentSpan span, final RE
497490
}
498491

499492
private Flow<Void> callIGCallbackResponseAndHeaders(
500-
AgentSpan span, RESPONSE carrier, int status, OutputStream responseBody) {
501-
return callIGCallbackResponseAndHeaders(span, carrier, status, responseGetter(), responseBody);
493+
AgentSpan span, RESPONSE carrier, int status) {
494+
return callIGCallbackResponseAndHeaders(span, carrier, status, responseGetter());
502495
}
503496

504497
public <RESP> Flow<Void> callIGCallbackResponseAndHeaders(
505498
AgentSpan span,
506499
RESP carrier,
507500
int status,
508-
AgentPropagation.ContextVisitor<RESP> contextVisitor,
509-
OutputStream responseBody) {
501+
AgentPropagation.ContextVisitor<RESP> contextVisitor) {
510502
CallbackProvider cbp = tracer().getCallbackProvider(RequestContextSlot.APPSEC);
511503
RequestContext requestContext = span.getRequestContext();
512504
if (cbp == null || requestContext == null) {
513505
return Flow.ResultFlow.empty();
514506
}
515507

516-
BiFunction<RequestContext, OutputStream, Flow<Void>> responseBodyDone =
517-
cbp.getCallback(EVENTS.responseBodyDone());
518-
if (null != responseBodyDone) {
519-
responseBodyDone.apply(requestContext, responseBody);
520-
}
521508
BiFunction<RequestContext, Integer, Flow<Void>> addrCallback =
522509
cbp.getCallback(EVENTS.responseStarted());
523510
if (null != addrCallback) {
@@ -530,8 +517,7 @@ public <RESP> Flow<Void> callIGCallbackResponseAndHeaders(
530517
IGKeyClassifier.create(
531518
requestContext,
532519
cbp.getCallback(EVENTS.responseHeader()),
533-
cbp.getCallback(EVENTS.responseHeaderDone()),
534-
cbp.getCallback(EVENTS.responseBodyDone()));
520+
cbp.getCallback(EVENTS.responseHeaderDone()));
535521
if (null != igKeyClassifier) {
536522
contextVisitor.forEachKey(carrier, igKeyClassifier);
537523
return igKeyClassifier.done();
@@ -615,8 +601,7 @@ protected static final class IGKeyClassifier implements AgentPropagation.KeyClas
615601
public static IGKeyClassifier create(
616602
RequestContext requestContext,
617603
TriConsumer<RequestContext, String, String> headerCallback,
618-
Function<RequestContext, Flow<Void>> doneCallback,
619-
BiFunction<RequestContext, OutputStream, Flow<Void>> bodyDoneCallback) {
604+
Function<RequestContext, Flow<Void>> doneCallback) {
620605
if (null == requestContext || null == headerCallback) {
621606
return null;
622607
}

‎dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@
1515
import datadog.trace.api.internal.TraceSegment;
1616
import datadog.trace.util.stacktrace.StackTraceEvent;
1717
import java.io.Closeable;
18-
import java.io.OutputStream;
1918
import java.util.*;
2019
import java.util.concurrent.ConcurrentHashMap;
2120
import java.util.concurrent.ConcurrentLinkedQueue;
@@ -99,7 +98,7 @@ public class AppSecRequestContext implements DataBundle, Closeable {
9998
private String inferredClientIp;
10099

101100
private volatile StoredBodySupplier storedRequestBodySupplier;
102-
private volatile OutputStream storedResponseBodySupplier;
101+
private volatile StoredBodySupplier storedResponseBodySupplier;
103102
private String dbType;
104103

105104
private int responseStatus;
@@ -455,7 +454,7 @@ void setStoredRequestBodySupplier(StoredBodySupplier storedRequestBodySupplier)
455454
this.storedRequestBodySupplier = storedRequestBodySupplier;
456455
}
457456

458-
void setStoredResponseBodySupplier(OutputStream storedResponseBodySupplier) {
457+
void setStoredResponseBodySupplier(StoredBodySupplier storedResponseBodySupplier) {
459458
this.storedResponseBodySupplier = storedResponseBodySupplier;
460459
}
461460

‎dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,6 @@
3737
import datadog.trace.bootstrap.instrumentation.api.URIDataAdapter;
3838
import datadog.trace.util.stacktrace.StackTraceEvent;
3939
import datadog.trace.util.stacktrace.StackUtils;
40-
import java.io.OutputStream;
4140
import java.net.URI;
4241
import java.net.URISyntaxException;
4342
import java.nio.charset.Charset;
@@ -660,7 +659,7 @@ private Flow<Void> onRequestBodyDone(RequestContext ctx_, StoredBodySupplier sup
660659
}
661660
}
662661

663-
private Flow<Void> onResponseBodyDone(RequestContext ctx_, OutputStream supplier) {
662+
private Flow<Void> onResponseBodyDone(RequestContext ctx_, StoredBodySupplier supplier) {
664663
AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC);
665664
if (ctx == null || ctx.isRawResBodyPublished()) {
666665
return NoopFlow.INSTANCE;
@@ -728,7 +727,7 @@ private Void onRequestBodyStart(RequestContext ctx_, StoredBodySupplier supplier
728727
return null;
729728
}
730729

731-
private Void onResponseBodyStart(RequestContext ctx_, OutputStream supplier) {
730+
private Void onResponseBodyStart(RequestContext ctx_, StoredBodySupplier supplier) {
732731
AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC);
733732
if (ctx == null) {
734733
return null;

‎dd-java-agent/instrumentation/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/appsec/BlockingResponseHelper.java

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -37,11 +37,7 @@ public static HttpResponse handleFinishForWaf(final AgentSpan span, final HttpRe
3737
}
3838
Flow<Void> flow =
3939
DECORATE.callIGCallbackResponseAndHeaders(
40-
span,
41-
response,
42-
response.status().intValue(),
43-
AkkaHttpServerHeaders.responseGetter(),
44-
null);
40+
span, response, response.status().intValue(), AkkaHttpServerHeaders.responseGetter());
4541
Flow.Action action = flow.getAction();
4642
if (action instanceof Flow.Action.RequestBlockingAction) {
4743
Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action;

‎dd-java-agent/instrumentation/jetty-7.0/src/main/java/datadog/trace/instrumentation/jetty70/JettyCommitResponseInstrumentation.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -84,7 +84,7 @@ static class CommitResponseAdvice {
8484

8585
Flow<Void> flow =
8686
DECORATE.callIGCallbackResponseAndHeaders(
87-
span, resp, resp.getStatus(), ExtractAdapter.Response.GETTER, null);
87+
span, resp, resp.getStatus(), ExtractAdapter.Response.GETTER);
8888
Flow.Action action = flow.getAction();
8989
if (action instanceof Flow.Action.RequestBlockingAction) {
9090
Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action;

‎dd-java-agent/instrumentation/jetty-7.6/src/main/java/datadog/trace/instrumentation/jetty76/JettyCommitResponseInstrumentation.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -98,7 +98,7 @@ static class CommitResponseAdvice {
9898

9999
Flow<Void> flow =
100100
DECORATE.callIGCallbackResponseAndHeaders(
101-
span, resp, resp.getStatus(), ExtractAdapter.Response.GETTER, null);
101+
span, resp, resp.getStatus(), ExtractAdapter.Response.GETTER);
102102
Flow.Action action = flow.getAction();
103103
if (action instanceof Flow.Action.RequestBlockingAction) {
104104
Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action;

‎dd-java-agent/instrumentation/jetty-9/src/main/java/datadog/trace/instrumentation/jetty9/JettyCommitResponseInstrumentation.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -119,7 +119,7 @@ static class CommitResponseAdvice {
119119

120120
Flow<Void> flow =
121121
DECORATE.callIGCallbackResponseAndHeaders(
122-
span, resp, resp.getStatus(), ExtractAdapter.Response.GETTER, null);
122+
span, resp, resp.getStatus(), ExtractAdapter.Response.GETTER);
123123
Flow.Action action = flow.getAction();
124124
if (action instanceof Flow.Action.RequestBlockingAction) {
125125
Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action;

‎dd-java-agent/instrumentation/jetty-9/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/JettyCommitResponseHelper.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,7 @@ public class JettyCommitResponseHelper {
6969
Response resp = connection.getResponse();
7070
Flow<Void> flow =
7171
JettyDecorator.DECORATE.callIGCallbackResponseAndHeaders(
72-
span, resp, resp.getStatus(), ExtractAdapter.Response.GETTER, null);
72+
span, resp, resp.getStatus(), ExtractAdapter.Response.GETTER);
7373
Flow.Action action = flow.getAction();
7474
if (action instanceof Flow.Action.RequestBlockingAction) {
7575
Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action;

‎dd-java-agent/instrumentation/jetty-9/src/main/java_jetty904/datadog/trace/instrumentation/jetty904/JettyCommitResponseHelper.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,7 @@ public class JettyCommitResponseHelper {
6969

7070
Flow<Void> flow =
7171
DECORATE.callIGCallbackResponseAndHeaders(
72-
span, resp, resp.getStatus(), ExtractAdapter.Response.GETTER, null);
72+
span, resp, resp.getStatus(), ExtractAdapter.Response.GETTER);
7373
Flow.Action action = flow.getAction();
7474
if (action instanceof Flow.Action.RequestBlockingAction) {
7575
Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action;

‎dd-java-agent/instrumentation/jetty-9/src/main/java_jetty93/datadog/trace/instrumentation/jetty93/JettyCommitResponseHelper.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,7 @@ public class JettyCommitResponseHelper {
6868
Response resp = connection.getResponse();
6969
Flow<Void> flow =
7070
DECORATE.callIGCallbackResponseAndHeaders(
71-
span, resp, resp.getStatus(), ExtractAdapter.Response.GETTER, null);
71+
span, resp, resp.getStatus(), ExtractAdapter.Response.GETTER);
7272
Flow.Action action = flow.getAction();
7373
if (action instanceof Flow.Action.RequestBlockingAction) {
7474
Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action;

‎dd-java-agent/instrumentation/jetty-9/src/main/java_jetty9421/datadog/trace/instrumentation/jetty9421/JettyCommitResponseHelper.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,7 @@ public class JettyCommitResponseHelper {
7171
Response resp = connection.getResponse();
7272
Flow<Void> flow =
7373
DECORATE.callIGCallbackResponseAndHeaders(
74-
span, resp, resp.getStatus(), ExtractAdapter.Response.GETTER, null);
74+
span, resp, resp.getStatus(), ExtractAdapter.Response.GETTER);
7575
Flow.Action action = flow.getAction();
7676
if (action instanceof Flow.Action.RequestBlockingAction) {
7777
Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action;

‎dd-java-agent/instrumentation/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/MaybeBlockResponseHandler.java

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -69,11 +69,7 @@ public void writeRequested(ChannelHandlerContext ctx, MessageEvent msg) throws E
6969

7070
Flow<Void> flow =
7171
DECORATE.callIGCallbackResponseAndHeaders(
72-
span,
73-
origResponse,
74-
origResponse.getStatus().getCode(),
75-
ResponseExtractAdapter.GETTER,
76-
null);
72+
span, origResponse, origResponse.getStatus().getCode(), ResponseExtractAdapter.GETTER);
7773
channelTraceContext.setAnalyzedResponse(true);
7874
Flow.Action action = flow.getAction();
7975
if (!(action instanceof Flow.Action.RequestBlockingAction)) {

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

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -88,11 +88,7 @@ public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise prm) thr
8888

8989
Flow<Void> flow =
9090
DECORATE.callIGCallbackResponseAndHeaders(
91-
span,
92-
origResponse,
93-
origResponse.getStatus().code(),
94-
ResponseExtractAdapter.GETTER,
95-
null);
91+
span, origResponse, origResponse.getStatus().code(), ResponseExtractAdapter.GETTER);
9692
markAnalyzedResponse(channel);
9793
Flow.Action action = flow.getAction();
9894
if (!(action instanceof Flow.Action.RequestBlockingAction)) {

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

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -85,11 +85,7 @@ public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise prm) thr
8585

8686
Flow<Void> flow =
8787
DECORATE.callIGCallbackResponseAndHeaders(
88-
span,
89-
origResponse,
90-
origResponse.getStatus().code(),
91-
ResponseExtractAdapter.GETTER,
92-
null);
88+
span, origResponse, origResponse.getStatus().code(), ResponseExtractAdapter.GETTER);
9389
markAnalyzedResponse(channel);
9490
Flow.Action action = flow.getAction();
9591
if (!(action instanceof Flow.Action.RequestBlockingAction)) {

0 commit comments

Comments
 (0)