diff --git a/dd-java-agent/instrumentation/synapse-3.0/src/main/java/datadog/trace/instrumentation/synapse3/SynapseClientInstrumentation.java b/dd-java-agent/instrumentation/synapse-3.0/src/main/java/datadog/trace/instrumentation/synapse3/SynapseClientInstrumentation.java index 19ca1eaab8d..dd57c58c552 100644 --- a/dd-java-agent/instrumentation/synapse-3.0/src/main/java/datadog/trace/instrumentation/synapse3/SynapseClientInstrumentation.java +++ b/dd-java-agent/instrumentation/synapse-3.0/src/main/java/datadog/trace/instrumentation/synapse3/SynapseClientInstrumentation.java @@ -21,6 +21,7 @@ import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import net.bytebuddy.asm.Advice; import org.apache.axis2.context.MessageContext; +import org.apache.http.HttpResponse; import org.apache.http.nio.NHttpClientConnection; import org.apache.synapse.transport.passthru.TargetContext; @@ -113,8 +114,8 @@ public static final class ClientResponseAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static ContextScope beginResponse( @Advice.Argument(0) final NHttpClientConnection connection) { - // check and remove context so it won't be finished twice - Context context = (Context) connection.getContext().removeAttribute(SYNAPSE_CONTEXT_KEY); + // don't remove stored context here because the response callback may run multiple times + Context context = (Context) connection.getContext().getAttribute(SYNAPSE_CONTEXT_KEY); if (null != context) { return context.attach(); } @@ -129,15 +130,30 @@ public static void responseReceived( if (null == scope) { return; } - AgentSpan span = spanFromContext(scope.context()); - DECORATE.onResponse(span, connection.getHttpResponse()); + final AgentSpan span = spanFromContext(scope.context()); + final HttpResponse httpResponse = connection.getHttpResponse(); + boolean isFinal = false; + if (httpResponse != null && httpResponse.getStatusLine() != null) { + int statusCode = httpResponse.getStatusLine().getStatusCode(); + // 1xx (e.g. 100 Continue) - not final response + isFinal = statusCode >= 200; + } + + if (isFinal) { + DECORATE.onResponse(span, httpResponse); + } if (null != error) { DECORATE.onError(span, error); } - DECORATE.beforeFinish(scope.context()); - scope.close(); - if (span != null) { - span.finish(); + if ((isFinal || error != null) + && connection.getContext().removeAttribute(SYNAPSE_CONTEXT_KEY) != null) { + DECORATE.beforeFinish(scope.context()); + scope.close(); + if (span != null) { + span.finish(); + } + } else { + scope.close(); } } } diff --git a/dd-java-agent/instrumentation/synapse-3.0/src/main/java/datadog/trace/instrumentation/synapse3/SynapseServerInstrumentation.java b/dd-java-agent/instrumentation/synapse-3.0/src/main/java/datadog/trace/instrumentation/synapse3/SynapseServerInstrumentation.java index 6bd0122e3c0..d16e628121c 100644 --- a/dd-java-agent/instrumentation/synapse-3.0/src/main/java/datadog/trace/instrumentation/synapse3/SynapseServerInstrumentation.java +++ b/dd-java-agent/instrumentation/synapse-3.0/src/main/java/datadog/trace/instrumentation/synapse3/SynapseServerInstrumentation.java @@ -17,6 +17,7 @@ import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import net.bytebuddy.asm.Advice; import org.apache.http.HttpRequest; +import org.apache.http.HttpResponse; import org.apache.http.nio.NHttpServerConnection; @AutoService(InstrumenterModule.class) @@ -92,8 +93,8 @@ public static final class ServerResponseAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static ContextScope beginResponse( @Advice.Argument(0) final NHttpServerConnection connection) { - // check and remove context so it won't be finished twice - Context context = (Context) connection.getContext().removeAttribute(SYNAPSE_CONTEXT_KEY); + // don't remove stored context here because the response callback may run multiple times + Context context = (Context) connection.getContext().getAttribute(SYNAPSE_CONTEXT_KEY); if (null != context) { return context.attach(); } @@ -108,14 +109,28 @@ public static void responseReady( if (null == scope) { return; } - AgentSpan span = spanFromContext(scope.context()); - DECORATE.onResponse(span, connection.getHttpResponse()); + final AgentSpan span = spanFromContext(scope.context()); + final HttpResponse httpResponse = connection.getHttpResponse(); + boolean isFinal = false; + if (httpResponse != null && httpResponse.getStatusLine() != null) { + isFinal = httpResponse.getStatusLine().getStatusCode() >= 200; + } + + if (isFinal) { + DECORATE.onResponse(span, httpResponse); + } if (null != error) { DECORATE.onError(span, error); } - DECORATE.beforeFinish(scope.context()); - scope.close(); - span.finish(); + + if ((isFinal || error != null) + && connection.getContext().removeAttribute(SYNAPSE_CONTEXT_KEY) != null) { + DECORATE.beforeFinish(scope.context()); + scope.close(); + span.finish(); + } else { + scope.close(); + } } }