[Feature Request] Instrument Java AWS-SDK clients

[pims]

Hi there,

I’ve been trying to instrument the Amazon S3 Client from the Java AWS-SDK.
New versions of the SDK offer hooks via the RequestHandler2 abstract class which implements the following interface:

public interface IRequestHandler2 {
    AmazonWebServiceRequest beforeExecution(AmazonWebServiceRequest request);
    AmazonWebServiceRequest beforeMarshalling(AmazonWebServiceRequest request);
    void beforeRequest(Request<?> request);
    HttpResponse beforeUnmarshalling(Request<?> request, HttpResponse httpResponse);
    void afterResponse(Request<?> request, Response<?> response);
    void afterError(Request<?> request, Response<?> response, Exception e);
}

I’ve tried something along those lines, but couldn't get it to work properly. @adriancole suggested raising the issue here.

public class ZipkinRequestHandler extends RequestHandler2 {
    final private Tracer tracer;
    final private CurrentTraceContext currentTraceContext;
    final private HttpClientHandler<Request, Response> handler;
    final private TraceContext.Injector<Request> injector;
    final private TraceContext.Extractor<Request> extractor;

    public static ZipkinRequestHandler create(final HttpTracing httpTracing,
                                              final HttpClientAdapter<Request, Response> adapter) {
        return new ZipkinRequestHandler(
                httpTracing.tracing().tracer(),
                httpTracing.tracing().currentTraceContext(),
                HttpClientHandler.create(httpTracing, adapter),
                httpTracing.tracing().propagation().injector(new Propagation.Setter<Request, String>() {
                    @Override
                    public void put(Request carrier, String key, String value) {
                        carrier.addHeader(key, value);
                    }
                }),
                httpTracing.tracing().propagation().extractor(new Propagation.Getter<Request, String>() {
                    @Override
                    public String get(Request carrier, String key) {
                        final Map<String, String> headers = carrier.getHeaders();
                        return headers.get(key);
                    }
                })

        );
    }

    private ZipkinRequestHandler(Tracer tracer, CurrentTraceContext ctc, HttpClientHandler<Request, Response> handler,
                                 TraceContext.Injector<Request> injector, TraceContext.Extractor<Request> extractor) {
        this.tracer = tracer;
        this.currentTraceContext = ctc;
        this.handler = handler;
        this.injector = injector;
        this.extractor = extractor;

    }

    @Override
    public void beforeRequest(Request<?> request) {
        TraceContext parent = currentTraceContext.get();
        try(CurrentTraceContext.Scope scope = currentTraceContext.newScope(parent)) {
            Span span = handler.handleSend(injector,request);
            span.annotate("start" + LocalDateTime.now().toString());
            System.out.println(LocalDateTime.now() + " beforeRequest " + span.toString());
        }


    }

    @Override
    public void afterResponse(Request<?> request, Response<?> response) {
        final Span span = tracer.joinSpan(extractor.extract(request).context());
        span.annotate("end-"+ LocalDateTime.now().toString());
        handler.handleReceive(response, null, span);
        System.out.println(LocalDateTime.now() + " afterResponse " + span.toString());
    }

    @Override
    public void afterError(Request<?> request, Response<?> response, Exception ex) {
        final Span span = tracer.joinSpan(extractor.extract(request).context());
        handler.handleReceive(null, ex, span);
        System.out.println("afterError " + span.toString());
    }
}

[devinsba]

Thanks for the detailed request! I've moved it to zipkin-aws so we can more effectively align it to other AWS work

openzipkin/zipkin-aws#53