brave.propagation.TraceContext.Injector Java Examples

/**
 * Injects the underlying context using B3 encoding by default.
 */
@Override public <C> void inject(SpanContext spanContext, Format<C> format, C carrier) {
  BraveSpanContext braveContext = ((BraveSpanContext) spanContext);
  if (carrier instanceof BinaryInject) {
    BinaryCodec.INSTANCE.inject(braveContext.unwrap(), (BinaryInject) carrier);
    return;
  }
  if (!(carrier instanceof TextMapInject)) {
    throw new UnsupportedOperationException(carrier + " not instanceof TextMapInject");
  }
  Kind kind = braveContext.kind;
  Injector<TextMapInject> injector = null;
  if (Kind.CLIENT.equals(kind)) {
    injector = formatToClientInjector.get(format);
  } else if (Kind.PRODUCER.equals(kind)) {
    injector = formatToProducerInjector.get(format);
  } else if (Kind.CONSUMER.equals(kind)) {
    injector = formatToConsumerInjector.get(format);
  }
  if (injector == null) injector = formatToInjector.get(format);
  if (injector == null) {
    throw new UnsupportedOperationException(format + " not in " + formatToInjector.keySet());
  }
  injector.inject(braveContext.unwrap(), (TextMapInject) carrier);
}
 

@Test public void twoFunctions_injects_remote() {
  when(remoteSetter.spanKind()).thenReturn(Kind.CLIENT);
  RemoteInjector<Request> remoteInjector =
      (RemoteInjector<Request>) twoFunctions.newInjector(remoteSetter);

  for (Kind kind : injectableKinds) {
    when(remoteSetter.spanKind()).thenReturn(kind);

    Injector<Request> nextRemoteInjector = twoFunctions.newInjector(remoteSetter);
    assertThat(nextRemoteInjector).isEqualTo(remoteInjector);
    assertThat(nextRemoteInjector).hasSameHashCodeAs(remoteInjector);
    assertThat(nextRemoteInjector).hasToString(remoteInjector.toString());

    assertThat(nextRemoteInjector).extracting("injectorFunction")
        .isEqualTo(remoteInjector.injectorFunction);

    nextRemoteInjector.inject(context, request);
    assertThat(oneCount.getAndSet(0)).isOne();
    assertThat(twoCount.getAndSet(0)).isOne();
  }
}
 

@Test public void oneFunction_injects_remote() {
  when(remoteSetter.spanKind()).thenReturn(Kind.CLIENT);
  RemoteInjector<Request> remoteInjector =
      (RemoteInjector<Request>) oneFunction.newInjector(remoteSetter);

  for (Kind kind : injectableKinds) {
    when(remoteSetter.spanKind()).thenReturn(kind);

    Injector<Request> nextRemoteInjector = oneFunction.newInjector(remoteSetter);
    assertThat(nextRemoteInjector).isEqualTo(remoteInjector);
    assertThat(nextRemoteInjector).hasSameHashCodeAs(remoteInjector);
    assertThat(nextRemoteInjector).hasToString(remoteInjector.toString());

    assertThat(nextRemoteInjector).extracting("injectorFunction")
        .isEqualTo(remoteInjector.injectorFunction);

    nextRemoteInjector.inject(context, request);
    assertThat(oneCount.getAndSet(0)).isOne();
  }
}
 

@Override public <C> void inject(SpanContext spanContext, Format<C> format, C carrier) {
  if (!(carrier instanceof TextMapInject)) {
    throw new UnsupportedOperationException(carrier + " not instanceof TextMapInject");
  }
  BraveSpanContext braveContext = ((BraveSpanContext) spanContext);
  Kind kind = braveContext.kind;
  Injector<TextMapInject> injector;
  if (Kind.CLIENT.equals(kind)) {
    injector = clientInjector;
  } else if (Kind.PRODUCER.equals(kind)) {
    injector = producerInjector;
  } else if (Kind.CONSUMER.equals(kind)) {
    injector = consumerInjector;
  } else {
    injector = this.injector;
  }
  injector.inject(braveContext.context, (TextMapInject) carrier);
}
 

/** This shows that we can encode arbitrary fields into a single header. */
@Test public void encodes_arbitrary_fields() {
  TraceContext context = factory.decorate(TraceContext.newBuilder().traceId(1).spanId(2).build());
  field1.updateValue(context, "1");
  field2.updateValue(context, "2");
  field3.updateValue(context, "3");

  Injector<Map<String, String>> injector = factory.get().injector(Map::put);
  Map<String, String> headers = new LinkedHashMap<>();
  injector.inject(context, headers);

  assertThat(headers).containsOnly(
      entry("b3", "0000000000000001-0000000000000002"),
      entry("one", "1"), // has its own header config which is still serialized
      entry("baggage", "one=1,three=3") // excluding the blacklist field including the dynamic one
  );
}
 

TracingProducer(Extractor<R> extractor, Injector<R> injector, JmsTracing jmsTracing) {
  this.jmsTracing = jmsTracing;
  this.tracer = jmsTracing.tracing.tracer();
  this.current = jmsTracing.tracing.currentTraceContext();
  this.extractor = extractor;
  this.injector = injector;
  this.sampler = jmsTracing.producerSampler;
  this.remoteServiceName = jmsTracing.remoteServiceName;
}
 

/**
 * @since 4.4
 * @deprecated Since 5.7, use {@link #handleSend(HttpClientRequest)}.
 */
@Deprecated public <C> Span handleSend(Injector<C> injector, C carrier, Req request, Span span) {
  if (request == null) throw new NullPointerException("carrier == null");
  if (span == null) throw new NullPointerException("span == null");

  injector.inject(span.context(), carrier);
  HttpClientRequest clientRequest;
  if (request instanceof HttpClientRequest) {
    clientRequest = (HttpClientRequest) request;
  } else {
    clientRequest = new HttpClientAdapters.FromRequestAdapter<>(adapter, request);
  }

  return handleStart(clientRequest, span);
}
 

void kindBasedFunctions_injects_remote(Kind kind, AtomicInteger kindCallCount) {
  when(remoteSetter.spanKind()).thenReturn(kind);

  Injector<Request> nextRemoteInjector = kindBasedFunctions.newInjector(remoteSetter);
  nextRemoteInjector.inject(context, request);

  for (AtomicInteger callCount : allCounts) {
    if (callCount.equals(kindCallCount)) {
      assertThat(callCount.getAndSet(0)).isOne();
    } else {
      assertThat(callCount.getAndSet(0)).isZero();
    }
  }
}
 

@Test public void producerUsesB3SingleNoParent_deferred() {
  // This injector won't know the type it is injecting until the call to inject()
  Injector<ProducerRequest> injector = Propagation.B3_STRING.injector(ProducerRequest::header);

  ProducerRequest request = new ProducerRequest();
  injector.inject(context, request);

  assertThat(request.headers)
    .hasSize(1)
    .containsEntry("b3", "0000000000000001-0000000000000003");
}
 

@Test public void producerUsesB3SingleNoParent() {
  // This injector needs no instanceof checks during inject()
  Injector<ProducerRequest> injector = Propagation.B3_STRING.injector(new ProducerSetter());

  ProducerRequest request = new ProducerRequest();
  injector.inject(context, request);

  assertThat(request.headers)
    .hasSize(1)
    .containsEntry("b3", "0000000000000001-0000000000000003");
}
 

@Override public <R> Injector<R> injector(Propagation.Setter<R, String> setter) {
  return delegate.injector(setter);
}
 

@Override public <R> Injector<R> injector(Setter<R, String> setter) {
  return new GrpcInjector<>(this, setter);
}
 

@Override
public <R> Injector<R> injector(Setter<R, String> setter) {
  if (setter == null) throw new NullPointerException("setter == null");
  return new AWSInjector<>(this, setter);
}
 

@Override public <R> Injector<R> injector(Setter<R, K> setter) {
  // No check if Setter is a RemoteSetter because this instance cannot have String keys while
  // RemoteSetter must have String keys
  return delegate.injector(new SetterAdapter<>(setter, map));
}
 

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

@Override public <R> Injector<R> injector(Setter<R, String> setter) {
  return new BaggageInjector<>(this, setter);
}
 

@Override public <R> Injector<R> injector(Setter<R, String> setter) {
  return injectorFactory.newInjector(setter);
}
 

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

@Override public <R> Injector<R> injector(Setter<R, String> setter) {
  return primary.injector(setter);
}
 

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

/**
 * @since 4.3
 * @deprecated Since 5.7, use {@link #handleSend(HttpClientRequest)}, as this allows more advanced
 * samplers to be used.
 */
@Deprecated public Span handleSend(Injector<Req> injector, Req request) {
  return handleSend(injector, request, request);
}
 

/**
 * @since 4.3
 * @deprecated Since 5.7, use {@link #handleSend(HttpClientRequest)}.
 */
@Deprecated public <C> Span handleSend(Injector<C> injector, C carrier, Req request) {
  return handleSend(injector, carrier, request, nextSpan(request));
}
 

/**
 * @since 4.4
 * @deprecated Since 5.7, use {@link #handleSend(HttpClientRequest, Span)}.
 */
@Deprecated public Span handleSend(Injector<Req> injector, Req request, Span span) {
  return handleSend(injector, request, request, span);
}