Add local and tag correlation fields

spring-boot-project/spring-boot-actuator-autoconfigure/src/main/java/org/springframework/boot/actuate/autoconfigure/tracing/BaggageTagSpanHandler.java

@@ -0,0 +1,41 @@
+/*
+ * Copyright 2012-2023 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.boot.actuate.autoconfigure.tracing;
+
+import brave.Tags;
+import brave.baggage.BaggageField;
+import brave.handler.MutableSpan;
+import brave.handler.SpanHandler;
+import brave.propagation.TraceContext;
+
+public class BaggageTagSpanHandler extends SpanHandler {
+
+	final BaggageField[] fieldsToTag;
+
+	BaggageTagSpanHandler(BaggageField[] fieldsToTag) {
+		this.fieldsToTag = fieldsToTag;
+	}
+
+	@Override
+	public boolean end(TraceContext context, MutableSpan span, Cause cause) {
+		for (BaggageField field : this.fieldsToTag) {
+			Tags.BAGGAGE_FIELD.tag(field, context, span);
+		}
+		return true;
+	}
+
+}

spring-boot-project/spring-boot-actuator-autoconfigure/src/main/java/org/springframework/boot/actuate/autoconfigure/tracing/BraveAutoConfiguration.java

@@ -244,6 +244,28 @@ BaggagePropagationCustomizer remoteFieldsBaggagePropagationCustomizer() {
 			};
 		}
 
+		@Bean
+		@Order(1)
+		BaggagePropagationCustomizer localFieldsBaggagePropagationCustomizer() {
+			return (builder) -> {
+				final List<String> localFields = this.tracingProperties.getBaggage().getLocalFields();
+				for (final String localFieldName : localFields) {
+					builder.add(BaggagePropagationConfig.SingleBaggageField.local(BaggageField.create(localFieldName)));
+				}
+			};
+		}
+
+		@Bean
+		@Order(2)
+		SpanHandler baggageTagSpanHandler() {
+			final List<String> tagFields = this.tracingProperties.getBaggage().getTagFields();
+
+			if (tagFields.isEmpty()) {
+				return SpanHandler.NOOP; // Brave ignores these
+			}
+			return new BaggageTagSpanHandler(tagFields.stream().map(BaggageField::create).toArray(BaggageField[]::new));
+		}
+
 		@Bean
 		@ConditionalOnMissingBean
 		Factory propagationFactory(BaggagePropagation.FactoryBuilder factoryBuilder) {

spring-boot-project/spring-boot-actuator-autoconfigure/src/main/java/org/springframework/boot/actuate/autoconfigure/tracing/TracingProperties.java

@@ -116,6 +116,17 @@ public static class Baggage {
 		 */
 		private List<String> remoteFields = new ArrayList<>();
 
+		/**
+		 * List of fields that should be accessible within the JVM process but not
+		 * propagated over the wire.
+		 */
+		private List<String> localFields = new ArrayList<>();
+
+		/**
+		 * List of fields that should automatically become tags.
+		 */
+		private List<String> tagFields = new ArrayList<>();
+
 		public boolean isEnabled() {
 			return this.enabled;
 		}
@@ -136,10 +147,26 @@ public List<String> getRemoteFields() {
 			return this.remoteFields;
 		}
 
+		public List<String> getLocalFields() {
+			return this.localFields;
+		}
+
+		public List<String> getTagFields() {
+			return this.tagFields;
+		}
+
 		public void setRemoteFields(List<String> remoteFields) {
 			this.remoteFields = remoteFields;
 		}
 
+		public void setLocalFields(List<String> localFields) {
+			this.localFields = localFields;
+		}
+
+		public void setTagFields(List<String> tagFields) {
+			this.tagFields = tagFields;
+		}
+
 		public static class Correlation {
 
 			/**

spring-boot-project/spring-boot-actuator-autoconfigure/src/test/java/org/springframework/boot/actuate/autoconfigure/tracing/BaggagePropagationIntegrationTests.java

@@ -227,6 +227,16 @@ public ApplicationContextRunner get() {
 							"management.tracing.baggage.remote-fields=x-vcap-request-id,country-code,bp",
 							"management.tracing.baggage.correlation.fields=country-code,bp");
 			}
+		},
+
+		BRAVE_LOCAL_FIELDS {
+			@Override
+			public ApplicationContextRunner get() {
+				return new ApplicationContextRunner()
+					.withConfiguration(AutoConfigurations.of(BraveAutoConfiguration.class))
+					.withPropertyValues("management.tracing.baggage.local-fields=country-code,bp",
+							"management.tracing.baggage.correlation.fields=country-code,bp");
+			}
 		}
 
 	}

spring-boot-project/spring-boot-actuator-autoconfigure/src/test/java/org/springframework/boot/actuate/autoconfigure/tracing/BraveAutoConfigurationTests.java

@@ -26,6 +26,7 @@
 import brave.SpanCustomizer;
 import brave.Tracer;
 import brave.Tracing;
+import brave.baggage.BaggageField;
 import brave.baggage.BaggagePropagation;
 import brave.baggage.CorrelationScopeConfig.SingleCorrelationField;
 import brave.handler.SpanHandler;
@@ -341,6 +342,22 @@ void compositeSpanHandlerUsesFilterPredicateAndReportersInOrder() {
 		});
 	}
 
+	@Test
+	void shouldCreateTagHandler() {
+		this.contextRunner.withPropertyValues("management.tracing.baggage.tag-fields=country-code,bp")
+			.run((context) -> assertThat(context.getBean(BaggageTagSpanHandler.class)).extracting("fieldsToTag")
+				.asInstanceOf(InstanceOfAssertFactories.array(BaggageField[].class))
+				.extracting(BaggageField::name)
+				.containsOnly("country-code", "bp"));
+	}
+
+	@Test
+	void noopOnNoTagFields() {
+		this.contextRunner.withPropertyValues("management.tracing.baggage.tag-fields=")
+			.run((context) -> assertThat(context.getBean("baggageTagSpanHandler", SpanHandler.class))
+				.isSameAs(SpanHandler.NOOP));
+	}
+
 	private void injectToMap(Map<String, String> map, String key, String value) {
 		map.put(key, value);
 	}