tracing-test.java

package io.vertx.tracing.jaeger;

import io.jaegertracing.internal.JaegerSpanContext;
import io.opentracing.Scope;
import io.opentracing.Span;
import io.opentracing.SpanContext;
import io.opentracing.Tracer;
import io.opentracing.propagation.Format;
import io.opentracing.propagation.TextMap;
import io.opentracing.tag.Tags;
import io.vertx.core.http.HttpClientRequest;
import io.vertx.core.http.HttpServerRequest;
import io.vertx.ext.jdbc.impl.actions.AbstractJDBCAction;
import io.vertx.ext.sql.SQLConnection;

import java.util.Iterator;
import java.util.Map;

public class JaegerTracer implements io.vertx.core.spi.tracing.Tracer<Scope, Scope> {

  private final Tracer tracer;

  public JaegerTracer(Tracer tracer) {
    this.tracer = tracer;
  }

  @Override
  public Scope beginInbound(Map<Object, Object> context, Object inbound) {
    if (inbound instanceof HttpServerRequest) {
      HttpServerRequest req = (HttpServerRequest) inbound;
      SpanContext sc = tracer.extract(Format.Builtin.HTTP_HEADERS, new TextMap() {
        @Override
        public Iterator<Map.Entry<String, String>> iterator() {
          return req.headers().iterator();
        }
        @Override
        public void put(String key, String value) {
          throw new UnsupportedOperationException();
        }
      });
      Tracer.SpanBuilder builder = tracer.buildSpan(req.method().name());
      if (sc != null) {
        builder = builder.asChildOf(sc);
      }
      Scope scope = builder
        .ignoreActiveSpan()
        .withTag(Tags.SPAN_KIND.getKey(), Tags.SPAN_KIND_SERVER)
        .withTag(Tags.HTTP_METHOD.getKey(), req.method().name())
        .withTag(Tags.COMPONENT.getKey(), "vertx")
        .withTag(Tags.HTTP_URL.getKey(), req.absoluteURI())
        .startActive(false);
      JaegerSpanContext jsc = (JaegerSpanContext) scope.span().context();
      // System.out.println("receive span " + Long.toHexString(jsc.getSpanId()) + " " + Long.toHexString(jsc.getParentId()) + " " + jsc.getTraceId());
      scope.close();
      context.put("ACTIVE_SCOPE", scope);
      return scope;
    }
    return null;
  }

  @Override
  public void endInbound(Map<Object, Object> context, Scope scope, Throwable failed) {
    if (scope != null) {
      context.remove("ACTIVE_SCOPE");
      Span span = scope.span();
      // Tags.HTTP_STATUS.set(span, request.response().getStatusCode());
      span.finish();
    }
  }

  @Override
  public Scope beginOutbound(Map<Object, Object> context, Object outbound) {
    Scope scope = (Scope) context.get("ACTIVE_SCOPE");
    if (scope != null) {
      if (outbound instanceof HttpClientRequest) {
        HttpClientRequest req = (HttpClientRequest) outbound;
        if (scope != null) {
          Scope child = tracer.buildSpan(req.method().name())
            .asChildOf(scope.span())
            .withTag(Tags.SPAN_KIND.getKey(), Tags.SPAN_KIND_CLIENT)
            .startActive(false);
          tracer.inject(child.span().context(), Format.Builtin.HTTP_HEADERS, new TextMap() {
            @Override
            public Iterator<Map.Entry<String, String>> iterator() {
              throw new UnsupportedOperationException();
            }
            @Override
            public void put(String key, String value) {
              req.headers().set(key, value);
            }
          });
          JaegerSpanContext jsc = (JaegerSpanContext) child.span().context();
          // System.out.println("send / report span " + Long.toHexString(jsc.getSpanId()) + " " + Long.toHexString(jsc.getParentId()) + " " + jsc.getTraceId());
          return child;
        }
      } else if (outbound instanceof AbstractJDBCAction) {
        AbstractJDBCAction action = (AbstractJDBCAction) outbound;
        Scope child = tracer.buildSpan("SQL")
          .asChildOf(scope.span())
          .withTag(Tags.SPAN_KIND.getKey(), Tags.SPAN_KIND_CLIENT)
          .withTag(Tags.DB_TYPE.getKey(), "sql")
          .withTag(Tags.DB_STATEMENT.getKey(), action.statement())
          .startActive(false);
        JaegerSpanContext jsc = (JaegerSpanContext) child.span().context();
        // System.out.println("report span " + Long.toHexString(jsc.getSpanId()) + " " + Long.toHexString(jsc.getParentId()) + " " + jsc.getTraceId());
        return child;
      }
    }
    return null;
  }

  @Override
  public void endOutbound(Map<Object, Object> context, Scope scope, Throwable failed) {
    if (scope != null) {
      scope.span().finish();
    }
  }
}