private void doWork(String spanName, int i) {
  try (Scope scope = tracer.spanBuilder(spanName).startScopedSpan()) {
    // Simulate some work.
    Span span = tracer.getCurrentSpan();

    try {
      Thread.sleep(10L);
    } catch (InterruptedException e) {
      span.setStatus(Status.INTERNAL.withDescription(e.toString()));
    }

    Map<String, AttributeValue> attributes = new HashMap<String, AttributeValue>();
    attributes.put("inner work iteration number", AttributeValue.longAttributeValue(i));
    span.addAnnotation("Invoking doWork", attributes);
  }
}
 

public boolean acquire(K key, final Runnable cacheHit) {
    try (Scope ss = tracer.spanBuilder("DefaultRateLimitedCache").startScopedSpan()) {
        Span span = tracer.getCurrentSpan();

        if (cache.get(key) != null) {
            cacheHit.run();
            span.addAnnotation("Found key in cache");
            return false;
        }

        span.addAnnotation("Acquiring rate limiter");
        rateLimiter.acquire();
        span.addAnnotation("Acquired rate limiter");

        if (cache.putIfAbsent(key, true) != null) {
            cacheHit.run();
            return false;
        }

        return true;
    }
}
 

private BulkMutation getOrAddBulkMutation(String tableName) {
    final Span span = tracer.spanBuilder("BigtableMutator.getOrAddBulkMutation").startSpan();
    try (Scope ws = tracer.withSpan(span)) {
        span.addAnnotation("Acquiring lock");
        synchronized (tableAccessLock) {
            span.addAnnotation("Lock acquired");

            if (tableToBulkMutation.containsKey(tableName)) {
                span.setStatus(Status.ALREADY_EXISTS.withDescription("Mutation exists in map"));
                span.end();
                return tableToBulkMutation.get(tableName);
            }

            final BulkMutation bulkMutation = session.createBulkMutation(
                session
                    .getOptions()
                    .getInstanceName()
                    .toTableName(tableName));

            tableToBulkMutation.put(tableName, bulkMutation);

            span.end();
            return bulkMutation;
        }
    }
}
 

private <T> AsyncFuture<T> run(
    final Function<ClusterNode.Group, AsyncFuture<T>> function,
    final Function<ClusterShard, Transform<Throwable, T>> catcher,
    final Collector<T, T> collector
) {
    final List<AsyncFuture<T>> futures = new ArrayList<>(shards.size());

    for (final ClusterShard shard : shards) {

        final Span span = tracer
            .spanBuilder("CoreQueryManager.run")
            .startSpan();

        span.putAttribute("shard", stringAttributeValue(shard.toString()));

        futures.add(shard
            .apply(function::apply, CoreQueryManager::retryTraceHandlerNoop)
            .catchFailed(catcher.apply(shard))
            .onFinished(span::end));
    }

    return async.collect(futures, collector);
}
 

@Test
public void testSpanStackCloseSpanFunction() {
  final org.springframework.cloud.sleuth.Span[] sleuthSpans = createSleuthSpans(4);
  // push all the spans
  for (int i = 0; i < sleuthSpans.length; i++) {
    OpenCensusSleuthSpanContextHolder.push(sleuthSpans[i], /* autoClose= */ false);
  }
  // pop all the spans, verify that given SpanFunction is called on the closed span.
  for (int i = sleuthSpans.length - 1; i >= 0; i--) {
    final int index = i;
    OpenCensusSleuthSpanContextHolder.close(
        new OpenCensusSleuthSpanContextHolder.SpanFunction() {
          @Override
          public void apply(org.springframework.cloud.sleuth.Span span) {
            assertThat(span).isEqualTo(sleuthSpans[index]);
          }
        });
  }
}
 

@Test
public void probabilitySampler_DifferentProbabilities_NotSampledParent() {
  final Sampler neverSample = Samplers.probabilitySampler(0.0);
  assertSamplerSamplesWithProbability(
      neverSample, notSampledSpanContext, Collections.<Span>emptyList(), 0.0);
  final Sampler alwaysSample = Samplers.probabilitySampler(1.0);
  assertSamplerSamplesWithProbability(
      alwaysSample, notSampledSpanContext, Collections.<Span>emptyList(), 1.0);
  final Sampler fiftyPercentSample = Samplers.probabilitySampler(0.5);
  assertSamplerSamplesWithProbability(
      fiftyPercentSample, notSampledSpanContext, Collections.<Span>emptyList(), 0.5);
  final Sampler twentyPercentSample = Samplers.probabilitySampler(0.2);
  assertSamplerSamplesWithProbability(
      twentyPercentSample, notSampledSpanContext, Collections.<Span>emptyList(), 0.2);
  final Sampler twoThirdsSample = Samplers.probabilitySampler(2.0 / 3.0);
  assertSamplerSamplesWithProbability(
      twoThirdsSample, notSampledSpanContext, Collections.<Span>emptyList(), 2.0 / 3.0);
}
 

@Test
public void probabilitySampler_DifferentProbabilities_SampledParent() {
  final Sampler neverSample = Samplers.probabilitySampler(0.0);
  assertSamplerSamplesWithProbability(
      neverSample, sampledSpanContext, Collections.<Span>emptyList(), 1.0);
  final Sampler alwaysSample = Samplers.probabilitySampler(1.0);
  assertSamplerSamplesWithProbability(
      alwaysSample, sampledSpanContext, Collections.<Span>emptyList(), 1.0);
  final Sampler fiftyPercentSample = Samplers.probabilitySampler(0.5);
  assertSamplerSamplesWithProbability(
      fiftyPercentSample, sampledSpanContext, Collections.<Span>emptyList(), 1.0);
  final Sampler twentyPercentSample = Samplers.probabilitySampler(0.2);
  assertSamplerSamplesWithProbability(
      twentyPercentSample, sampledSpanContext, Collections.<Span>emptyList(), 1.0);
  final Sampler twoThirdsSample = Samplers.probabilitySampler(2.0 / 3.0);
  assertSamplerSamplesWithProbability(
      twoThirdsSample, sampledSpanContext, Collections.<Span>emptyList(), 1.0);
}
 

@Test
public void startChildSpan_SampledLinkedParent() {
  Span rootSpanUnsampled =
      SpanBuilderImpl.createWithParent(SPAN_NAME, null, spanBuilderOptions)
          .setSampler(Samplers.neverSample())
          .startSpan();
  assertThat(rootSpanUnsampled.getContext().getTraceOptions().isSampled()).isFalse();
  Span rootSpanSampled =
      SpanBuilderImpl.createWithParent(SPAN_NAME, null, spanBuilderOptions)
          .setSampler(Samplers.alwaysSample())
          .startSpan();
  assertThat(rootSpanSampled.getContext().getTraceOptions().isSampled()).isTrue();
  // Sampled because the linked parent is sampled.
  Span childSpan =
      SpanBuilderImpl.createWithParent(SPAN_NAME, rootSpanUnsampled, spanBuilderOptions)
          .setParentLinks(Collections.singletonList(rootSpanSampled))
          .startSpan();
  assertThat(childSpan.getContext().isValid()).isTrue();
  assertThat(childSpan.getContext().getTraceId())
      .isEqualTo(rootSpanUnsampled.getContext().getTraceId());
  assertThat(childSpan.getContext().getTraceOptions().isSampled()).isTrue();
}
 

@Test
public void startChildSpan_WithSpecifiedSampler() {
  Span rootSpan =
      SpanBuilderImpl.createWithParent(SPAN_NAME, null, spanBuilderOptions)
          .setSampler(Samplers.alwaysSample())
          .startSpan();
  assertThat(rootSpan.getContext().isValid()).isTrue();
  assertThat(rootSpan.getContext().getTraceOptions().isSampled()).isTrue();
  // Apply the given sampler for child spans.
  Span childSpan =
      SpanBuilderImpl.createWithParent(SPAN_NAME, rootSpan, spanBuilderOptions)
          .setSampler(Samplers.neverSample())
          .startSpan();
  assertThat(childSpan.getContext().isValid()).isTrue();
  assertThat(childSpan.getContext().getTraceId()).isEqualTo(rootSpan.getContext().getTraceId());
  assertThat(childSpan.getContext().getTraceOptions().isSampled()).isFalse();
}
 

@Test
public void testResponseFilter() throws Exception {
  Span span = new FakeSpan(SpanContext.INVALID, null);
  TagContext tagContext = mock(TagContext.class);

  HttpRequestContext context = createHttpRequestContext(span, tagContext);

  ClientRequestContext requestContext = mock(ClientRequestContext.class);
  when(requestContext.getProperty("opencensus.context")).thenReturn(context);

  ClientResponseContext responseContext = mock(ClientResponseContext.class);

  filter.filter(requestContext, responseContext);

  verify(requestContext).getProperty("opencensus.context");
  verify(responseContext, times(1)).getStatus();
}
 

@Test
public void testSpanStackSimple() {
  org.springframework.cloud.sleuth.Span[] sleuthSpans = createSleuthSpans(4);
  // push all the spans
  for (int i = 0; i < sleuthSpans.length; i++) {
    OpenCensusSleuthSpanContextHolder.push(sleuthSpans[i], /* autoClose= */ false);
    assertThat(OpenCensusSleuthSpanContextHolder.getCurrentSpan()).isEqualTo(sleuthSpans[i]);
    assertSpanEquals(tracer.getCurrentSpan(), sleuthSpans[i]);
  }
  // pop all the spans
  for (int i = sleuthSpans.length - 1; i >= 0; i--) {
    assertThat(OpenCensusSleuthSpanContextHolder.getCurrentSpan()).isEqualTo(sleuthSpans[i]);
    assertSpanEquals(tracer.getCurrentSpan(), sleuthSpans[i]);
    OpenCensusSleuthSpanContextHolder.close();
  }
}
 

static Object proceed(
    ProceedingJoinPoint call, Tracer tracer, String spanName, String... annotations)
    throws Throwable {
  Scope scope = tracer.spanBuilder(spanName).startScopedSpan();
  try {
    for (String annotation : annotations) {
      tracer.getCurrentSpan().addAnnotation(annotation);
    }

    return call.proceed();

  } catch (Throwable t) {
    Map<String, AttributeValue> attributes = new HashMap<String, AttributeValue>();
    String message = t.getMessage();
    attributes.put(
        "message", AttributeValue.stringAttributeValue(message == null ? "null" : message));
    attributes.put("type", AttributeValue.stringAttributeValue(t.getClass().toString()));

    Span span = tracer.getCurrentSpan();
    span.addAnnotation("error", attributes);
    span.setStatus(Status.UNKNOWN);
    throw t;
  } finally {
    scope.close();
  }
}
 

@Override
public Span startSpan() {
  if (remoteParentSpanContext != null) {
    return startSpanInternal(
        remoteParentSpanContext,
        Boolean.TRUE,
        name,
        sampler,
        parentLinks,
        recordEvents,
        kind,
        null);
  } else {
    // This is not a child of a remote Span. Get the parent SpanContext from the parent Span if
    // any.
    SpanContext parentContext = null;
    Boolean hasRemoteParent = null;
    if (parent != null) {
      parentContext = parent.getContext();
      hasRemoteParent = Boolean.FALSE;
    }
    return startSpanInternal(
        parentContext, hasRemoteParent, name, sampler, parentLinks, recordEvents, kind, parent);
  }
}
 

@Test
public void startChildSpan() {
  Span rootSpan =
      SpanBuilderImpl.createWithParent(SPAN_NAME, null, spanBuilderOptions).startSpan();
  assertThat(rootSpan.getContext().isValid()).isTrue();
  assertThat(rootSpan.getOptions().contains(Options.RECORD_EVENTS)).isTrue();
  assertThat(rootSpan.getContext().getTraceOptions().isSampled()).isTrue();
  assertThat(((RecordEventsSpanImpl) rootSpan).toSpanData().getHasRemoteParent()).isNull();
  Span childSpan =
      SpanBuilderImpl.createWithParent(SPAN_NAME, rootSpan, spanBuilderOptions).startSpan();
  assertThat(childSpan.getContext().isValid()).isTrue();
  assertThat(childSpan.getContext().getTraceId()).isEqualTo(rootSpan.getContext().getTraceId());
  assertThat(((RecordEventsSpanImpl) childSpan).toSpanData().getParentSpanId())
      .isEqualTo(rootSpan.getContext().getSpanId());
  assertThat(((RecordEventsSpanImpl) childSpan).toSpanData().getHasRemoteParent()).isFalse();
  assertThat(((RecordEventsSpanImpl) childSpan).getTimestampConverter())
      .isEqualTo(((RecordEventsSpanImpl) rootSpan).getTimestampConverter());
}
 

@Before
@SuppressWarnings("unchecked")
public void setUp() throws Exception {
  when(spyClientSpanBuilder.startSpan()).thenReturn(spyClientSpan);
  when(tracer.spanBuilderWithExplicitParent(anyString(), ArgumentMatchers.<Span>any()))
      .thenReturn(spyClientSpanBuilder);
  when(spyServerSpanBuilder.startSpan()).thenReturn(spyServerSpan);
  when(tracer.spanBuilderWithRemoteParent(anyString(), ArgumentMatchers.<SpanContext>any()))
      .thenReturn(spyServerSpanBuilder);
  when(mockTracingPropagationHandler.toByteArray(any(SpanContext.class)))
      .thenReturn(binarySpanContext);
  when(mockTracingPropagationHandler.fromByteArray(any(byte[].class)))
      .thenReturn(fakeClientSpanContext);
  censusStats =
      new CensusStatsModule(
          tagger, tagCtxSerializer, statsRecorder, fakeClock.getStopwatchSupplier(),
          true, true, true, false /* real-time */);
  censusTracing = new CensusTracingModule(tracer, mockTracingPropagationHandler);
}
 

private Transform(
    final FullQuery.Request request,
    final boolean failOnLimits,
    final OptionalLimit seriesLimit,
    final OptionalLimit groupLimit,
    final QuotaWatcher quotaWatcher,
    final DataInMemoryReporter dataInMemoryReporter,
    final Span parentSpan
) {
    this.aggregation = request.aggregation();
    this.range = request.range();
    this.options = request.options();
    this.source = request.source();

    this.failOnLimits = failOnLimits;
    this.seriesLimit = seriesLimit;
    this.groupLimit = groupLimit;

    this.namedWatch = QueryTrace.watch(QUERY);
    this.quotaWatcher = quotaWatcher;

    this.dataInMemoryReporter = dataInMemoryReporter;
    this.parentSpan = parentSpan;

    final Features features = request.features();
    this.bucketStrategy = options
        .bucketStrategy()
        .orElseGet(
            () -> features.withFeature(Feature.END_BUCKET, () -> BucketStrategy.END,
                () -> BucketStrategy.START));
}
 

private static void recordMessageEvent(
    Span span, MessageEvent.Type type,
    int seqNo, long optionalWireSize, long optionalUncompressedSize) {
  MessageEvent.Builder eventBuilder = MessageEvent.builder(type, seqNo);
  if (optionalUncompressedSize != -1) {
    eventBuilder.setUncompressedMessageSize(optionalUncompressedSize);
  }
  if (optionalWireSize != -1) {
    eventBuilder.setCompressedMessageSize(optionalWireSize);
  }
  span.addMessageEvent(eventBuilder.build());
}
 

final void addSpanRequestAttributes(Span span, Q request, HttpExtractor<Q, P> extractor) {
  putAttributeIfNotEmptyOrNull(
      span, HttpTraceAttributeConstants.HTTP_USER_AGENT, extractor.getUserAgent(request));
  putAttributeIfNotEmptyOrNull(
      span, HttpTraceAttributeConstants.HTTP_HOST, extractor.getHost(request));
  putAttributeIfNotEmptyOrNull(
      span, HttpTraceAttributeConstants.HTTP_METHOD, extractor.getMethod(request));
  putAttributeIfNotEmptyOrNull(
      span, HttpTraceAttributeConstants.HTTP_PATH, extractor.getPath(request));
  putAttributeIfNotEmptyOrNull(
      span, HttpTraceAttributeConstants.HTTP_ROUTE, extractor.getRoute(request));
  putAttributeIfNotEmptyOrNull(
      span, HttpTraceAttributeConstants.HTTP_URL, extractor.getUrl(request));
}
 

private AsyncFuture<WriteMetric> writeTyped(
    final Series series,
    final BigtableDataClient client,
    final MetricCollection g,
    final Span parentSpan
) throws IOException {
    switch (g.getType()) {
        case POINT:
            return writeBatch(POINTS, series, client, g.getDataAs(Point.class),
                d -> serializeValue(d.getValue()), parentSpan);
        default:
            return async.resolved(new WriteMetric(
                new QueryError("Unsupported metric type: " + g.getType())));
    }
}
 

@Override
public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(
    final ServerCall<ReqT, RespT> call, final Metadata headers,
    final ServerCallHandler<ReqT, RespT> next) {

    final Span span = tracer.getCurrentSpan();
    span.putAttributes(environmentMetadata.toAttributes());
    span.putAttribute("protocol", stringAttributeValue("grpc"));
    span.putAttribute("span.kind", stringAttributeValue("server"));

    tracingConfig.getTags().forEach(
        (key, value) -> span.putAttribute(key, stringAttributeValue(value)));

    return next.startCall(call, headers);
}
 

private static void doWork() {
  Span rootSpan = tracer.spanBuilderWithExplicitParent("MyRootSpan", null).startSpan();
  rootSpan.addAnnotation("Annotation to the root Span before child is created.");
  Span childSpan = tracer.spanBuilderWithExplicitParent("MyChildSpan", rootSpan).startSpan();
  childSpan.addAnnotation("Annotation to the child Span");
  childSpan.end();
  rootSpan.addAnnotation("Annotation to the root Span after child is ended.");
  rootSpan.end();
}
 

private static void doSomeMoreWork() {
  // Create a child Span of the current Span.
  Span span = tracer.spanBuilder("MyChildSpan").startSpan();
  try (Scope ws = tracer.withSpan(span)) {
    doSomeOtherWork();
  }
  span.end();
}
 

org.springframework.cloud.sleuth.Span[] createSleuthSpans(int len) {
  org.springframework.cloud.sleuth.Span[] spans = new org.springframework.cloud.sleuth.Span[len];
  for (int i = 0; i < len; i++) {
    spans[i] = createSleuthSpan(i * 10 + 1, i * 10 + 2, i * 10 + 3, /* exportable= */ true);
  }
  return spans;
}
 

private static void performWork(Span parent) {
  SpanBuilder spanBuilder =
      tracer
          .spanBuilderWithExplicitParent("internal_work", parent)
          .setRecordEvents(true)
          .setSampler(Samplers.alwaysSample());
  try (Scope scope = spanBuilder.startScopedSpan()) {
    Span span = tracer.getCurrentSpan();
    span.putAttribute("my_attribute", AttributeValue.stringAttributeValue("blue"));
    span.addAnnotation("Performing work.");
    sleepFor(20); // Working hard here.
    span.addAnnotation("Done work.");
  }
}
 

private static void sleepFor(int milliseconds) {
  try {
    Thread.sleep(milliseconds);
  } catch (InterruptedException e) {
    Span span = tracer.getCurrentSpan();
    span.addAnnotation("Exception thrown when performing work " + e.getMessage());
    span.setStatus(Status.UNKNOWN);
  }
}
 

private static LogEntry getEnhancedLogEntry(LoggingEnhancer loggingEnhancer, Span span) {
  Scope scope = tracer.withSpan(span);
  try {
    LogEntry.Builder builder = LogEntry.newBuilder(null);
    loggingEnhancer.enhanceLogEntry(builder);
    return builder.build();
  } finally {
    scope.close();
  }
}
 

@Override
public final boolean shouldSample(
    @Nullable SpanContext parentContext,
    @Nullable Boolean hasRemoteParent,
    TraceId traceId,
    SpanId spanId,
    String name,
    @Nullable List<Span> parentLinks) {
  // If the parent is sampled keep the sampling decision.
  if (parentContext != null && parentContext.getTraceOptions().isSampled()) {
    return true;
  }
  if (parentLinks != null) {
    // If any parent link is sampled keep the sampling decision.
    for (Span parentLink : parentLinks) {
      if (parentLink.getContext().getTraceOptions().isSampled()) {
        return true;
      }
    }
  }
  // Always sample if we are within probability range. This is true even for child spans (that
  // may have had a different sampling decision made) to allow for different sampling policies,
  // and dynamic increases to sampling probabilities for debugging purposes.
  // Note use of '<' for comparison. This ensures that we never sample for probability == 0.0,
  // while allowing for a (very) small chance of *not* sampling if the id == Long.MAX_VALUE.
  // This is considered a reasonable tradeoff for the simplicity/performance requirements (this
  // code is executed in-line for every Span creation).
  return Math.abs(traceId.getLowerLong()) < getIdUpperBound();
}
 

@Override
public boolean shouldSample(
    @Nullable SpanContext parentContext,
    @Nullable Boolean hasRemoteParent,
    TraceId traceId,
    SpanId spanId,
    String name,
    List<Span> parentLinks) {
  return false;
}
 

@Test
public void testFromSleuthUnsampled() {
  org.springframework.cloud.sleuth.Span sleuthSpan =
      createSleuthSpan(21, 22, 23, /* exportable= */ false);
  OpenCensusSleuthSpanContextHolder.setCurrentSpan(sleuthSpan);
  assertThat(OpenCensusSleuthSpanContextHolder.isTracing()).isTrue();
  assertThat(OpenCensusSleuthSpanContextHolder.getCurrentSpan()).isEqualTo(sleuthSpan);
  assertSpanEquals(tracer.getCurrentSpan(), sleuthSpan);
  assertThat(tracer.getCurrentSpan().getContext().getTraceOptions().isSampled()).isFalse();
  OpenCensusSleuthSpanContextHolder.close();
}
 

@Override
public AsyncFuture<WriteMetric> write(
    final WriteMetric.Request request, final Span parentSpan
) {
    return connection.doto(c -> {
        final Series series = request.getSeries();
        final List<AsyncFuture<WriteMetric>> results = new ArrayList<>();

        final BigtableDataClient client = c.getDataClient();

        final MetricCollection g = request.getData();
        results.add(writeTyped(series, client, g, parentSpan));
        return async.collect(results, WriteMetric.reduce());
    });
}