RequestIdAwarePropagation.java

import brave.propagation.Propagation;
import brave.propagation.TraceContext;
import brave.propagation.TraceContextOrSamplingFlags;
import brave.propagation.TraceIdContext;
import java.util.List;
import java.util.Locale;
import java.util.UUID;

/**
 * <p>Tries to extract a trace ID and span ID using the supplied default implementation, such as
 * {@code X-B3-TraceId}, {@code X-B3-SpanId}, etc. Otherwise, tries to parse {@code request_id} into
 * a trace ID.
 */
public final class RequestIdAwarePropagation<K> implements Propagation<K> {
  public static Factory newFactory(Factory delegate, String requestIdFieldName) {
    if (delegate == null) throw new NullPointerException("delegate == null");
    if (requestIdFieldName == null) throw new NullPointerException("requestIdFieldName == null");
    return new Factory(delegate, requestIdFieldName.toLowerCase(Locale.ROOT));
  }

  static class Factory extends Propagation.Factory {
    final Factory delegate;
    final String requestIdFieldName;

    Factory(Factory delegate, String requestIdFieldName) {
      this.delegate = delegate;
      this.requestIdFieldName = requestIdFieldName;
    }

    @Override public <K> Propagation<K> create(KeyFactory<K> keyFactory) {
      return new RequestIdAwarePropagation<>(this, keyFactory);
    }

    @Override public boolean supportsJoin() {
      return delegate.supportsJoin();
    }

    @Override public boolean requires128BitTraceId() {
      return delegate.requires128BitTraceId();
    }

    @Override public String toString() {
      return "RequestIdAwarePropagationFactory{" + requestIdFieldName + "}";
    }
  }

  final Propagation<K> delegate;
  final K requestIdKey;

  RequestIdAwarePropagation(Factory factory, KeyFactory<K> keyFactory) {
    this.delegate = factory.delegate.create(keyFactory);
    this.requestIdKey = keyFactory.create(factory.requestIdFieldName);
  }

  @Override public List<K> keys() {
    return delegate.keys();
  }

  @Override public <C> TraceContext.Injector<C> injector(Setter<C, K> setter) {
    return delegate.injector(setter);
  }

  @Override public <C> TraceContext.Extractor<C> extractor(Getter<C, K> getter) {
    if (getter == null) throw new NullPointerException("getter == null");
    return new FallbackToRequestIdExtractor<>(this, getter);
  }

  final class FallbackToRequestIdExtractor<C> implements TraceContext.Extractor<C> {

    final TraceContext.Extractor<C> delegate;
    final Getter<C, K> getter;
    final K requestIdKey;

    FallbackToRequestIdExtractor(RequestIdAwarePropagation<K> propagation, Getter<C, K> getter) {
      this.delegate = propagation.delegate.extractor(getter);
      this.getter = getter;
      this.requestIdKey = propagation.requestIdKey;
    }

    @Override public TraceContextOrSamplingFlags extract(C carrier) {
      TraceContextOrSamplingFlags extracted = delegate.extract(carrier);
      if (extracted != TraceContextOrSamplingFlags.EMPTY) return extracted;
      String requestId = getter.get(carrier, requestIdKey);
      if (requestId == null) return extracted;
      try {
        UUID uuid = UUID.fromString(requestId);
        return TraceContextOrSamplingFlags.create(TraceIdContext.newBuilder()
            .traceIdHigh(uuid.getMostSignificantBits())
            .traceId(uuid.getLeastSignificantBits())
            .build());
      } catch (IllegalArgumentException e) {
        return extracted;
      }
    }
  }

  @Override public String toString() {
    return "RequestIdAwarePropagation{" + requestIdKey + "}";
  }
}