com.undefinedlabs.scope.SpanBuilder Maven / Gradle / Ivy

Go to download
package com.undefinedlabs.scope;

import org.apache.commons.lang3.StringUtils;
import com.undefinedlabs.scope.statistics.Statistics;
import io.opentracing.References;
import io.opentracing.Tracer;
import io.opentracing.tag.Tag;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class SpanBuilder implements Tracer.SpanBuilder {

  private final ScopeTracer scopeTracer;
  private final String operationName;
  private final Map tags;
  private long startTimeMicroseconds;
  private long startTimeNanoseconds;
  private boolean ignoreActiveSpan;
  private List references;

  public SpanBuilder(final ScopeTracer scopeTracer, final String operationName) {
    this.scopeTracer = scopeTracer;
    this.operationName = StringUtils.abbreviate(operationName, 255);
    references = Collections.EMPTY_LIST;
    tags = new HashMap<>();
    ignoreActiveSpan = false;
  }

  @Override
  public Tracer.SpanBuilder asChildOf(final io.opentracing.SpanContext parent) {
    return addReference(References.CHILD_OF, parent);
  }

  @Override
  public Tracer.SpanBuilder asChildOf(final io.opentracing.Span parent) {
    return addReference(References.CHILD_OF, parent != null ? parent.context() : null);
  }

  @Override
  public Tracer.SpanBuilder addReference(
      final String referenceType, final io.opentracing.SpanContext referenced) {
    synchronized (this) {
      if (referenceType == null
          || !(referenced instanceof SpanContext)
          || (!References.CHILD_OF.equals(referenceType)
              && !References.FOLLOWS_FROM.equals(referenceType))) {
        return this;
      }

      // Optimization when there is only one parent.
      if (references.isEmpty()) {
        references =
            Collections.singletonList(new SpanReference((SpanContext) referenced, referenceType));
      } else {
        if (references.size() == 1) {
          references = new ArrayList<>(references);
        }
        references.add(new SpanReference((SpanContext) referenced, referenceType));
      }
    }

    return this;
  }

  @Override
  public Tracer.SpanBuilder ignoreActiveSpan() {
    synchronized (this) {
      ignoreActiveSpan = true;
    }

    return this;
  }

  @Override
  public Tracer.SpanBuilder withTag(final String key, final String value) {
    synchronized (this) {
      tags.put(key, value);
    }

    return this;
  }

  @Override
  public Tracer.SpanBuilder withTag(final String key, final boolean value) {
    synchronized (this) {
      tags.put(key, value);
    }

    return this;
  }

  @Override
  public Tracer.SpanBuilder withTag(final String key, final Number value) {
    synchronized (this) {
      tags.put(key, value);
    }

    return this;
  }

  @Override
  public  Tracer.SpanBuilder withTag(final Tag tag, final T value) {
    synchronized (this) {
      tags.put(tag.getKey(), value);
    }

    return this;
  }

  @Override
  public Tracer.SpanBuilder withStartTimestamp(final long microseconds) {
    startTimeMicroseconds = microseconds;
    startTimeNanoseconds = microseconds * 1000;
    return this;
  }

  @Override
  public io.opentracing.Span start() {
    synchronized (this) {

      // Check if active span should be established as CHILD_OF relationship
      if (!ignoreActiveSpan && references.isEmpty() && scopeTracer.activeSpan() != null) {
        asChildOf(scopeTracer.activeSpan());
      }

      // If there is no previous references, new context is created.
      final SpanContext context =
          references.isEmpty() || !references.get(0).getSpanContext().hasTrace()
              ? createNewContext()
              : createChildContext();

      boolean computeDurationViaNanoTicks = false;

      if (startTimeMicroseconds == 0) {
        startTimeMicroseconds = scopeTracer.clock().currentTimeMicros();

        if (!scopeTracer.clock().isMicrosAccurate()) {
          startTimeNanoseconds = scopeTracer.clock().currentNanoTicks();
          computeDurationViaNanoTicks = true;
        }
      }

      final Span span =
          new Span(
              scopeTracer,
              operationName,
              context,
              startTimeMicroseconds,
              startTimeNanoseconds,
              computeDurationViaNanoTicks,
              tags,
              references);

      scopeTracer.registerStartedSpan(span);
      Statistics.INSTANCE.registerStartedSpan(span);
      return span;
    }
  }

  private SpanContext createNewContext() {
    // Random long
    final TraceId spanId = TraceId.generateRandomId();
    final TraceId traceId = spanId;

    return new SpanContext(traceId, spanId, null, getBaggage());
  }

  private SpanContext createChildContext() {
    final SpanReference preferredReference = preferredReference();
    final TraceId parentTraceId = preferredReference.getSpanContext().getTraceId();
    final TraceId spanId = TraceId.generateRandomId();
    final TraceId parentSpanId = preferredReference.getSpanContext().getSpanId();

    return new SpanContext(parentTraceId, spanId, parentSpanId, getBaggage());
  }

  private SpanReference preferredReference() {
    final SpanReference preferredReference = references.get(0);

    for (final SpanReference reference : references) {
      if (References.CHILD_OF.equals(reference.getType())
          && !References.CHILD_OF.equals(preferredReference.getType())) {
        return reference;
      }
    }

    return preferredReference;
  }

  private Map getBaggage() {
    Map baggage = null;

    if (references.size() == 1) {
      return references.get(0).getSpanContext().getBaggage();
    }

    for (final SpanReference reference : references) {
      if (reference.getSpanContext().getBaggage() != null) {
        if (baggage == null) {
          baggage = new HashMap<>();
        }
        baggage.putAll(reference.getSpanContext().getBaggage());
      }
    }

    return baggage;
  }
}