Tracing utilities

[zeitlinger]

The Opentelemetry APIs are rather difficult to use - therefore I've developed the following utility method, which is part of a company internal library - so it has already been tested by some users:

Does it make sense to include some/all method in the API - or in a contrib library - or not at all?

package org.zalando.observability;

import static io.opentelemetry.api.trace.SpanKind.CONSUMER;
import static io.opentelemetry.api.trace.SpanKind.SERVER;

import io.opentelemetry.api.GlobalOpenTelemetry;
import io.opentelemetry.api.baggage.Baggage;
import io.opentelemetry.api.baggage.BaggageBuilder;
import io.opentelemetry.api.trace.Span;
import io.opentelemetry.api.trace.SpanBuilder;
import io.opentelemetry.api.trace.SpanKind;
import io.opentelemetry.api.trace.StatusCode;
import io.opentelemetry.api.trace.Tracer;
import io.opentelemetry.context.Context;
import io.opentelemetry.context.Scope;
import io.opentelemetry.context.propagation.TextMapGetter;
import java.util.HashMap;
import java.util.Locale;
import java.util.Map;
import java.util.function.BiConsumer;
import java.util.function.Supplier;
import java.util.stream.Collectors;

/** Utility methods for tracing. */
public class Tracing {

  // only visible for compatibility
  public static Supplier<Tracer> tracerSupplier =
      () -> GlobalOpenTelemetry.get().getTracer("observability-sdk-utils");

  private static final TextMapGetter<Map<String, String>> TEXT_MAP_GETTER =
      new TextMapGetter<Map<String, String>>() {
        @Override
        public Iterable<String> keys(Map<String, String> carrier) {
          return carrier.values();
        }

        @Override
        public String get(Map<String, String> carrier, String key) {
          //noinspection ConstantConditions
          return carrier.get(key);
        }
      };

  private Tracing() {}

  /**
   * Marks the current span as error.
   *
   * @param description what went wrong
   */
  public static void setSpanError(String description) {
    setSpanError(Span.current(), description);
  }

  /**
   * Marks the current span as error.
   *
   * @param exception the exception that caused the error
   */
  public static void setSpanError(Throwable exception) {
    setSpanError(Span.current(), exception);
  }

  /**
   * Marks the current span as error.
   *
   * @param description what went wrong
   * @param exception the exception that caused the error
   */
  public static void setSpanError(String description, Throwable exception) {
    setSpanError(Span.current(), description, exception);
  }

  /**
   * Marks a span as error.
   *
   * @param span the span
   * @param description what went wrong
   */
  public static void setSpanError(Span span, String description) {
    span.setStatus(StatusCode.ERROR, description);
  }

  /**
   * Marks a span as error.
   *
   * @param span the span
   * @param exception the exception that caused the error
   */
  public static void setSpanError(Span span, Throwable exception) {
    span.setStatus(StatusCode.ERROR);
    span.recordException(exception);
  }

  /**
   * Marks a span as error.
   *
   * @param span the span
   * @param description what went wrong
   * @param exception the exception that caused the error
   */
  public static void setSpanError(Span span, String description, Throwable exception) {
    span.setStatus(StatusCode.ERROR, description);
    span.recordException(exception);
  }

  /**
   * Runs a block of code with a new span - ending the span at the end and recording any exception.
   *
   * @param operation operation name of the new span
   */
  public static <T> T trace(String operation, CheckedSupplier<T> block) {
    return trace(tracerSupplier.get().spanBuilder(operation).startSpan(), block);
  }

  /**
   * Runs a block of code with a new span - ending the span at the end and recording any exception.
   */
  public static <T> T trace(Span span, CheckedSupplier<T> block) {
    return traceBlock(span, block, Tracing::setSpanError);
  }

  private static <T> T traceBlock(
      Span span, CheckedSupplier<T> block, BiConsumer<Span, Exception> handleException) {
    try (Scope ignore = span.makeCurrent()) {
      return block.get();
    } catch (Exception e) {
      // not just RuntimeException, because Kotlin can easily throw non-RuntimeExceptions
      handleException.accept(span, e);
      sneakyThrow(e);
      return null;
    } catch (Throwable t) {
      // if it not an Exception, it must be a throwable and we could have an OutOfMemoryError, so we
      // don't do anything with the span
      sneakyThrow(t);
      return null;
    } finally {
      span.end();
    }
  }

  /**
   * Trace a block of code using a server span.
   *
   * <p>The span context will be extracted from the <code>transport</code>, which you usually get
   * from HTTP headers of the metadata of an event you're processing.
   */
  public static <T> T traceServerSpan(
      Map<String, String> transport, SpanBuilder spanBuilder, CheckedSupplier<T> block) {
    return extractAndRun(SERVER, transport, spanBuilder, block, Tracing::setSpanError);
  }

  /**
   * Trace a block of code using a server span.
   *
   * <p>The span context will be extracted from the <code>transport</code>, which you usually get
   * from HTTP headers of the metadata of an event you're processing.
   */
  public static <T> T traceServerSpan(
      Map<String, String> transport,
      SpanBuilder spanBuilder,
      CheckedSupplier<T> block,
      BiConsumer<Span, Exception> handleException) {
    return extractAndRun(SERVER, transport, spanBuilder, block, handleException);
  }

  /**
   * Trace a block of code using a consumer span.
   *
   * <p>The span context will be extracted from the <code>transport</code>, which you usually get
   * from HTTP headers of the metadata of an event you're processing.
   */
  public static <T> T traceConsumerSpan(
      Map<String, String> transport, SpanBuilder spanBuilder, CheckedSupplier<T> block) {
    return extractAndRun(CONSUMER, transport, spanBuilder, block, Tracing::setSpanError);
  }

  /**
   * Trace a block of code using a consumer span.
   *
   * <p>The span context will be extracted from the <code>transport</code>, which you usually get
   * from HTTP headers of the metadata of an event you're processing.
   */
  public static <T> T traceConsumerSpan(
      Map<String, String> transport,
      SpanBuilder spanBuilder,
      CheckedSupplier<T> block,
      BiConsumer<Span, Exception> handleException) {
    return extractAndRun(CONSUMER, transport, spanBuilder, block, handleException);
  }

  private static <T> T extractAndRun(
      SpanKind spanKind,
      Map<String, String> transport,
      SpanBuilder spanBuilder,
      CheckedSupplier<T> block,
      BiConsumer<Span, Exception> handleException) {
    try (Scope ignore = extractContext(transport).makeCurrent()) {
      return traceBlock(spanBuilder.setSpanKind(spanKind).startSpan(), block, handleException);
    }
  }

  /**
   * Injects the current context into a string map, which can then be added to HTTP headers or the
   * metadata of an event.
   */
  public static Map<String, String> injectContext() {
    Map<String, String> transport = new HashMap<>();
    //noinspection ConstantConditions
    GlobalOpenTelemetry.get()
        .getPropagators()
        .getTextMapPropagator()
        .inject(Context.current(), transport, Map::put);
    return transport;
  }

  /**
   * Extract the context from a string map, which you get from HTTP headers of the metadata of an
   * event you're processing.
   */
  public static Context extractContext(Map<String, String> transport) {
    Context current = Context.current();
    //noinspection ConstantConditions
    if (transport == null) {
      return current;
    }
    // HTTP headers are case-insensitive. As we're using Map, which is case-sensitive, we need to
    // normalize all the keys
    Map<String, String> normalizedTransport =
        transport.entrySet().stream()
            .collect(
                Collectors.toMap(
                    entry -> entry.getKey().toLowerCase(Locale.ROOT), Map.Entry::getValue));
    return GlobalOpenTelemetry.get()
        .getPropagators()
        .getTextMapPropagator()
        .extract(current, normalizedTransport, TEXT_MAP_GETTER);
  }

  /** Sets baggage items which are active in given block. */
  public static <T> T setBaggage(Map<String, String> baggage, CheckedSupplier<T> block) {
    BaggageBuilder builder = Baggage.current().toBuilder();
    baggage.forEach(builder::put);
    Context context = builder.build().storeInContext(Context.current());
    try (Scope ignore = context.makeCurrent()) {
      return block.get();
    } catch (Throwable e) {
      sneakyThrow(e);
      return null;
    }
  }

  private static <E extends Throwable> void sneakyThrow(Throwable e) throws E {
    //noinspection unchecked
    throw (E) e;
  }
}

package org.zalando.observability;

/** A supplier that can throw anything. */
@FunctionalInterface
public interface CheckedSupplier<T> {
  T get() throws Throwable;
}

[jkwatson]

This seems like it could be a good addition the opentelemetry-java-contrib libraries. Having a suite of helpers like this could definitely be useful, but I'm not sure we'd want to put this directly into the core repository at this point.

[zeitlinger]

Just discovered another utility class. Would it make sense to add to this one?