brave.baggage.CorrelationScopeConfig.SingleCorrelationField Java Examples

@Test public void allProperties() {
  context = new XmlBeans(""
    + "<util:constant id=\"traceId\" static-field=\"brave.baggage.BaggageFields.TRACE_ID\"/>\n"
    + "<bean id=\"traceIdCorrelationConfig\" class=\"brave.spring.beans.SingleCorrelationFieldFactoryBean\">\n"
    + "  <property name=\"baggageField\" ref=\"traceId\"/>\n"
    + "  <property name=\"name\" value=\"X-B3-TraceId\"/>\n"
    + "  <property name=\"dirty\" value=\"true\"/>\n"
    + "  <property name=\"flushOnUpdate\" value=\"true\"/>\n"
    + "</bean>\n"
  );

  assertThat(context.getBean("traceIdCorrelationConfig", CorrelationScopeConfig.class))
    .usingRecursiveComparison()
    .isEqualTo(SingleCorrelationField.newBuilder(BaggageFields.TRACE_ID)
      .name("X-B3-TraceId")
      .dirty()
      .flushOnUpdate()
      .build());
}
 

@Test public void configs() {
  context = new XmlBeans(""
    + "<util:constant id=\"traceId\" static-field=\"brave.baggage.BaggageFields.TRACE_ID\"/>\n"
    + "<bean id=\"correlationDecorator\" class=\"brave.spring.beans.CorrelationScopeDecoratorFactoryBean\">\n"
    + "  <property name=\"builder\">\n"
    + "    <bean class=\"brave.context.log4j12.MDCScopeDecorator\" factory-method=\"newBuilder\"/>\n"
    + "  </property>\n"
    + "  <property name=\"configs\">\n"
    + "    <list>\n"
    + "      <bean class=\"brave.spring.beans.SingleCorrelationFieldFactoryBean\">\n"
    + "        <property name=\"baggageField\" ref=\"traceId\"/>\n"
    + "        <property name=\"name\" value=\"X-B3-TraceId\"/>\n"
    + "      </bean>\n"
    + "    </list>\n"
    + "  </property>\n"
    + "</bean>"
  );

  ScopeDecorator decorator = context.getBean("correlationDecorator", ScopeDecorator.class);
  assertThat(decorator).extracting("field")
    .usingRecursiveComparison()
    .isEqualTo(
      SingleCorrelationField.newBuilder(TRACE_ID).name("X-B3-TraceId").build());
}
 

@Bean
CorrelationScopeCustomizer makeCorrelationFieldsDirty() {
	return b -> {
		Set<CorrelationScopeConfig> configs = b.configs();
		b.clear();

		for (CorrelationScopeConfig config : configs) {
			if (config instanceof SingleCorrelationField) {
				SingleCorrelationField field = (SingleCorrelationField) config;
				if (!field.readOnly()) {
					config = field.toBuilder().dirty().build();
				}
			}
			b.add(config);
		}
	};
}
 

@Test public void defaultFields() {
  context = new XmlBeans(""
    + "<bean id=\"correlationDecorator\" class=\"brave.spring.beans.CorrelationScopeDecoratorFactoryBean\">\n"
    + "  <property name=\"builder\">\n"
    + "    <bean class=\"brave.context.log4j12.MDCScopeDecorator\" factory-method=\"newBuilder\"/>\n"
    + "  </property>\n"
    + "</bean>"
  );

  assertThat(context.getBean("correlationDecorator", CorrelationScopeDecorator.class))
    .extracting("fields").asInstanceOf(InstanceOfAssertFactories.ARRAY)
    .containsExactly(
      SingleCorrelationField.create(TRACE_ID),
      SingleCorrelationField.create(SPAN_ID)
    );
}
 

/** @since 5.11 */
public Builder add(CorrelationScopeConfig config) {
  if (config == null) throw new NullPointerException("config == null");
  if (!(config instanceof SingleCorrelationField)) {
    throw new UnsupportedOperationException("dynamic fields not yet supported");
  }
  SingleCorrelationField field = (SingleCorrelationField) config;
  if (fields.contains(field)) {
    throw new IllegalArgumentException(
      "Baggage Field already added: " + field.baggageField.name);
  }
  if (allNames.contains(field.name)) {
    throw new IllegalArgumentException("Correlation name already in use: " + field.name);
  }
  fields.add(field);
  return this;
}
 

SingleCorrelationField(Builder builder) { // sealed to this package
  baggageField = builder.baggageField;
  name = builder.name;
  dirty = builder.dirty;
  flushOnUpdate = builder.flushOnUpdate;
  readOnly = baggageField.context instanceof BaggageContext.ReadOnly;
}
 

@Test public void leastProperties() {
  context = new XmlBeans(""
    + "<util:constant id=\"traceId\" static-field=\"brave.baggage.BaggageFields.TRACE_ID\"/>\n"
    + "<bean id=\"traceIdCorrelationConfig\" class=\"brave.spring.beans.SingleCorrelationFieldFactoryBean\">\n"
    + "  <property name=\"baggageField\" ref=\"traceId\"/>\n"
    + "</bean>\n"
  );

  assertThat(context.getBean("traceIdCorrelationConfig", CorrelationScopeConfig.class))
    .usingRecursiveComparison()
    .isEqualTo(SingleCorrelationField.create(BaggageFields.TRACE_ID));
}
 

@Override public SingleCorrelationField getObject() {
  SingleCorrelationField.Builder builder = SingleCorrelationField.newBuilder(baggageField);
  if (name != null) builder.name(name);
  if (dirty) builder.dirty();
  if (flushOnUpdate) builder.flushOnUpdate();
  return builder.build();
}
 

Multiple(
  Scope delegate,
  CorrelationContext context,
  SingleCorrelationField[] fields,
  String[] valuesToRevert,
  int shouldRevert
) {
  super(context);
  this.delegate = delegate;
  this.fields = fields;
  this.valuesToRevert = valuesToRevert;
  this.shouldRevert = shouldRevert;
}
 

Single(
  Scope delegate,
  CorrelationContext context,
  SingleCorrelationField field,
  @Nullable String valueToRevert,
  boolean shouldRevert
) {
  super(context);
  this.delegate = delegate;
  this.field = field;
  this.valueToRevert = valueToRevert;
  this.shouldRevert = shouldRevert;
}
 

@Override public Scope decorateScope(@Nullable TraceContext traceContext, Scope scope) {
  int dirty = 0;
  boolean flushOnUpdate = false;

  String[] valuesToRevert = new String[fields.length];
  for (int i = 0; i < fields.length; i++) {
    SingleCorrelationField field = fields[i];
    String valueToRevert = context.getValue(field.name);
    String currentValue = field.baggageField.getValue(traceContext);

    if (scope != Scope.NOOP || !field.readOnly) {
      if (!equal(valueToRevert, currentValue)) {
        context.update(field.name, currentValue);
        dirty = setBit(dirty, i);
      }
    }

    // Always revert fields that could be updated in the context directly
    if (field.dirty) dirty = setBit(dirty, i);
    if (field.flushOnUpdate) flushOnUpdate = true;

    valuesToRevert[i] = valueToRevert;
  }

  if (dirty == 0 && !flushOnUpdate) return scope;

  // If there was or could be a value update, we need to track values to revert.
  CorrelationUpdateScope updateScope =
    new CorrelationUpdateScope.Multiple(scope, context, fields, valuesToRevert, dirty);
  return flushOnUpdate ? new CorrelationFlushScope(updateScope) : updateScope;
}
 

/** @return {@link ScopeDecorator#NOOP} if no baggage fields were added. */
public final ScopeDecorator build() {
  int fieldCount = fields.size();
  if (fieldCount == 0) return ScopeDecorator.NOOP;
  if (fieldCount == 1) return new Single(context, fields.iterator().next());
  if (fieldCount > 32) throw new IllegalArgumentException("over 32 baggage fields");
  return new Multiple(context, fields.toArray(new SingleCorrelationField[0]));
}
 

/** Internal constructor used by subtypes. */
protected Builder(CorrelationContext context) {
  if (context == null) throw new NullPointerException("context == null");
  this.context = context;
  add(SingleCorrelationField.create(BaggageFields.TRACE_ID));
  add(SingleCorrelationField.create(BaggageFields.SPAN_ID));
}
 

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

/**
 * Returns a scope decorator that configures {@link BaggageFields#TRACE_ID}, {@link
 * BaggageFields#PARENT_ID}, {@link BaggageFields#SPAN_ID} and {@link BaggageFields#SAMPLED}
 *
 * @since 5.2
 * @deprecated since 5.11 use {@link #get()} or {@link #newBuilder()}
 */
@Deprecated public static CurrentTraceContext.ScopeDecorator create() {
  return new Builder()
    .clear()
    .add(SingleCorrelationField.create(BaggageFields.TRACE_ID))
    .add(SingleCorrelationField.create(BaggageFields.PARENT_ID))
    .add(SingleCorrelationField.create(BaggageFields.SPAN_ID))
    .add(SingleCorrelationField.create(BaggageFields.SAMPLED))
    .build();
}
 

/**
 * Returns a scope decorator that configures {@link BaggageFields#TRACE_ID}, {@link
 * BaggageFields#PARENT_ID}, {@link BaggageFields#SPAN_ID} and {@link BaggageFields#SAMPLED}
 *
 * @since 5.2
 * @deprecated since 5.11 use {@link #get()} or {@link #newBuilder()}
 */
@Deprecated public static CurrentTraceContext.ScopeDecorator create() {
  return new Builder()
    .clear()
    .add(SingleCorrelationField.create(BaggageFields.TRACE_ID))
    .add(SingleCorrelationField.create(BaggageFields.PARENT_ID))
    .add(SingleCorrelationField.create(BaggageFields.SPAN_ID))
    .add(SingleCorrelationField.create(BaggageFields.SAMPLED))
    .build();
}
 

/**
 * Returns a scope decorator that configures {@link BaggageFields#TRACE_ID}, {@link
 * BaggageFields#PARENT_ID}, {@link BaggageFields#SPAN_ID} and {@link BaggageFields#SAMPLED}
 *
 * @since 5.2
 * @deprecated since 5.11 use {@link #get()} or {@link #newBuilder()}
 */
@Deprecated public static CurrentTraceContext.ScopeDecorator create() {
  return new Builder()
    .clear()
    .add(SingleCorrelationField.create(BaggageFields.TRACE_ID))
    .add(SingleCorrelationField.create(BaggageFields.PARENT_ID))
    .add(SingleCorrelationField.create(BaggageFields.SPAN_ID))
    .add(SingleCorrelationField.create(BaggageFields.SAMPLED))
    .build();
}
 

@Test
public void canMakeAllCorrelationFieldsDirty() {
	this.contextRunner
			.withPropertyValues(
					"spring.sleuth.baggage.correlation-fields=country-code")
			.withUserConfiguration(DirtyCorrelationFieldConfiguration.class)
			.run((context) -> assertThat(
					context.getBean(CorrelationScopeDecorator.class))
							.extracting("fields")
							.asInstanceOf(array(SingleCorrelationField[].class))
							.filteredOn(c -> !c.readOnly())
							.extracting(SingleCorrelationField::dirty)
							.containsExactly(true));
}
 

@Test
public void canAddOldCorrelationFieldsForLogScraping() {
	this.contextRunner
			.withUserConfiguration(
					OldCorrelationFieldsForLogScrapingConfiguration.class)
			.run((context) -> assertThat(
					context.getBean(CorrelationScopeDecorator.class))
							.extracting("fields")
							.asInstanceOf(array(SingleCorrelationField[].class))
							.extracting(SingleCorrelationField::name)
							.containsExactly("traceId", "spanId", "parentId",
									"spanExportable"));
}
 

@Test
public void should_only_include_whitelist() {
	assertThat(this.scopeDecorator).extracting("fields")
			.asInstanceOf(
					InstanceOfAssertFactories.array(SingleCorrelationField[].class))
			.extracting(SingleCorrelationField::name)
			.containsOnly("traceId", "spanId", "bp", COUNTRY_CODE.name()); // x-vcap-request-id
																			// is not
																			// in the
																			// whitelist
}
 

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

/** Returns true for any config with the same baggage field. */
@Override public boolean equals(Object o) {
  if (o == this) return true;
  if (!(o instanceof SingleCorrelationField)) return false;
  return baggageField.equals(((SingleCorrelationField) o).baggageField);
}
 

Builder(SingleCorrelationField input) {
  baggageField = input.baggageField;
  name = input.name;
  dirty = input.dirty;
  flushOnUpdate = input.flushOnUpdate;
}
 

/** @since 5.11 */
public static SingleCorrelationField create(BaggageField baggageField) {
  return new Builder(baggageField).build();
}
 

Single(CorrelationContext context, SingleCorrelationField field) {
  super(context);
  this.field = field;
}
 

Multiple(CorrelationContext context, SingleCorrelationField[] fields) {
  super(context);
  this.fields = fields;
}
 

/** Allows log patterns to use {@code %{traceId}} {@code %{spanId}} and {@code %{userName}} */
@Bean ScopeDecorator correlationScopeDecorator() {
  return ThreadContextScopeDecorator.newBuilder()
      .add(SingleCorrelationField.create(USER_NAME)).build();
}
 

/** Allows log patterns to use {@code %{traceId}} {@code %{spanId}} and {@code %{userName}} */
@Bean ScopeDecorator correlationScopeDecorator() {
  return MDCScopeDecorator.newBuilder()
      .add(SingleCorrelationField.create(USER_NAME)).build();
}
 

@Bean
CorrelationScopeCustomizer addParentAndSpanExportable() {
	return b -> b.add(SingleCorrelationField.create(BaggageFields.PARENT_ID))
			.add(SingleCorrelationField.newBuilder(BaggageFields.SAMPLED)
					.name("spanExportable").build());
}