com.undefinedlabs.scope.propagation.internal.TextMapCodec Maven / Gradle / Ivy

Go to download
package com.undefinedlabs.scope.propagation.internal;

import com.undefinedlabs.scope.SpanContext;
import com.undefinedlabs.scope.TraceId;
import org.apache.commons.lang3.StringUtils;
import com.undefinedlabs.scope.propagation.Codec;
import com.undefinedlabs.scope.utils.UrlUtils;
import io.opentracing.propagation.TextMap;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class TextMapCodec implements Codec {

  // OpenTracing
  private static final String PREFIX_TRACE_STATE = "ot-tracer-";
  protected static final String PREFIX_BAGGAGE = "ot-baggage-";
  protected static final String OT_FIELD_TRACE_ID = PREFIX_TRACE_STATE + "traceid";
  protected static final String OT_FIELD_SPAN_ID = PREFIX_TRACE_STATE + "spanid";
  protected static final String OT_FIELD_PARENT_SPAN_ID = PREFIX_TRACE_STATE + "parentspanid";
  protected static final String OT_FIELD_SAMPLED = PREFIX_TRACE_STATE + "sampled";

  // OpenTelemetry
  protected static final String OTEL_TRACEPARENT = "traceparent";
  protected static final String OTEL_TRACESTATE = "tracestate";
  private static final String OTEL_VERSION = "00";
  private static final String OTEL_TRACEPARENT_DELIMITER = "-";
  private static final String OTEL_FLAGS_SAMPLED = "01";
  private static final String ZEROS_IDS = "00000000000000000000000000000000";
  protected static final String OTEL_TRACESTATE_VENDOR = "scope";

  private final boolean useUrlEncoding;

  private TextMapCodec(final Builder builder) {
    useUrlEncoding = builder.useUrlEncoding;
  }

  @Override
  public void inject(final SpanContext spanContext, final TextMap carrier) {
    // OpenTracing
    carrier.put(
        OT_FIELD_TRACE_ID,
        useUrlEncoding ? UrlUtils.encodeUtf8(spanContext.toTraceId()) : spanContext.toTraceId());
    carrier.put(
        OT_FIELD_SPAN_ID,
        useUrlEncoding
            ? UrlUtils.encodeUtf8(spanContext.toSpanId64bitStr())
            : spanContext.toSpanId64bitStr());
    carrier.put(
        OT_FIELD_PARENT_SPAN_ID,
        useUrlEncoding
            ? UrlUtils.encodeUtf8(spanContext.toParentSpanId64bitStr())
            : spanContext.toParentSpanId64bitStr());
    carrier.put(OT_FIELD_SAMPLED, "true");

    final List otelTracestateList = new ArrayList<>();
    for (final Map.Entry baggageItem : spanContext.baggageItems()) {
      final String key =
          useUrlEncoding ? UrlUtils.encodeUtf8(baggageItem.getKey()) : baggageItem.getKey();
      final String value =
          useUrlEncoding ? UrlUtils.encodeUtf8(baggageItem.getValue()) : baggageItem.getValue();

      carrier.put(PREFIX_BAGGAGE + key, value);
      otelTracestateList.add(OTEL_TRACESTATE_VENDOR + "/" + key.replace(".", "-") + "=" + value);
    }

    // OpenTelemetry
    final String otelSb =
        OTEL_VERSION
            + OTEL_TRACEPARENT_DELIMITER
            + spanContext.toTraceId()
            + OTEL_TRACEPARENT_DELIMITER
            + spanContext.toSpanId()
            + OTEL_TRACEPARENT_DELIMITER
            + OTEL_FLAGS_SAMPLED;
    carrier.put(OTEL_TRACEPARENT, otelSb);
    carrier.put(OTEL_TRACESTATE, StringUtils.join(otelTracestateList, ","));
  }

  @Override
  public SpanContext extract(final TextMap carrier) {
    final Map otelIdsMap = new HashMap<>();
    final Map otracingIdsMap = new HashMap<>();
    final Map baggage = new HashMap<>();

    boolean hasTracestate = false;
    for (final Map.Entry carrierEntry : carrier) {
      if (carrierEntry.getKey().toLowerCase().equals(OTEL_TRACEPARENT)) {
        final String traceparent =
            useUrlEncoding ? UrlUtils.decodeUtf8(carrierEntry.getValue()) : carrierEntry.getValue();
        final String[] traceparentChunks = traceparent.split(OTEL_TRACEPARENT_DELIMITER);
        otelIdsMap.put(OT_FIELD_TRACE_ID, traceparentChunks[1]);
        otelIdsMap.put(OT_FIELD_SPAN_ID, traceparentChunks[2]);
      }

      if (carrierEntry.getKey().toLowerCase().equals(OTEL_TRACESTATE)) {
        hasTracestate = true;
        final String otelTracestate =
            useUrlEncoding ? UrlUtils.decodeUtf8(carrierEntry.getValue()) : carrierEntry.getValue();
        final String[] otelTracestateChunks = otelTracestate.split(",");
        for (final String chunk : otelTracestateChunks) {
          final String[] keyValue = chunk.split("=");
          if (keyValue[0].toLowerCase().startsWith(OTEL_TRACESTATE_VENDOR)) {
            baggage.put(
                keyValue[0]
                    .toLowerCase()
                    .replace(OTEL_TRACESTATE_VENDOR + "/", "")
                    .replace("-", "."),
                keyValue[1]);
          }
        }
      }

      if (carrierEntry.getKey().startsWith(PREFIX_BAGGAGE) && !hasTracestate) {
        baggage.put(
            carrierEntry.getKey().substring(PREFIX_BAGGAGE.length()),
            useUrlEncoding
                ? UrlUtils.decodeUtf8(carrierEntry.getValue())
                : carrierEntry.getValue());
      } else {
        otracingIdsMap.put(
            carrierEntry.getKey(),
            useUrlEncoding
                ? UrlUtils.decodeUtf8(carrierEntry.getValue())
                : carrierEntry.getValue());
      }
    }

    return !otelIdsMap.isEmpty()
        ? createSpanContext(otelIdsMap, baggage)
        : createSpanContext(otracingIdsMap, baggage);
  }

  private SpanContext createSpanContext(
      final Map idsMap, final Map baggage) {
    final String traceIdLowerCaseHex = idsMap.get(OT_FIELD_TRACE_ID);
    final String spanIdLowerCaseHex = idsMap.get(OT_FIELD_SPAN_ID);
    final String parentSpanIdLowerCaseHex = idsMap.get(OT_FIELD_PARENT_SPAN_ID);

    return new SpanContext(
        (traceIdLowerCaseHex == null)
            ? TraceId.getInvalid()
            : TraceId.fromLowerCaseBase16(
                ZEROS_IDS.substring(traceIdLowerCaseHex.length()) + traceIdLowerCaseHex, 0),
        (spanIdLowerCaseHex == null)
            ? TraceId.getInvalid()
            : TraceId.fromLowerCaseBase16(
                ZEROS_IDS.substring(spanIdLowerCaseHex.length()) + spanIdLowerCaseHex, 0),
        (parentSpanIdLowerCaseHex == null)
            ? null
            : TraceId.fromLowerCaseBase16(
                ZEROS_IDS.substring(parentSpanIdLowerCaseHex.length()) + parentSpanIdLowerCaseHex,
                0),
        baggage);
  }

  public static Builder newBuilder() {
    return new Builder();
  }

  public static class Builder {

    private boolean useUrlEncoding;

    public Builder withUseUrlEncoding(final boolean useUrlEncoding) {
      this.useUrlEncoding = useUrlEncoding;
      return this;
    }

    public TextMapCodec build() {
      return new TextMapCodec(this);
    }
  }
}