Streamline the interactions with the MDC and the OTel Context. (#34)

* Streamline the interactions with the MDC and the OTel Context.
Resolves #28

* remove an errant println

* formatting

Author: jkwatson

spring-cloud-sleuth-otel-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/otel/OtelLogConfiguration.java

@@ -19,14 +19,13 @@
 import io.opentelemetry.exporter.logging.LoggingSpanExporter;
 import org.slf4j.MDC;
 
-import org.springframework.beans.factory.BeanFactory;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
 import org.springframework.boot.context.properties.EnableConfigurationProperties;
 import org.springframework.cloud.sleuth.autoconfig.SleuthBaggageProperties;
 import org.springframework.cloud.sleuth.otel.bridge.Slf4jApplicationListener;
-import org.springframework.cloud.sleuth.otel.bridge.Slf4jBaggageSpanProcessor;
+import org.springframework.cloud.sleuth.otel.bridge.Slf4jBaggageApplicationListener;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 
@@ -55,9 +54,8 @@ Slf4jApplicationListener otelSlf4jApplicationListener() {
 
 		@Bean
 		@ConditionalOnMissingBean
-		Slf4jBaggageSpanProcessor otelSlf4jBaggageSpanProcessor(BeanFactory beanFactory,
-				SleuthBaggageProperties baggageProperties) {
-			return new Slf4jBaggageSpanProcessor(baggageProperties.getCorrelationFields(), beanFactory);
+		Slf4jBaggageApplicationListener otelSlf4jBaggageApplicationListener(SleuthBaggageProperties baggageProperties) {
+			return new Slf4jBaggageApplicationListener(baggageProperties.getCorrelationFields());
 		}
 
 	}

spring-cloud-sleuth-otel/src/main/java/org/springframework/cloud/sleuth/otel/bridge/OtelBaggageInScope.java

@@ -28,8 +28,6 @@
 import org.springframework.cloud.sleuth.BaggageInScope;
 import org.springframework.cloud.sleuth.CurrentTraceContext;
 import org.springframework.cloud.sleuth.TraceContext;
-import org.springframework.context.ApplicationEvent;
-import org.springframework.context.ApplicationEventPublisher;
 
 /**
  * OpenTelemetry implementation of a {@link BaggageInScope}.
@@ -43,19 +41,16 @@ class OtelBaggageInScope implements BaggageInScope {
 
 	private final CurrentTraceContext currentTraceContext;
 
-	private final ApplicationEventPublisher publisher;
-
 	private final List<String> tagFields;
 
 	private final AtomicReference<Entry> entry = new AtomicReference<>();
 
 	private final AtomicReference<Scope> scope = new AtomicReference<>();
 
 	OtelBaggageInScope(OtelBaggageManager otelBaggageManager, CurrentTraceContext currentTraceContext,
-			ApplicationEventPublisher publisher, List<String> tagFields, Entry entry) {
+			List<String> tagFields, Entry entry) {
 		this.otelBaggageManager = otelBaggageManager;
 		this.currentTraceContext = currentTraceContext;
-		this.publisher = publisher;
 		this.tagFields = tagFields;
 		this.entry.set(entry);
 	}
@@ -109,7 +104,6 @@ private BaggageInScope doSet(TraceContext context, String value) {
 		if (this.tagFields.stream().map(String::toLowerCase).anyMatch(s -> s.equals(entry().getKey()))) {
 			currentSpan.setAttribute(entry().getKey(), value);
 		}
-		this.publisher.publishEvent(new BaggageChanged(this, baggage, entry().getKey(), value));
 		Entry previous = entry();
 		this.entry.set(new Entry(previous.getKey(), value, previous.getMetadata()));
 		return this;
@@ -143,35 +137,4 @@ public void close() {
 		}
 	}
 
-	static class BaggageChanged extends ApplicationEvent {
-
-		/**
-		 * Baggage with the new entry.
-		 */
-		Baggage baggage;
-
-		/**
-		 * Baggage entry name.
-		 */
-		String name;
-
-		/**
-		 * Baggage entry value.
-		 */
-		String value;
-
-		BaggageChanged(OtelBaggageInScope source, Baggage baggage, String name, String value) {
-			super(source);
-			this.baggage = baggage;
-			this.name = name;
-			this.value = value;
-		}
-
-		@Override
-		public String toString() {
-			return "BaggageChanged{" + "name='" + name + '\'' + ", value='" + value + '\'' + '}';
-		}
-
-	}
-
 }

spring-cloud-sleuth-otel/src/main/java/org/springframework/cloud/sleuth/otel/bridge/OtelBaggageManager.java

@@ -137,7 +137,7 @@ private Entry entryForName(String name, io.opentelemetry.api.baggage.Baggage bag
 	}
 
 	private BaggageInScope otelBaggage(Entry entry) {
-		return new OtelBaggageInScope(this, this.currentTraceContext, this.publisher, this.tagFields, entry);
+		return new OtelBaggageInScope(this, this.currentTraceContext, this.tagFields, entry);
 	}
 
 	@Override
@@ -157,7 +157,7 @@ private BaggageInScope baggageWithValue(String name, String value) {
 				.anyMatch(s -> s.equals(name.toLowerCase()));
 		BaggageEntryMetadata entryMetadata = BaggageEntryMetadata.create(propagationString(remoteField));
 		Entry entry = new Entry(name, value, entryMetadata);
-		return new OtelBaggageInScope(this, this.currentTraceContext, this.publisher, this.tagFields, entry);
+		return new OtelBaggageInScope(this, this.currentTraceContext, this.tagFields, entry);
 	}
 
 	private String propagationString(boolean remoteField) {

spring-cloud-sleuth-otel/src/main/java/org/springframework/cloud/sleuth/otel/bridge/OtelCurrentTraceContext.java

@@ -46,17 +46,16 @@ public OtelCurrentTraceContext(ApplicationEventPublisher publisher) {
 		ContextStorage.addWrapper(contextStorage -> new ContextStorage() {
 			@Override
 			public io.opentelemetry.context.Scope attach(Context context) {
-				Span currentSpan = Span.fromContextOrNull(Context.current());
+				Context currentContext = Context.current();
 				io.opentelemetry.context.Scope scope = contextStorage.attach(context);
 				if (scope == io.opentelemetry.context.Scope.noop()) {
 					return scope;
 				}
-				Span attachingSpan = Span.fromContext(context);
-				publisher.publishEvent(new ScopeAttached(this, attachingSpan));
+				publisher.publishEvent(new ScopeAttached(this, context));
 				return () -> {
 					scope.close();
 					publisher.publishEvent(new ScopeClosed(this));
-					publisher.publishEvent(new ScopeRestored(this, currentSpan));
+					publisher.publishEvent(new ScopeRestored(this, currentContext));
 				};
 			}
 
@@ -147,22 +146,30 @@ static class ScopeAttached extends ApplicationEvent {
 		/**
 		 * Span corresponding to the attached scope. Might be {@code null}.
 		 */
-		final Span span;
+		final Context context;
 
 		/**
 		 * Create a new {@code ApplicationEvent}.
 		 * @param source the object on which the event initially occurred or with which
 		 * the event is associated (never {@code null})
-		 * @param span corresponding trace context
+		 * @param context corresponding trace context
 		 */
-		ScopeAttached(Object source, @Nullable Span span) {
+		ScopeAttached(Object source, @Nullable Context context) {
 			super(source);
-			this.span = span;
+			this.context = context;
+		}
+
+		Span getSpan() {
+			return Span.fromContextOrNull(context);
+		}
+
+		Baggage getBaggage() {
+			return Baggage.fromContextOrNull(context);
 		}
 
 		@Override
 		public String toString() {
-			return "ScopeAttached{span=" + span + "}";
+			return "ScopeAttached{context: [span: " + getSpan() + "] [baggage: " + getBaggage() + "]}";
 		}
 
 	}
@@ -172,22 +179,30 @@ static class ScopeRestored extends ApplicationEvent {
 		/**
 		 * Span corresponding to the scope being restored. Might be {@code null}.
 		 */
-		final Span span;
+		final Context context;
 
 		/**
 		 * Create a new {@code ApplicationEvent}.
 		 * @param source the object on which the event initially occurred or with which
 		 * the event is associated (never {@code null})
-		 * @param span corresponding trace context
+		 * @param context corresponding trace context
 		 */
-		ScopeRestored(Object source, @Nullable Span span) {
+		ScopeRestored(Object source, @Nullable Context context) {
 			super(source);
-			this.span = span;
+			this.context = context;
+		}
+
+		Span getSpan() {
+			return Span.fromContextOrNull(context);
+		}
+
+		Baggage getBaggage() {
+			return Baggage.fromContextOrNull(context);
 		}
 
 		@Override
 		public String toString() {
-			return "ScopeRestored{span=" + span + "}";
+			return "ScopeRestored{context: [span: " + getSpan() + "] [baggage: " + getBaggage() + "]}";
 		}
 
 	}

spring-cloud-sleuth-otel/src/main/java/org/springframework/cloud/sleuth/otel/bridge/Slf4jApplicationListener.java

@@ -16,6 +16,7 @@
 
 package org.springframework.cloud.sleuth.otel.bridge;
 
+import io.opentelemetry.api.trace.Span;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.slf4j.MDC;
@@ -31,19 +32,21 @@ private void onScopeAttached(OtelCurrentTraceContext.ScopeAttached event) {
 		if (log.isTraceEnabled()) {
 			log.trace("Got scope changed event [" + event + "]");
 		}
-		if (event.span != null) {
-			MDC.put("traceId", event.span.getSpanContext().getTraceId());
-			MDC.put("spanId", event.span.getSpanContext().getSpanId());
+		Span span = event.getSpan();
+		if (span != null) {
+			MDC.put("traceId", span.getSpanContext().getTraceId());
+			MDC.put("spanId", span.getSpanContext().getSpanId());
 		}
 	}
 
 	private void onScopeRestored(OtelCurrentTraceContext.ScopeRestored event) {
 		if (log.isTraceEnabled()) {
-			log.trace("Got scope changed event [" + event + "]");
+			log.trace("Got scope restored event [" + event + "]");
 		}
-		if (event.span != null) {
-			MDC.put("traceId", event.span.getSpanContext().getTraceId());
-			MDC.put("spanId", event.span.getSpanContext().getSpanId());
+		Span span = event.getSpan();
+		if (span != null) {
+			MDC.put("traceId", span.getSpanContext().getTraceId());
+			MDC.put("spanId", span.getSpanContext().getSpanId());
 		}
 	}
 

spring-cloud-sleuth-otel/src/main/java/org/springframework/cloud/sleuth/otel/bridge/Slf4jBaggageApplicationListener.java

@@ -0,0 +1,90 @@
+/*
+ * Copyright 2013-2020 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.cloud.sleuth.otel.bridge;
+
+import java.util.List;
+import java.util.stream.Collectors;
+
+import io.opentelemetry.api.baggage.Baggage;
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+import org.slf4j.MDC;
+
+import org.springframework.context.ApplicationEvent;
+import org.springframework.context.ApplicationListener;
+
+public class Slf4jBaggageApplicationListener implements ApplicationListener<ApplicationEvent> {
+
+	private static final Log log = LogFactory.getLog(Slf4jBaggageApplicationListener.class);
+
+	private final List<String> lowerCaseCorrelationFields;
+
+	private final List<String> correlationFields;
+
+	public Slf4jBaggageApplicationListener(List<String> correlationFields) {
+		this.lowerCaseCorrelationFields = correlationFields.stream().map(String::toLowerCase)
+				.collect(Collectors.toList());
+		this.correlationFields = correlationFields;
+	}
+
+	private void onScopeAttached(OtelCurrentTraceContext.ScopeAttached event) {
+		if (log.isTraceEnabled()) {
+			log.trace("Got scope attached event [" + event + "]");
+		}
+		if (event.getBaggage() != null) {
+			putEntriesIntoMdc(event.getBaggage());
+		}
+	}
+
+	private void onScopeRestored(OtelCurrentTraceContext.ScopeRestored event) {
+		if (log.isTraceEnabled()) {
+			log.trace("Got scope restored event [" + event + "]");
+		}
+		if (event.getBaggage() != null) {
+			putEntriesIntoMdc(event.getBaggage());
+		}
+	}
+
+	private void putEntriesIntoMdc(Baggage baggage) {
+		baggage.forEach((key, baggageEntry) -> {
+			if (lowerCaseCorrelationFields.contains(key.toLowerCase())) {
+				MDC.put(key, baggageEntry.getValue());
+			}
+		});
+	}
+
+	private void onScopeClosed(OtelCurrentTraceContext.ScopeClosed event) {
+		if (log.isTraceEnabled()) {
+			log.trace("Got scope closed event [" + event + "]");
+		}
+		correlationFields.forEach(MDC::remove);
+	}
+
+	@Override
+	public void onApplicationEvent(ApplicationEvent event) {
+		if (event instanceof OtelCurrentTraceContext.ScopeAttached) {
+			onScopeAttached((OtelCurrentTraceContext.ScopeAttached) event);
+		}
+		else if (event instanceof OtelCurrentTraceContext.ScopeClosed) {
+			onScopeClosed((OtelCurrentTraceContext.ScopeClosed) event);
+		}
+		else if (event instanceof OtelCurrentTraceContext.ScopeRestored) {
+			onScopeRestored((OtelCurrentTraceContext.ScopeRestored) event);
+		}
+	}
+
+}