brave.baggage.BaggagePropagationConfig.SingleBaggageField Java Examples

@Test public void allProperties() {
  context = new XmlBeans(""
    + "<bean id=\"userId\" class=\"brave.baggage.BaggageField\" factory-method=\"create\">\n"
    + "  <constructor-arg><value>userId</value></constructor-arg>\n"
    + "</bean>\n"
    + "<bean id=\"userIdBaggageConfig\" class=\"brave.spring.beans.SingleBaggageFieldFactoryBean\">\n"
    + "  <property name=\"field\" ref=\"userId\"/>\n"
    + "  <property name=\"keyNames\">\n"
    + "    <list>\n"
    + "      <value>user-id</value>\n"
    + "    </list>\n"
    + "  </property>\n"
    + "</bean>\n"
  );

  assertThat(context.getBean("userIdBaggageConfig", BaggagePropagationConfig.class))
    .usingRecursiveComparison()
    .isEqualTo(SingleBaggageField.newBuilder(BaggageField.create("userId"))
      .addKeyName("user-id")
      .build());
}
 

@Test public void clear_and_add() {
  SingleBaggageField requestIdConfig = SingleBaggageField.newBuilder(vcapRequestId)
    .addKeyName("request-id")
    .addKeyName("request_id")
    .build();

  SingleBaggageField traceIdConfig = SingleBaggageField.remote(amznTraceId);
  BaggagePropagation.FactoryBuilder builder = newFactoryBuilder(B3Propagation.FACTORY)
    .add(requestIdConfig)
    .add(traceIdConfig);

  Set<BaggagePropagationConfig> configs = builder.configs();

  builder.clear();

  configs.forEach(builder::add);

  assertThat(builder)
    .usingRecursiveComparison()
    .isEqualTo(newFactoryBuilder(B3Propagation.FACTORY)
      .add(requestIdConfig)
      .add(traceIdConfig));
}
 

@Test public void extract_no_overridden_key_names() {
  BaggageField userId = BaggageField.create("userId");
  BaggageField sessionId = BaggageField.create("sessionId");

  factory = newFactoryBuilder(B3Propagation.FACTORY)
    .add(SingleBaggageField.local(userId))
    .add(SingleBaggageField.remote(sessionId))
    .build();
  initialize();

  injector.inject(context, request);
  request.put("userid", "bob");
  request.put("sessionid", "12345");

  context = extractor.extract(request).context();

  assertThat(userId.getValue(context)).isNull();
  assertThat(sessionId.getValue(context)).isEqualTo("12345");
}
 

Factory(FactoryBuilder factoryBuilder) {
  this.delegateFactory = factoryBuilder.delegate;
  this.delegate = delegateFactory.get();

  // Don't add another "extra" if there are only local fields
  List<String> extractKeyNames = Lists.ensureImmutable(factoryBuilder.extractKeyNames);
  this.extra = !extractKeyNames.isEmpty() ? new Extra(extractKeyNames) : null;

  // Associate baggage fields with any remote propagation keys
  this.configs = factoryBuilder.configs.toArray(new BaggagePropagationConfig[0]);

  List<BaggageField> fields = new ArrayList<>();
  Set<String> localFieldNames = new LinkedHashSet<>();
  int maxDynamicFields = 0;
  for (BaggagePropagationConfig config : factoryBuilder.configs) {
    maxDynamicFields += config.maxDynamicFields;
    if (config instanceof SingleBaggageField) {
      BaggageField field = ((SingleBaggageField) config).field;
      fields.add(field);
      if (config.baggageCodec == BaggageCodec.NOOP) localFieldNames.add(field.name());
    }
  }
  this.baggageFactory = BaggageFields.newFactory(fields, maxDynamicFields);
  this.localFieldNames = localFieldNames.toArray(new String[0]);
}
 

public TracedBaggage() {
  super(Tracing.newBuilder()
    .propagationFactory(BaggagePropagation.newFactoryBuilder(B3Propagation.FACTORY)
      .add(SingleBaggageField.remote(REQUEST_ID))
      .add(SingleBaggageField.newBuilder(COUNTRY_CODE)
        .addKeyName("baggage-country-code")
        .build())
      .add(SingleBaggageField.newBuilder(USER_ID)
        .addKeyName("baggage-user-id")
        .build())
      .build())
    .addSpanHandler(new SpanHandler() {
      // intentionally not NOOP to ensure spans report
    })
    .build());
}
 

protected ITRemote() {
  CurrentTraceContext.Builder currentTraceContextBuilder = currentTraceContextBuilder();
  if (currentTraceContextBuilder instanceof StrictCurrentTraceContext.Builder) {
    currentTraceContext = currentTraceContextBuilder.build();
    checkForLeakedScopes = (Closeable) currentTraceContext;
  } else {
    StrictScopeDecorator strictScopeDecorator = StrictScopeDecorator.create();
    currentTraceContext = currentTraceContextBuilder
      .addScopeDecorator(strictScopeDecorator).build();
    checkForLeakedScopes = strictScopeDecorator;
  }
  propagationFactory = BaggagePropagation.newFactoryBuilder(B3Propagation.FACTORY)
    .add(SingleBaggageField.newBuilder(BAGGAGE_FIELD)
      .addKeyName(BAGGAGE_FIELD_KEY)
      .build()).build();
  tracing = tracingBuilder(Sampler.ALWAYS_SAMPLE).build();
}
 

/** Returns a wrapper of the delegate if there are no fields to propagate. */
public Factory build() {
  Set<String> extraKeyNames = new LinkedHashSet<>();
  for (Map.Entry<String, Set<String>> entry : nameToKeyNames.entrySet()) {
    BaggageField field = BaggageField.create(entry.getKey());
    if (redactedNames.contains(field.name())) {
      baggageFactory.add(SingleBaggageField.local(field));
    } else {
      extraKeyNames.addAll(entry.getValue());
      SingleBaggageField.Builder builder = SingleBaggageField.newBuilder(field);
      for (String keyName : entry.getValue()) {
        builder.addKeyName(keyName);
      }
      baggageFactory.add(builder.build());
    }
  }
  return new Factory(baggageFactory.build(), extraKeyNames.toArray(new String[0]));
}
 

/** Redaction prevents named fields from being written downstream. */
@Test public void inject_no_key_names() {
  BaggageField userId = BaggageField.create("userId");
  BaggageField sessionId = BaggageField.create("sessionId");

  factory = newFactoryBuilder(B3SinglePropagation.FACTORY)
    .add(SingleBaggageField.local(userId))
    .add(SingleBaggageField.remote(sessionId))
    .build();
  initialize();

  userId.updateValue(context, "bob");
  sessionId.updateValue(context, "12345");

  injector.inject(context, request);

  assertThat(request)
    .doesNotContainKey("userid")
    .containsEntry("sessionid", "12345");
}
 

@Test public void extract_field_multiple_key_names() {
  // switch to case insensitive as this example is about http :P
  request = new TreeMap<>(String.CASE_INSENSITIVE_ORDER);

  BaggageField userId = BaggageField.create("userId");
  BaggageField sessionId = BaggageField.create("sessionId");

  SingleBaggageField userIdConfig = SingleBaggageField.newBuilder(userId)
    .addKeyName("baggage-userId")
    .addKeyName("baggage_userId")
    .build();

  SingleBaggageField sessionIdConfig = SingleBaggageField.newBuilder(sessionId)
    .addKeyName("baggage-sessionId")
    .addKeyName("baggage_sessionId")
    .build();

  factory = newFactoryBuilder(B3Propagation.FACTORY)
    .add(userIdConfig)
    .add(sessionIdConfig)
    .build();
  initialize();

  injector.inject(context, request);
  request.put("baggage-userId", "bob");
  request.put("baggage-sessionId", "12345");

  context = extractor.extract(request).context();

  assertThat(userId.getValue(context)).isEqualTo("bob");
  assertThat(sessionId.getValue(context)).isEqualTo("12345");
}
 

@Test public void newFactory_sharingRemoteName() {
  BaggagePropagation.FactoryBuilder builder = newFactoryBuilder(B3Propagation.FACTORY);
  SingleBaggageField userName =
    SingleBaggageField.newBuilder(BaggageField.create("userName")).addKeyName("baggage").build();
  SingleBaggageField userId =
    SingleBaggageField.newBuilder(BaggageField.create("userId")).addKeyName("baggage").build();
  builder.add(userName);
  assertThatThrownBy(() -> builder.add(userId))
    .isInstanceOf(IllegalArgumentException.class)
    .hasMessage("Propagation key already in use: baggage");
}
 

@Test public void inject_field_multiple_key_names() {
  BaggageField userId = BaggageField.create("userId");
  BaggageField sessionId = BaggageField.create("sessionId");

  SingleBaggageField userIdConfig = SingleBaggageField.newBuilder(userId)
    .addKeyName("userId")
    .addKeyName("baggage-userId")
    .addKeyName("baggage_userId")
    .build();

  SingleBaggageField sessionIdConfig = SingleBaggageField.newBuilder(sessionId)
    .addKeyName("sessionId")
    .addKeyName("baggage-sessionId")
    .addKeyName("baggage_sessionId")
    .build();

  factory = newFactoryBuilder(B3SinglePropagation.FACTORY)
    .add(userIdConfig)
    .add(sessionIdConfig)
    .build();
  initialize();

  userId.updateValue(context, "bob");
  sessionId.updateValue(context, "12345");

  injector.inject(context, request);

  // NOTE: the labels are downcased
  assertThat(request).containsOnly(
    entry("b3", B3SingleFormat.writeB3SingleFormat(context)),
    entry("userid", "bob"),
    entry("sessionid", "12345"),
    entry("baggage-userid", "bob"),
    entry("baggage-sessionid", "12345"),
    entry("baggage_userid", "bob"),
    entry("baggage_sessionid", "12345")
  );
}
 

@Test public void dupesNotOk() {
  SingleBaggageField userIdConfig = SingleBaggageField.local(BaggageField.create("userId"));
  BaggagePropagation.FactoryBuilder builder = newFactoryBuilder(B3Propagation.FACTORY)
    .add(userIdConfig);
  assertThatThrownBy(() -> builder.add(userIdConfig))
    .isInstanceOf(IllegalArgumentException.class);
}
 

@Test public void allKeyNames_baggagePropagation() {
  Propagation.Factory factory = BaggagePropagation.newFactoryBuilder(B3SinglePropagation.FACTORY)
    .add(SingleBaggageField.local(BaggageField.create("redacted"))) // local shouldn't return
    .add(SingleBaggageField.remote(BaggageField.create("user-id")))
    .add(SingleBaggageField.remote(BaggageField.create("session-id"))).build();

  assertThat(BaggagePropagation.allKeyNames(factory.get()))
    .containsExactly("b3", "user-id", "session-id");
}
 

@Test public void allKeyNames_baggagePropagation_noRemote() {
  Propagation.Factory factory = BaggagePropagation.newFactoryBuilder(B3SinglePropagation.FACTORY)
      .add(SingleBaggageField.local(BaggageField.create("redacted"))) // local shouldn't return
      .add(SingleBaggageField.local(BaggageField.create("user-id")))
      .add(SingleBaggageField.local(BaggageField.create("session-id"))).build();

  assertThat(BaggagePropagation.allKeyNames(factory.get()))
      .containsExactly("b3");
}
 

public TracedBaggage() {
  super(TracingFilter.create(Tracing.newBuilder()
    .propagationFactory(BaggagePropagation.newFactoryBuilder(B3Propagation.FACTORY)
      .add(SingleBaggageField.remote(BAGGAGE_FIELD)).build())
    .spanReporter(Reporter.NOOP)
    .build()));
}
 

@Override public Set<Object> getSingletons() {
  return new LinkedHashSet<>(asList(new Resource(), TracingApplicationEventListener.create(
    HttpTracing.create(Tracing.newBuilder()
      .propagationFactory(BaggagePropagation.newFactoryBuilder(B3Propagation.FACTORY)
        .add(SingleBaggageField.remote(BAGGAGE_FIELD)).build())
      .spanReporter(Reporter.NOOP)
      .build())
  )));
}
 

@Override public SingleBaggageField getObject() {
  SingleBaggageField.Builder builder = SingleBaggageField.newBuilder(field);
  if (keyNames != null) {
    for (String keyName : keyNames) {
      builder.addKeyName(keyName);
    }
  }
  return builder.build();
}
 

@Test public void leastProperties() {
  context = new XmlBeans(""
    + "<bean id=\"userId\" class=\"brave.baggage.BaggageField\" factory-method=\"create\">\n"
    + "  <constructor-arg><value>userId</value></constructor-arg>\n"
    + "</bean>\n"
    + "<bean id=\"userIdBaggageConfig\" class=\"brave.spring.beans.SingleBaggageFieldFactoryBean\">\n"
    + "  <property name=\"field\" ref=\"userId\"/>\n"
    + "</bean>\n"
  );

  assertThat(context.getBean("userIdBaggageConfig", BaggagePropagationConfig.class))
    .usingRecursiveComparison()
    .isEqualTo(SingleBaggageField.local(BaggageField.create("userId")));
}
 

/** Less overhead and distraction for the edge case of correlation-only. */
@Test public void extract_baggage_onlyOneExtraWhenNothingRemote() {
  Propagation.Factory factory = newFactoryBuilder(B3Propagation.FACTORY)
      .add(SingleBaggageField.local(vcapRequestId))
      .add(SingleBaggageField.local(amznTraceId)).build();
  extractor = factory.get().extractor(Map::get);

  TraceContextOrSamplingFlags extracted = extractor.extract(request);
  assertThat(extracted.extra())
      .hasSize(1)
      .noneMatch(Extra.class::isInstance);
}
 

void init(Tracing.Builder tracingBuilder) {
  if (brave != null) brave.close();
  brave = tracingBuilder
      .localServiceName("tracer")
      .currentTraceContext(currentTraceContext)
      .addSpanHandler(spans)
      .propagationFactory(BaggagePropagation.newFactoryBuilder(B3Propagation.FACTORY)
          .add(SingleBaggageField.remote(BaggageField.create("client-id")))
          .build()).build();
  tracer = BraveTracer.create(brave);
}
 

SingleBaggageField(Builder builder) { // sealed to this package
  super(builder.keyNames.isEmpty()
      ? BaggageCodec.NOOP
      : SingleFieldBaggageCodec.single(builder.field, builder.keyNames), 0);
  field = builder.field;
  keyNames = builder.keyNames.isEmpty() ? Collections.emptySet()
      : Collections.unmodifiableSet(new LinkedHashSet<>(builder.keyNames));
}
 

@Bean
BaggagePropagationCustomizer countryCodeBaggageConfig() {
	return fb -> fb.add(
			SingleBaggageField.newBuilder(BaggageField.create("country-code"))
					.addKeyName("baggage-country-code")
					.addKeyName("baggage_country-code").build());
}
 

/** @since 5.11 */
public SingleBaggageField build() {
  return new SingleBaggageField(this);
}
 

Builder(SingleBaggageField input) {
  this.field = input.field;
  this.keyNames = new ArrayList<>(input.keyNames());
}
 

/** Configures propagation for {@link #USER_NAME}, using the remote header "user_name" */
@Bean Propagation.Factory propagationFactory() {
  return BaggagePropagation.newFactoryBuilder(B3Propagation.FACTORY)
      .add(SingleBaggageField.newBuilder(USER_NAME).addKeyName("user_name").build())
      .build();
}
 

/** Configures propagation for {@link #USER_NAME}, using the remote header "user_name" */
@Bean Propagation.Factory propagationFactory() {
  return BaggagePropagation.newFactoryBuilder(B3Propagation.FACTORY)
      .add(SingleBaggageField.newBuilder(USER_NAME).addKeyName("user_name").build())
      .build();
}
 

@Bean
BaggagePropagationCustomizer requestId() {
	return fb -> fb.add(
			SingleBaggageField.remote(BaggageField.create("x-vcap-request-id")));
}
 

@Bean
BaggagePropagationCustomizer countryCode() {
	return fb -> fb
			.add(SingleBaggageField.remote(BaggageField.create("country-code")));
}
 

@Bean
BaggagePropagationConfig notInProperties() {
	return SingleBaggageField.remote(BaggageField.create("bar"));
}
 

/**
 * Configures this field for remote propagation using its lower-case {@link BaggageField#name()}
 * as the only {@linkplain #keyNames() propagation key name}.
 *
 * @see #local(BaggageField)
 * @see #newBuilder(BaggageField) to use different propagation key names.
 */
public static SingleBaggageField remote(BaggageField field) {
  return new Builder(field).addKeyName(field.lcName).build();
}