Update aws-java-sdk-1.11.0 to new agent api

This commit is contained in:
Trask Stalnaker 2019-10-19 11:57:38 -07:00
parent 887633f086
commit 09cdc177cf
3 changed files with 18 additions and 15 deletions

View File

@ -12,7 +12,7 @@ import com.amazonaws.Request;
import com.amazonaws.handlers.RequestHandler2;
import com.google.auto.service.AutoService;
import datadog.trace.agent.tooling.Instrumenter;
import io.opentracing.Scope;
import datadog.trace.instrumentation.api.AgentScope;
import java.util.Map;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.method.MethodDescription;
@ -60,7 +60,7 @@ public class AWSHttpClientInstrumentation extends Instrumenter.Default {
@Advice.Argument(value = 0, optional = true) final Request<?> request,
@Advice.Thrown final Throwable throwable) {
if (throwable != null) {
final Scope scope = request.getHandlerContext(TracingRequestHandler.SCOPE_CONTEXT_KEY);
final AgentScope scope = request.getHandlerContext(TracingRequestHandler.SCOPE_CONTEXT_KEY);
if (scope != null) {
request.addHandlerContext(TracingRequestHandler.SCOPE_CONTEXT_KEY, null);
DECORATE.onError(scope.span(), throwable);
@ -96,7 +96,8 @@ public class AWSHttpClientInstrumentation extends Instrumenter.Default {
@Advice.FieldValue("request") final Request<?> request,
@Advice.Thrown final Throwable throwable) {
if (throwable != null) {
final Scope scope = request.getHandlerContext(TracingRequestHandler.SCOPE_CONTEXT_KEY);
final AgentScope scope =
request.getHandlerContext(TracingRequestHandler.SCOPE_CONTEXT_KEY);
if (scope != null) {
request.addHandlerContext(TracingRequestHandler.SCOPE_CONTEXT_KEY, null);
DECORATE.onError(scope.span(), throwable);

View File

@ -5,7 +5,7 @@ import com.amazonaws.Request;
import com.amazonaws.Response;
import datadog.trace.agent.decorator.HttpClientDecorator;
import datadog.trace.api.DDTags;
import io.opentracing.Span;
import datadog.trace.instrumentation.api.AgentSpan;
import java.net.URI;
import java.net.URISyntaxException;
import java.util.Map;
@ -20,7 +20,7 @@ public class AwsSdkClientDecorator extends HttpClientDecorator<Request, Response
private final Map<Class, String> operationNames = new ConcurrentHashMap<>();
@Override
public Span onRequest(final Span span, final Request request) {
public AgentSpan onRequest(final AgentSpan span, final Request request) {
// Call super first because we override the resource name below.
super.onRequest(span, request);
@ -40,7 +40,7 @@ public class AwsSdkClientDecorator extends HttpClientDecorator<Request, Response
}
@Override
public Span onResponse(final Span span, final Response response) {
public AgentSpan onResponse(final AgentSpan span, final Response response) {
if (response.getAwsResponse() instanceof AmazonWebServiceResponse) {
final AmazonWebServiceResponse awsResp = (AmazonWebServiceResponse) response.getAwsResponse();
span.setTag("aws.requestId", awsResp.getRequestId());

View File

@ -1,5 +1,7 @@
package datadog.trace.instrumentation.aws.v0;
import static datadog.trace.instrumentation.api.AgentTracer.activateSpan;
import static datadog.trace.instrumentation.api.AgentTracer.startSpan;
import static datadog.trace.instrumentation.aws.v0.AwsSdkClientDecorator.DECORATE;
import com.amazonaws.AmazonWebServiceRequest;
@ -7,8 +9,8 @@ import com.amazonaws.Request;
import com.amazonaws.Response;
import com.amazonaws.handlers.HandlerContextKey;
import com.amazonaws.handlers.RequestHandler2;
import io.opentracing.Scope;
import io.opentracing.util.GlobalTracer;
import datadog.trace.instrumentation.api.AgentScope;
import datadog.trace.instrumentation.api.AgentSpan;
/** Tracing Request Handler */
public class TracingRequestHandler extends RequestHandler2 {
@ -16,7 +18,7 @@ public class TracingRequestHandler extends RequestHandler2 {
// Note: aws1.x sdk doesn't have any truly async clients so we can store scope in request context
// safely.
public static final HandlerContextKey<Scope> SCOPE_CONTEXT_KEY =
public static final HandlerContextKey<AgentScope> SCOPE_CONTEXT_KEY =
new HandlerContextKey<>("DatadogScope");
@Override
@ -26,15 +28,15 @@ public class TracingRequestHandler extends RequestHandler2 {
@Override
public void beforeRequest(final Request<?> request) {
final Scope scope = GlobalTracer.get().buildSpan("aws.command").startActive(true);
DECORATE.afterStart(scope.span());
DECORATE.onRequest(scope.span(), request);
request.addHandlerContext(SCOPE_CONTEXT_KEY, scope);
final AgentSpan span = startSpan("aws.command");
DECORATE.afterStart(span);
DECORATE.onRequest(span, request);
request.addHandlerContext(SCOPE_CONTEXT_KEY, activateSpan(span, true));
}
@Override
public void afterResponse(final Request<?> request, final Response<?> response) {
final Scope scope = request.getHandlerContext(SCOPE_CONTEXT_KEY);
final AgentScope scope = request.getHandlerContext(SCOPE_CONTEXT_KEY);
if (scope != null) {
request.addHandlerContext(SCOPE_CONTEXT_KEY, null);
DECORATE.onResponse(scope.span(), response);
@ -45,7 +47,7 @@ public class TracingRequestHandler extends RequestHandler2 {
@Override
public void afterError(final Request<?> request, final Response<?> response, final Exception e) {
final Scope scope = request.getHandlerContext(SCOPE_CONTEXT_KEY);
final AgentScope scope = request.getHandlerContext(SCOPE_CONTEXT_KEY);
if (scope != null) {
request.addHandlerContext(SCOPE_CONTEXT_KEY, null);
DECORATE.onError(scope.span(), e);