Java Code Examples for brave.propagation.TraceContextOrSamplingFlags#context()

/**
 * Returns span from TraceContext or dubbo attachments.
 */
public Span nextSpan(TraceContextOrSamplingFlags extracted) {
  TraceContext parent = currentTraceContext.get();
  //If spanInScope is closed we can use dubbo attachments
  if (parent == null) {
    parent = extracted.context();
  }
  if (parent != null) {
    return tracer.newChild(parent);
  }
  return tracer.newTrace(SamplingFlags.NOT_SAMPLED);
}
 

/** Creates a potentially noop span representing this request */
Span nextSpan(TraceContextOrSamplingFlags extracted, HttpServerRequest request) {
  Boolean sampled = extracted.sampled();
  // only recreate the context if the http sampler made a decision
  if (sampled == null && (sampled = sampler.trySample(request)) != null) {
    extracted = extracted.sampled(sampled.booleanValue());
  }
  return extracted.context() != null
    ? tracer.joinSpan(extracted.context())
    : tracer.nextSpan(extracted);
}
 

protected TraceContext extract(RecordedRequest request) {
  TraceContextOrSamplingFlags extracted = extractor.extract(request);
  assertThat(extracted.context())
    .withFailMessage("Expected to extract a trace context from %s", request.getHeaders())
    .isNotNull();
  return extracted.context();
}
 

/** Creates a potentially noop span representing this request */
Span nextSpan(TraceContextOrSamplingFlags extracted, RpcServerRequest request) {
  Boolean sampled = extracted.sampled();
  // only recreate the context if the RPC sampler made a decision
  if (sampled == null && (sampled = sampler.trySample(request)) != null) {
    extracted = extracted.sampled(sampled.booleanValue());
  }
  return extracted.context() != null
      ? tracer.joinSpan(extracted.context())
      : tracer.nextSpan(extracted);
}
 

/**
 * Use this to create a span for processing the given record. Note: the result has no name and is
 * not started.
 *
 * <p>This creates a child from identifiers extracted from the record headers, or a new span if
 * one couldn't be extracted.
 */
public Span nextSpan(ConsumerRecord<?, ?> record) {
  // Eventhough the type is ConsumerRecord, this is not a (remote) consumer span. Only "poll"
  // events create consumer spans. Since this is a processor span, we use the normal sampler.
  TraceContextOrSamplingFlags extracted =
    extractAndClearTraceIdHeaders(processorExtractor, record.headers(), record.headers());
  Span result = tracer.nextSpan(extracted);
  if (extracted.context() == null && !result.isNoop()) {
    addTags(record, result);
  }
  return result;
}
 

void verifyRoundTrip(TraceContextOrSamplingFlags expected) {
  TraceContextOrSamplingFlags extracted = propagation.extractor(mapEntry).extract(map);

  assertThat(extracted)
    .isEqualTo(expected);

  Map<String, String> injected = new LinkedHashMap<>();
  if (expected.context() != null) {
    propagation.injector(mapEntry).inject(expected.context(), injected);
  } else {
    inject(injected, expected.samplingFlags());
  }

  assertThat(map).isEqualTo(injected);
}
 

/**
 * Use this to create a span for processing the given message. Note: the result has no
 * name and is not started. This creates a child from identifiers extracted from the
 * message headers, or a new span if one couldn't be extracted.
 * @param message message to use for span creation
 * @return span to be created
 */
public Span nextSpan(Message<?> message) {
	MessageHeaderAccessor headers = mutableHeaderAccessor(message);
	TraceContextOrSamplingFlags extracted = this.extractor.extract(headers);
	headers.setImmutable();
	Span result = this.tracer.nextSpan(extracted);
	if (extracted.context() == null && !result.isNoop()) {
		addTags(message, result, null);
	}
	if (log.isDebugEnabled()) {
		log.debug("Created a new span " + result);
	}
	return result;
}
 

/**
 * Resumes or forks
 */
private void startSpan(ActivityRuntimeContext context) {
    Tracing tracing = TraceHelper.getTracing("mdw-activity");
    Tracer tracer = tracing.tracer();
    Span span = tracer.currentSpan();
    if (span == null) {
        // async brave server if b3 requestHeaders populated (subspan)
        Object requestHeaders = context.getValue("requestHeaders");
        if (requestHeaders instanceof Map) {
            Map<?, ?> headers = (Map<?, ?>) requestHeaders;
            if (headers.containsKey("x-b3-traceid")) {
                TraceContext.Extractor<Map<?, ?>> extractor =
                        tracing.propagation().extractor((map, key) -> {
                            Object val = map.get(key.toLowerCase());
                            return val == null ? null : val.toString();
                        });
                TraceContextOrSamplingFlags extracted = extractor.extract(headers);
                span = extracted.context() != null
                        ? tracer.joinSpan(extracted.context())
                        : tracer.nextSpan(extracted);
                span.name("m" + context.getMasterRequestId()).kind(Span.Kind.SERVER);
                span = tracer.newChild(span.context()).name("p" + context.getProcessInstanceId());
            }
        }
    }
    if (span == null) {
        // create a new root span for async start
        span = tracer.nextSpan().name("a" + context.getActivityInstanceId());
    }
    span.start().flush(); // async
    Tracer.SpanInScope spanInScope = tracer.withSpanInScope(span);
    context.getRuntimeAttributes().put(Tracer.SpanInScope.class, spanInScope);
}
 

Span nextSpan(TraceContextOrSamplingFlags extracted) {
  if (extracted.sampled() == null) {
    extracted = extracted.sampled(false);
  }
  return extracted.context() != null
      ? tracer.joinSpan(extracted.context())
      : tracer.nextSpan(extracted);
}
 

@Override public String getValue(BaggageField field, TraceContextOrSamplingFlags extracted) {
  if (extracted.context() != null) return getValue(field, extracted.context());
  if (extracted.traceIdContext() != null) return extracted.traceIdContext().traceIdString();
  return null;
}
 

@Override public String getValue(BaggageField field, TraceContextOrSamplingFlags extracted) {
  if (extracted.context() != null) return getValue(field, extracted.context());
  return null;
}
 

static BraveSpanContext create(TraceContextOrSamplingFlags extractionResult) {
  return extractionResult.context() != null
      ? new BraveSpanContext.Complete(extractionResult.context())
      : new BraveSpanContext.Incomplete(extractionResult);
}
 

public static List<BaggageField> getAllFields(TraceContextOrSamplingFlags extracted) {
  if (extracted.context() != null) return getAllFields(extracted.context());
  return getAllFields(extracted.extra());
}
 

public static Map<String, String> getAllValues(TraceContextOrSamplingFlags extracted) {
  if (extracted.context() != null) return getAllValues(extracted.context());
  return getAllValues(extracted.extra());
}
 

@Nullable
public static BaggageField getFieldByName(TraceContextOrSamplingFlags extracted, String name) {
  if (extracted.context() != null) return getFieldByName(extracted.context(), name);
  return getFieldByName(getAllFields(extracted.extra()), name);
}
 

@Override public String getValue(BaggageField field, TraceContextOrSamplingFlags extracted) {
  if (extracted.context() != null) return getValue(field, extracted.context());
  return getValue(field, extracted.extra());
}
 

@Override public boolean updateValue(BaggageField field, TraceContextOrSamplingFlags extracted,
  @Nullable String value) {
  if (extracted.context() != null) return updateValue(field, extracted.context(), value);
  return updateValue(field, extracted.extra(), value);
}
 

static BraveSpanContext create(TraceContextOrSamplingFlags extractionResult) {
  return extractionResult.context() != null
    ? new BraveSpanContext(extractionResult.context())
    : new BraveSpanContext.Incomplete(extractionResult);
}
 

Incomplete(TraceContextOrSamplingFlags extractionResult) {
  super(extractionResult.context());
  this.extractionResult = extractionResult;
}
 

@Override
public void onNext(StreamingPullResponse value) {
  if (ctx == null) {
    ctx = RequestContext.current();
  }

  streamReconnectBackoff = INITIAL_CHANNEL_RECONNECT_BACKOFF;

  receivedMessages.increment(value.getReceivedMessagesCount());

  AtomicInteger pendingAcks = new AtomicInteger(value.getReceivedMessagesCount());

  for (ReceivedMessage message : value.getReceivedMessagesList()) {
    TraceContextOrSamplingFlags contextOrFlags = traceExtractor.extract(message.getMessage());

    // Add an artificial span modeling the time spent within Pub/Sub until getting here.
    Span span =
        contextOrFlags.context() != null
            ? tracer.joinSpan(contextOrFlags.context())
            // We want each message to be a new trace rather than having a long trace for the
            // entire stream.
            : tracer.newTrace();

    span.kind(Kind.SERVER)
        .name("google.pubsub.v1.Publisher.Publish")
        .tag("subscription", options.getSubscription())
        .start(Timestamps.toMicros(message.getMessage().getPublishTime()))
        .finish();

    StreamObserver<StreamingPullRequest> requestObserver = this.requestObserver;

    long startTimeNanos = System.nanoTime();
    options
        .getMessageReceiver()
        .receiveMessage(
            message.getMessage(),
            new AckReplyConsumer() {
              @Override
              public void ack() {
                releaseAndRecord();

                ackedMessages.increment();

                checkNotNull(requestObserver, "onNext called before start()");
                requestObserver.onNext(
                    StreamingPullRequest.newBuilder().addAckIds(message.getAckId()).build());
              }

              @Override
              public void nack() {
                releaseAndRecord();

                nackedMessages.increment();

                checkNotNull(requestObserver, "onNext called before start()");
                requestObserver.onNext(
                    StreamingPullRequest.newBuilder()
                        .addModifyDeadlineAckIds(message.getAckId())
                        .addModifyDeadlineSeconds(0)
                        .build());
              }

              private void releaseAndRecord() {
                if (options.getUnsafeWrapBuffers() && pendingAcks.decrementAndGet() == 0) {
                  GrpcUnsafeBufferUtil.releaseBuffer(value, ctx);
                }

                messageProcessingTime.record(
                    Duration.ofNanos(System.nanoTime() - startTimeNanos));
              }
            });
  }
}