io.opentracing.contrib.jdbi.OpenTracingCollector Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of jdbi-opentracing Show documentation
Root module for OpenTracing JDBI Instrumentation
The newest version!
package io.opentracing.contrib.jdbi;

import io.opentracing.Span;
import io.opentracing.Tracer;
import org.skife.jdbi.v2.SQLStatement;
import org.skife.jdbi.v2.StatementContext;
import org.skife.jdbi.v2.TimingCollector;

/**
 * OpenTracingCollector is a JDBI TimingCollector that creates OpenTracing Spans for each JDBI SQLStatement.
 *
 * Example usage:
 * 
{@code
 * io.opentracing.Tracer tracer = ...;
 * DBI dbi = ...;
 *
 * // One time only: bind OpenTracing to the DBI instance as a TimingCollector.
 * dbi.setTimingCollector(new OpenTracingCollector(tracer));
 *
 * // Elsewhere, anywhere a `Handle` is available:
 * Handle handle = ...;
 * Span parentSpan = ...;  // optional
 *
 * // Create statements as usual with your `handle` instance.
 *  Query> statement = handle.createQuery("SELECT COUNT(*) FROM accounts");
 *
 * // If a parent Span is available, establish the relationship via setParent.
 * OpenTracingCollector.setParent(statement, parent);
 *
 * // Use JDBI as per usual, and Spans will be created for every SQLStatement automatically.
 * List> results = statement.list();
 * }
 */
public class OpenTracingCollector implements TimingCollector {
    public final static String PARENT_SPAN_ATTRIBUTE_KEY = "io.opentracing.parent";

    private final Tracer tracer;
    private final SpanDecorator spanDecorator;
    private final ActiveSpanSource activeSpanSource;

    public OpenTracingCollector(Tracer tracer) {
        this(tracer, SpanDecorator.DEFAULT);
    }

    public OpenTracingCollector(Tracer tracer, SpanDecorator spanDecorator) {
        this(tracer, spanDecorator, null);
    }
    public OpenTracingCollector(Tracer tracer, ActiveSpanSource spanSource) {
        this(tracer, SpanDecorator.DEFAULT, spanSource);
    }

    public OpenTracingCollector(Tracer tracer, SpanDecorator operationNamer, ActiveSpanSource spanSource) {
        this.tracer = tracer;
        this.spanDecorator = operationNamer;
        this.activeSpanSource = spanSource;
    }

    public void collect(long elapsedNanos, StatementContext statementContext) {
        long nowMicros = System.currentTimeMillis() * 1000;
        Tracer.SpanBuilder builder = tracer
                .buildSpan(spanDecorator.generateOperationName(statementContext))
                .withStartTimestamp(nowMicros - (elapsedNanos / 1000));
        Span parent = (Span)statementContext.getAttribute(PARENT_SPAN_ATTRIBUTE_KEY);
        if (parent == null && this.activeSpanSource != null) {
            parent = this.activeSpanSource.activeSpan(statementContext);
        }
        if (parent != null) {
            builder = builder.asChildOf(parent);
        }
        Span collectSpan = builder.start();
        spanDecorator.decorateSpan(collectSpan, elapsedNanos, statementContext);
        try {
            collectSpan.log("SQL query finished", statementContext.getRawSql());
        } finally {
            collectSpan.finish(nowMicros);
        }
    }

    /**
     * Establish an explicit parent relationship for the (child) Span associated with a SQLStatement.
     *
     * @param statement the JDBI SQLStatement which will act as the child of `parent`
     * @param parent the parent Span for `statement`
     */
    public static void setParent(SQLStatement statement, Span parent) {
        statement.getContext().setAttribute(PARENT_SPAN_ATTRIBUTE_KEY, parent);
    }

    /**
     * SpanDecorator allows the OpenTracingCollector user to control the precise naming and decoration of OpenTracing
     * Spans emitted by the collector.
     *
     * @see OpenTracingCollector#OpenTracingCollector(Tracer, SpanDecorator)
     */
    public interface SpanDecorator {
        public static SpanDecorator DEFAULT = new SpanDecorator() {
            public String generateOperationName(StatementContext ctx) {
                return "DBI Statement";
            }

            @Override
            public void decorateSpan(Span jdbiSpan, long elapsedNanos, StatementContext ctx) {
                // (by default, do nothing)
            }
        };

        /**
         * Transform an DBI StatementContext into an OpenTracing Span operation name.
         *
         * @param ctx the StatementContext passed to TimingCollector.collect()
         * @return an operation name suitable for the associated OpenTracing Span
         */
        public String generateOperationName(StatementContext ctx);

        /**
         * Get the active Span (to use as a parent for any DBI Spans). Implementations may or may not need to refer
         * to the StatementContext.
         *
         * @param jdbiSpan the JDBI Span to decorate (before `finish` is called)
         * @param elapsedNanos the elapsedNanos passed to TimingCollector.collect()
         * @param ctx the StatementContext passed to TimingCollector.collect()
         */
        public void decorateSpan(Span jdbiSpan, long elapsedNanos, StatementContext ctx);
    }

    /**
     * An abstract API that allows the OpenTracingCollector to customize how parent Spans are discovered.
     *
     * For instance, if Spans are stored in a thread-local variable, an ActiveSpanSource could access them like so:
     * Example usage:
     * 
{@code
     * public class SomeClass {
     *     // Thread local variable containing each thread's ID
     *     private static final ThreadLocal activeSpan =
     *         new ThreadLocal() {
     *             protected Integer initialValue() {
     *                 return null;
     *             }
     *         };
     * };
     *
     * ... elsewhere ...
     * ActiveSpanSource spanSource = new ActiveSpanSource() {
     *     public Span activeSpan(StatementContext ctx) {
     *         // (In this example we ignore `ctx` entirely)
     *         return activeSpan.get();
     *     }
     * };
     * OpenTracingCollector otColl = new OpenTracingCollector(tracer, spanSource);
     * ...
     * }
     */
    public interface ActiveSpanSource {
        /**
         * Get the active Span (to use as a parent for any DBI Spans). Implementations may or may not need to refer
         * to the StatementContext.
         *
         * @param ctx the StatementContext that needs to be collected and traced
         * @return the currently active Span (for this thread, etc), or null if no such Span could be found.
         */
        public Span activeSpan(StatementContext ctx);
    }
}