brave.baggage.BaggageField 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());
}
 

@Bean
SpanHandler baggageTagSpanHandler(
		@Qualifier(WHITELISTED_KEYS) List<String> whiteListedKeys,
		SleuthBaggageProperties sleuthBaggageProperties) {

	Set<String> tagFields = redirectOldPropertyToNew(WHITELISTED_KEYS,
			whiteListedKeys, "spring.sleuth.baggage.tag-fields",
			sleuthBaggageProperties.getTagFields());

	if (tagFields.isEmpty()) {
		return SpanHandler.NOOP; // Brave ignores these
	}

	return new BaggageTagSpanHandler(tagFields.stream().map(BaggageField::create)
			.toArray(BaggageField[]::new));
}
 

@RequestMapping(
		value = "/order",
		method = POST)
String order(HttpEntity<String> body) {
	String processIdFromHeaders = body.getHeaders().getFirst(PROCESS_ID_HEADER_NAME);
	String processId = StringUtils.hasText(body.getHeaders().getFirst(PROCESS_ID_HEADER_NAME)) ?
			processIdFromHeaders :
			new JdkIdGenerator().generateId().toString();
	log.info("Making new order with [{}] and processid [{}].", body.getBody(), processId);
	Span span = this.tracer.nextSpan().name("inside_presenting").start();
	Tracer.SpanInScope ws = tracer.withSpanInScope(span);
	try {
		String testCommunicationType = BaggageField.getByName("TEST-COMMUNICATION-TYPE").getValue();
		log.info("Found the following communication type [{}]", testCommunicationType);
		switch (testCommunicationType) {
		case "FEIGN":
			return useFeignToCallAggregation(body, processId);
		default:
			return useRestTemplateToCallAggregation(body, processId);
		}
	} finally {
		span.finish();
		ws.close();
	}
}
 

@Test public void extract_prefixed() {
  factory = newFactoryBuilder(B3Propagation.FACTORY)
    .addField("x-vcap-request-id")
    .addPrefixedFields("baggage-", asList("country-code"))
    .build();
  initialize();

  injector.inject(context, request);
  request.put("baggage-country-code", "FO");
  request.put("x-vcap-request-id", uuid);

  TraceContextOrSamplingFlags extracted = extractor.extract(request);
  assertThat(extracted.context().toBuilder().extra(Collections.emptyList()).build())
    .isEqualTo(context);
  assertThat(extracted.context().extra())
    .hasSize(2);

  assertThat(BaggageField.getByName(extracted, "country-code").getValue(extracted))
    .isEqualTo("FO");
  assertThat(BaggageField.getByName(extracted, "x-vcap-request-id").getValue(extracted))
    .isEqualTo(uuid);
}
 

/** 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]));
}
 

private void notifyPresentingService(String correlationId) {
    log.info("Calling presenting from maturing");
    Span scope = this.tracer.nextSpan().name("calling_presenting_from_maturing").start();
    try (Tracer.SpanInScope ws = tracer.withSpanInScope(scope)) {
        String testCommunicationType = BaggageField.getByName("TEST-COMMUNICATION-TYPE").getValue();
        log.info("Found the following communication type [{}]", testCommunicationType);
        switch (testCommunicationType) {
        case "FEIGN":
            callPresentingViaFeign(correlationId);
            break;
        default:
            useRestTemplateToCallPresenting(correlationId);
        }
    } finally {
        scope.finish();
    }
}
 

@Test public void extract_two() {
  injector.inject(context, request);
  request.put("x-amzn-trace-id", awsTraceId);
  request.put("x-vcap-request-id", uuid);

  TraceContextOrSamplingFlags extracted = extractor.extract(request);
  assertThat(extracted.context().toBuilder().extra(Collections.emptyList()).build())
    .isEqualTo(context);
  assertThat(extracted.context().extra())
    .hasSize(2);

  assertThat(BaggageField.getByName(extracted, "x-amzn-trace-id").getValue(extracted))
    .isEqualTo(awsTraceId);
  assertThat(BaggageField.getByName(extracted, "x-vcap-request-id").getValue(extracted))
    .isEqualTo(uuid);
}
 

@Test public void inject_prefixed() {
  factory = newFactoryBuilder(B3Propagation.FACTORY)
    .addField("x-vcap-request-id")
    .addPrefixedFields("baggage-", asList("country-code"))
    .build();
  initialize();

  BaggageField.getByName(context, "x-vcap-request-id").updateValue(context, uuid);
  BaggageField.getByName(context, "country-code").updateValue(context, "FO");

  injector.inject(context, request);

  assertThat(request)
    .containsEntry("baggage-country-code", "FO")
    .containsEntry("x-vcap-request-id", uuid);
}
 

private void notifyPresentingService(String processId) {
    Span scope = this.tracer.nextSpan().name("calling_presenting").start();
    try (Tracer.SpanInScope ws = tracer.withSpanInScope(scope)) {
        String testCommunicationType = BaggageField.getByName("TEST-COMMUNICATION-TYPE").getValue();
        log.info("Found the following communication type [{}]", testCommunicationType);
        if (testCommunicationType.equals("FEIGN")) {
            callPresentingViaFeign(processId);
        }
        else {
            useRestTemplateToCallPresenting(processId);
        }
    } finally {
        scope.finish();
    }
}
 

private String useFeignToCallAggregation(HttpEntity<String> body, String processId) {
	String testCommunicationType = BaggageField.getByName("TEST-COMMUNICATION-TYPE").getValue();
	log.info("Found the following communication type [{}]", testCommunicationType);
	return aggregationServiceClient.getIngredients(body.getBody(),
			processId,
			testCommunicationType);
}
 

@Bean
@ConditionalOnMissingBean(CorrelationScopeDecorator.class)
@ConditionalOnBean(CorrelationScopeDecorator.Builder.class)
@ConditionalOnProperty(value = "spring.sleuth.baggage.correlation-enabled",
		matchIfMissing = true)
ScopeDecorator correlationScopeDecorator(
		@Qualifier(WHITELISTED_MDC_KEYS) List<String> whiteListedMDCKeys,
		SleuthBaggageProperties sleuthBaggageProperties,
		@Nullable List<CorrelationScopeCustomizer> correlationScopeCustomizers) {

	Set<String> correlationFields = redirectOldPropertyToNew(WHITELISTED_MDC_KEYS,
			whiteListedMDCKeys, "spring.sleuth.baggage.correlation-fields",
			sleuthBaggageProperties.getCorrelationFields());

	// Add fields from properties
	CorrelationScopeDecorator.Builder builder = MDCScopeDecorator.newBuilder();
	for (String field : correlationFields) {
		builder.add(SingleCorrelationField.newBuilder(BaggageField.create(field))
				.build());
	}

	// handle user overrides
	if (correlationScopeCustomizers != null) {
		for (CorrelationScopeCustomizer customizer : correlationScopeCustomizers) {
			customizer.customize(builder);
		}
	}
	return builder.build();
}
 

void notifyPresenting(String processId) {
    log.info("I'm inside bottling. Notifying presenting");
    String testCommunicationType = BaggageField.getByName("TEST-COMMUNICATION-TYPE").getValue();
    log.info("Found the following communication type [{}]", testCommunicationType);
    if (testCommunicationType.equals("FEIGN")) {
        callPresentingViaFeign(processId);
    }
    else {
        useRestTemplateToCallPresenting(processId);
    }
}
 

List<Ingredient> collectIngredients(Order order, String processId) {
	String testCommunicationType = BaggageField.getByName("TEST-COMMUNICATION-TYPE").getValue();
	log.info("Found the following communication type [{}]", testCommunicationType);
	switch (testCommunicationType) {
		case "FEIGN":
			return callViaFeign(order, processId);
		default:
			return callViaRestTemplate(order, processId);
	}
}
 

@Test
void should_use_baggageBean() {
	this.contextRunner.withUserConfiguration(WithBaggageBeans.class, Baggage.class)
			.run((context -> {
				final Baggage bean = context.getBean(Baggage.class);
				BDDAssertions.then(bean.fields).containsOnly(
						BaggageField.create("country-code"),
						BaggageField.create("x-vcap-request-id"));
			}));
}
 

private void handleEvents(Event event, Map<String, Object> headers) {
	log.info("Received the following message with headers [{}] and body [{}]", headers, event);
	String testCommunicationType = BaggageField.getByName("TEST-COMMUNICATION-TYPE").getValue();
	log.info("Found the following communication type [{}]", testCommunicationType);
	Span newSpan = tracer.nextSpan().name("inside_reporting").start();
	try (Tracer.SpanInScope ws = tracer.withSpanInScope(newSpan)) {
		reportingRepository.createOrUpdate(event);
		newSpan.annotate("savedEvent");
		log.info("Saved event to the db", headers, event);
	} finally {
		newSpan.finish();
	}
}
 

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);
}
 

@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")));
}
 

/** This is a NOOP unless {@link BaggagePropagation} is in use */
@Override public BraveSpan setBaggageItem(String key, String value) {
  BaggageField field = BaggageField.getByName(delegate.context(), key);
  if (field == null) return this;
  field.updateValue(delegate.context(), value);
  return this;
}
 

@Override
public boolean end(TraceContext context, MutableSpan span, Cause cause) {
	for (BaggageField field : fieldsToTag) {
		Tags.BAGGAGE_FIELD.tag(field, context, span);
	}
	return true;
}
 

@Test
void should_use_local_keys_from_properties() {
	this.contextRunner.withPropertyValues("spring.sleuth.baggage.local-fields=bp")
			.withUserConfiguration(Baggage.class).run((context -> {
				final Baggage bean = context.getBean(Baggage.class);
				BDDAssertions.then(bean.fields)
						.containsExactly(BaggageField.create("bp"));
			}));
}
 

@Test
void should_combine_baggage_beans_and_properties() {
	this.contextRunner.withPropertyValues("spring.sleuth.baggage.local-fields=bp")
			.withUserConfiguration(WithBaggageBeans.class, Baggage.class)
			.run((context -> {
				final Baggage bean = context.getBean(Baggage.class);
				BDDAssertions.then(bean.fields).containsOnly(
						BaggageField.create("country-code"),
						BaggageField.create("x-vcap-request-id"),
						BaggageField.create("bp"));
			}));
}
 

@Autowired
Baggage(Tracing tracing) {
	// When predefined baggage fields exist, the result !=
	// TraceContextOrSamplingFlags.EMPTY
	TraceContextOrSamplingFlags emptyExtraction = tracing.propagation()
			.extractor((c, k) -> null).extract(Boolean.TRUE);
	fields = BaggageField.getAll(emptyExtraction);
}
 

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

@Test public void extract_extra() {
  injector.inject(context, request);
  request.put("x-amzn-trace-id", awsTraceId);

  TraceContextOrSamplingFlags extracted = extractor.extract(request);
  assertThat(extracted.context().toBuilder().extra(Collections.emptyList()).build())
    .isEqualTo(context);
  assertThat(extracted.context().extra())
    .hasSize(2);

  assertThat(BaggageField.getByName(extracted, "x-amzn-trace-id").getValue(extracted))
    .isEqualTo(awsTraceId);
}
 

@Test public void inject_two() {
  BaggageField.getByName(context, "x-vcap-request-id").updateValue(context, uuid);
  BaggageField.getByName(context, "x-amzn-trace-id").updateValue(context, awsTraceId);

  injector.inject(context, request);

  assertThat(request)
    .containsEntry("x-amzn-trace-id", awsTraceId)
    .containsEntry("x-vcap-request-id", uuid);
}
 

@Nullable static BaggageField getFieldByName(List<BaggageField> fields, String name) {
  if (name == null) throw new NullPointerException("name == null");
  name = name.trim();
  if (name.isEmpty()) throw new IllegalArgumentException("name is empty");
  for (BaggageField field : fields) {
    if (name.equals(field.name())) {
      return field;
    }
  }
  return null;
}
 

@Nullable static String getValue(BaggageField field, List<Object> extraList) {
  BaggageFields extra = findExtra(BaggageFields.class, extraList);
  if (extra == null) return null;
  return extra.getValue(field);
}
 

@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, String value) {
  return put(field, value);
}
 

static boolean updateValue(BaggageField field, List<Object> extraList, @Nullable String value) {
  BaggageFields extra = findExtra(BaggageFields.class, extraList);
  return extra != null && extra.updateValue(field, value);
}