improves traceQueue cleaning logic (#2241)

spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/reactor/TraceReactorAutoConfiguration.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013-2021 the original author or authors.
+ * Copyright 2013-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.
@@ -18,9 +18,6 @@ package org.springframework.cloud.sleuth.autoconfig.instrument.reactor;
 
 import java.io.Closeable;
 import java.io.IOException;
-import java.util.AbstractQueue;
-import java.util.Iterator;
-import java.util.Queue;
 import java.util.function.Function;
 
 import org.apache.commons.logging.Log;
@@ -42,8 +39,6 @@ import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
 import org.springframework.boot.context.properties.EnableConfigurationProperties;
 import org.springframework.cloud.context.scope.refresh.RefreshScope;
 import org.springframework.cloud.context.scope.refresh.RefreshScopeRefreshedEvent;
-import org.springframework.cloud.sleuth.CurrentTraceContext;
-import org.springframework.cloud.sleuth.TraceContext;
 import org.springframework.cloud.sleuth.Tracer;
 import org.springframework.cloud.sleuth.autoconfig.brave.BraveAutoConfiguration;
 import org.springframework.cloud.sleuth.instrument.reactor.ReactorSleuth;
@@ -260,7 +255,7 @@ class HookRegisteringBeanFactoryPostProcessor implements BeanFactoryPostProcesso
 		if (log.isTraceEnabled()) {
 			log.trace("Decorating queues");
 		}
-		Hooks.addQueueWrapper(SLEUTH_TRACE_REACTOR_KEY, queue -> traceQueue(springContext, queue));
+		Hooks.addQueueWrapper(SLEUTH_TRACE_REACTOR_KEY, queue -> ReactorSleuth.traceQueue(springContext, queue));
 	}
 
 	@Override
@@ -274,84 +269,6 @@ class HookRegisteringBeanFactoryPostProcessor implements BeanFactoryPostProcesso
 		Schedulers.resetOnScheduleHook(TraceReactorAutoConfiguration.SLEUTH_REACTOR_EXECUTOR_SERVICE_KEY);
 	}
 
-	private static Queue<?> traceQueue(ConfigurableApplicationContext springContext, Queue<?> queue) {
-		if (!springContext.isActive()) {
-			return queue;
-		}
-		CurrentTraceContext currentTraceContext = springContext.getBean(CurrentTraceContext.class);
-		@SuppressWarnings("unchecked")
-		Queue envelopeQueue = queue;
-		return new AbstractQueue<Object>() {
-
-			@Override
-			public int size() {
-				return envelopeQueue.size();
-			}
-
-			@Override
-			public boolean offer(Object o) {
-				TraceContext traceContext = currentTraceContext.context();
-				return envelopeQueue.offer(new Envelope(o, traceContext));
-			}
-
-			@Override
-			public Object poll() {
-				Object object = envelopeQueue.poll();
-				if (object == null) {
-					// to clear thread-local
-					currentTraceContext.maybeScope(null);
-					return null;
-				}
-				else if (object instanceof Envelope) {
-					Envelope envelope = (Envelope) object;
-					restoreTheContext(envelope);
-					return envelope.body;
-				}
-				return object;
-			}
-
-			private void restoreTheContext(Envelope envelope) {
-				if (envelope.traceContext != null) {
-					currentTraceContext.maybeScope(envelope.traceContext);
-				}
-			}
-
-			@Override
-			public Object peek() {
-				Object peek = queue.peek();
-				if (peek instanceof Envelope) {
-					Envelope envelope = (Envelope) peek;
-					restoreTheContext(envelope);
-					return (envelope).body;
-				}
-				return peek;
-			}
-
-			@Override
-			@SuppressWarnings("unchecked")
-			public Iterator<Object> iterator() {
-				Iterator<?> iterator = queue.iterator();
-				return new Iterator<Object>() {
-					@Override
-					public boolean hasNext() {
-						return iterator.hasNext();
-					}
-
-					@Override
-					public Object next() {
-						Object next = iterator.next();
-						if (next instanceof Envelope) {
-							Envelope envelope = (Envelope) next;
-							restoreTheContext(envelope);
-							return (envelope).body;
-						}
-						return next;
-					}
-				};
-			}
-		};
-	}
-
 	@Override
 	public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
 		if (log.isTraceEnabled()) {
@@ -367,17 +284,4 @@ class HookRegisteringBeanFactoryPostProcessor implements BeanFactoryPostProcesso
 		springContext = (ConfigurableApplicationContext) applicationContext;
 	}
 
-	static class Envelope {
-
-		final Object body;
-
-		final TraceContext traceContext;
-
-		Envelope(Object body, TraceContext traceContext) {
-			this.body = body;
-			this.traceContext = traceContext;
-		}
-
-	}
-
 }

spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ReactorSleuth.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013-2021 the original author or authors.
+ * Copyright 2013-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.
@@ -16,6 +16,9 @@
 
 package org.springframework.cloud.sleuth.instrument.reactor;
 
+import java.util.AbstractQueue;
+import java.util.Iterator;
+import java.util.Queue;
 import java.util.function.BiConsumer;
 import java.util.function.BiFunction;
 import java.util.function.Function;
@@ -323,6 +326,9 @@ public abstract class ReactorSleuth {
 					try (CurrentTraceContext.Scope scope = currentTraceContext.maybeScope(traceContext)) {
 						delegate.run();
 					}
+					// extra step to ensure context is cleared when publishOn or similar
+					// operators leaks different context and leaves it uncleared
+					currentTraceContext.maybeScope(null);
 				};
 			}
 			return delegate;
@@ -632,6 +638,125 @@ public abstract class ReactorSleuth {
 		return newSpan;
 	}
 
+	public static Queue<?> traceQueue(ConfigurableApplicationContext springContext, Queue<?> queue) {
+		if (!springContext.isActive()) {
+			return queue;
+		}
+		CurrentTraceContext currentTraceContext = springContext.getBean(CurrentTraceContext.class);
+		@SuppressWarnings("unchecked")
+		Queue envelopeQueue = queue;
+		return new AbstractQueue<Object>() {
+
+			boolean cleanOnNull;
+
+			boolean hasPrevious = false;
+
+			Thread lastReader;
+
+			@Override
+			public int size() {
+				return envelopeQueue.size();
+			}
+
+			@Override
+			public boolean offer(Object o) {
+				TraceContext traceContext = currentTraceContext.context();
+				return envelopeQueue.offer(new Envelope(o, traceContext));
+			}
+
+			@Override
+			public Object poll() {
+				Object object = envelopeQueue.poll();
+				if (object == null) {
+					if (cleanOnNull) {
+						// to clear thread-local if was just restored
+						currentTraceContext.maybeScope(null);
+					}
+					cleanOnNull = true;
+					lastReader = Thread.currentThread();
+					hasPrevious = false;
+					return null;
+				}
+				else if (object instanceof Envelope) {
+					Envelope envelope = (Envelope) object;
+					restoreTheContext(envelope);
+					hasPrevious = true;
+					return envelope.body;
+				}
+				hasPrevious = true;
+				return object;
+			}
+
+			private void restoreTheContext(Envelope envelope) {
+				TraceContext traceContext = envelope.traceContext;
+				if (traceContext != null) {
+					if (!traceContext.equals(currentTraceContext.context())) {
+						if (!hasPrevious || !Thread.currentThread().equals(this.lastReader)) {
+							// means context was restored form the envelope, thus it has
+							// to be cleared
+							cleanOnNull = true;
+							lastReader = Thread.currentThread();
+						}
+						currentTraceContext.maybeScope(traceContext);
+					}
+					else if (!hasPrevious || !Thread.currentThread().equals(this.lastReader)) {
+						// means same context was already available, no need to clean
+						// anything
+						cleanOnNull = false;
+						lastReader = Thread.currentThread();
+					}
+				}
+			}
+
+			@Override
+			public Object peek() {
+				Object peek = queue.peek();
+				if (peek instanceof Envelope) {
+					Envelope envelope = (Envelope) peek;
+					restoreTheContext(envelope);
+					return (envelope).body;
+				}
+				return peek;
+			}
+
+			@Override
+			@SuppressWarnings("unchecked")
+			public Iterator<Object> iterator() {
+				Iterator<?> iterator = queue.iterator();
+				return new Iterator<Object>() {
+					@Override
+					public boolean hasNext() {
+						return iterator.hasNext();
+					}
+
+					@Override
+					public Object next() {
+						Object next = iterator.next();
+						if (next instanceof Envelope) {
+							Envelope envelope = (Envelope) next;
+							restoreTheContext(envelope);
+							return (envelope).body;
+						}
+						return next;
+					}
+				};
+			}
+		};
+	}
+
+	static class Envelope {
+
+		final Object body;
+
+		final TraceContext traceContext;
+
+		Envelope(Object body, TraceContext traceContext) {
+			this.body = body;
+			this.traceContext = traceContext;
+		}
+
+	}
+
 }
 
 class SleuthContextOperator<T> implements Subscription, CoreSubscriber<T>, Scannable {

tests/brave/spring-cloud-sleuth-instrumentation-reactor-tests/src/test/java/org/springframework/cloud/sleuth/brave/instrument/reactor/QueueWrapperTests.java

@@ -0,0 +1,45 @@
+/*
+ * Copyright 2013-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.cloud.sleuth.brave.instrument.reactor;
+
+import org.springframework.cloud.sleuth.CurrentTraceContext;
+import org.springframework.cloud.sleuth.TraceContext;
+import org.springframework.cloud.sleuth.brave.bridge.BraveAccessor;
+
+/**
+ * @author Oleh Dokuka
+ */
+public class QueueWrapperTests extends org.springframework.cloud.sleuth.instrument.reactor.QueueWrapperTests {
+
+	brave.propagation.CurrentTraceContext traceContext = brave.propagation.CurrentTraceContext.Default.create();
+
+	CurrentTraceContext currentTraceContext = BraveAccessor.currentTraceContext(traceContext);
+
+	TraceContext context = BraveAccessor
+			.traceContext(brave.propagation.TraceContext.newBuilder().traceId(1).spanId(1).sampled(true).build());
+
+	@Override
+	protected CurrentTraceContext currentTraceContext() {
+		return this.currentTraceContext;
+	}
+
+	@Override
+	protected TraceContext context() {
+		return this.context;
+	}
+
+}

tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/QueueWrapperTests.java

@@ -0,0 +1,247 @@
+/*
+ * Copyright 2013-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.cloud.sleuth.instrument.reactor;
+
+import java.util.Objects;
+import java.util.Queue;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+
+import org.assertj.core.api.Assertions;
+import org.assertj.core.presentation.StandardRepresentation;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import reactor.core.publisher.Hooks;
+import reactor.core.scheduler.Schedulers;
+import reactor.util.concurrent.Queues;
+
+import org.springframework.cloud.sleuth.CurrentTraceContext;
+import org.springframework.cloud.sleuth.TraceContext;
+import org.springframework.context.annotation.AnnotationConfigApplicationContext;
+
+import static org.springframework.cloud.sleuth.instrument.reactor.ReactorSleuth.traceQueue;
+
+/**
+ * @author Oleh Dokuka
+ */
+public abstract class QueueWrapperTests {
+
+	static {
+		// AssertJ will recognise QueueSubscription implements queue and try to invoke
+		// iterator. That's not allowed, and will cause an exception
+		// Fuseable$QueueSubscription.NOT_SUPPORTED_MESSAGE.
+		// This ensures AssertJ uses normal toString.
+		StandardRepresentation.registerFormatterForType(ScopePassingSpanSubscriber.class, Objects::toString);
+	}
+
+	protected abstract CurrentTraceContext currentTraceContext();
+
+	protected abstract TraceContext context();
+
+	AnnotationConfigApplicationContext springContext = new AnnotationConfigApplicationContext();
+
+	@BeforeEach
+	public void setup() {
+		Hooks.removeQueueWrappers();
+		Hooks.resetOnLastOperator();
+		Schedulers.resetOnScheduleHooks();
+	}
+
+	@AfterEach
+	public void close() {
+		springContext.close();
+		Hooks.removeQueueWrappers();
+		Hooks.resetOnLastOperator();
+		Schedulers.resetOnScheduleHooks();
+	}
+
+	@Test
+	void checkContextIsRestoredAndOnNullCleaned() {
+		springContext.registerBean(CurrentTraceContext.class, this::currentTraceContext);
+		springContext.refresh();
+
+		final Queue queue = traceQueue(this.springContext, Queues.get(128).get());
+
+		TraceContext context;
+		try (CurrentTraceContext.Scope ws = currentTraceContext().newScope(context())) {
+			context = currentTraceContext().context();
+			queue.offer(1);
+		}
+
+		Assertions.assertThat(queue.poll()).isEqualTo(1);
+		Assertions.assertThat(currentTraceContext().context()).isNotNull().isEqualTo(context);
+
+		Assertions.assertThat(queue.poll()).isNull();
+		Assertions.assertThat(currentTraceContext().context()).isNull();
+	}
+
+	@Test
+	void checkContextIsNotCleanOnNullCleanedIfContextWasAvailableOnThread() {
+		springContext.registerBean(CurrentTraceContext.class, this::currentTraceContext);
+		springContext.refresh();
+
+		final Queue queue = traceQueue(this.springContext, Queues.get(128).get());
+
+		TraceContext context;
+		CurrentTraceContext.Scope ws = currentTraceContext().newScope(context());
+		context = currentTraceContext().context();
+		queue.offer(1);
+
+		Assertions.assertThat(queue.poll()).isEqualTo(1);
+		Assertions.assertThat(currentTraceContext().context()).isNotNull().isEqualTo(context);
+
+		Assertions.assertThat(queue.poll()).isNull();
+		Assertions.assertThat(currentTraceContext().context()).isNotNull().isEqualTo(context);
+
+		ws.close();
+		Assertions.assertThat(currentTraceContext().context()).isNull();
+	}
+
+	@Test
+	void checkContextIsRestoredAndOnNullCleanedInCaseOfSubsequentPolls() {
+		springContext.registerBean(CurrentTraceContext.class, this::currentTraceContext);
+		springContext.refresh();
+
+		final Queue queue = traceQueue(this.springContext, Queues.get(128).get());
+
+		TraceContext context;
+		try (CurrentTraceContext.Scope ws = currentTraceContext().newScope(context())) {
+			context = currentTraceContext().context();
+			queue.offer(1);
+			queue.offer(2);
+		}
+
+		Assertions.assertThat(queue.poll()).isEqualTo(1);
+		Assertions.assertThat(currentTraceContext().context()).isNotNull().isEqualTo(context);
+
+		Assertions.assertThat(queue.poll()).isEqualTo(2);
+		Assertions.assertThat(currentTraceContext().context()).isNotNull().isEqualTo(context);
+
+		Assertions.assertThat(queue.poll()).isNull();
+		Assertions.assertThat(currentTraceContext().context()).isNull();
+	}
+
+	@Test
+	void checkContextIsRestoredAndOnNullCleanedInCaseOfSubsequentPollsByAnotherThread() throws InterruptedException {
+		springContext.registerBean(CurrentTraceContext.class, this::currentTraceContext);
+		springContext.refresh();
+
+		final Queue queue = traceQueue(this.springContext, Queues.get(128).get());
+
+		TraceContext context;
+		try (CurrentTraceContext.Scope ws = currentTraceContext().newScope(context())) {
+			context = currentTraceContext().context();
+			queue.offer(1);
+			queue.offer(2);
+		}
+
+		Assertions.assertThat(queue.poll()).isEqualTo(1);
+		Assertions.assertThat(currentTraceContext().context()).isNotNull().isEqualTo(context);
+
+		CountDownLatch latch = new CountDownLatch(1);
+		new Thread(() -> {
+			Assertions.assertThat(queue.poll()).isEqualTo(2);
+			Assertions.assertThat(currentTraceContext().context()).isNotNull().isEqualTo(context);
+
+			Assertions.assertThat(queue.poll()).isNull();
+			Assertions.assertThat(currentTraceContext().context()).isNull();
+			latch.countDown();
+		}).start();
+
+		Assertions.assertThat(latch.await(10, TimeUnit.SECONDS)).isTrue();
+
+		Assertions.assertThat(queue.poll()).isNull();
+		Assertions.assertThat(currentTraceContext().context()).isNull();
+	}
+
+	@Test
+	void checkContextIsNotCleanOnNullCleanedIfContextWasAvailableOnThreadAnotherThreadCase()
+			throws InterruptedException {
+		springContext.registerBean(CurrentTraceContext.class, this::currentTraceContext);
+		springContext.refresh();
+
+		final Queue queue = traceQueue(this.springContext, Queues.get(128).get());
+
+		TraceContext context;
+		try (CurrentTraceContext.Scope ws = currentTraceContext().newScope(context())) {
+			context = currentTraceContext().context();
+			queue.offer(1);
+			queue.offer(2);
+		}
+
+		Assertions.assertThat(queue.poll()).isEqualTo(1);
+		Assertions.assertThat(currentTraceContext().context()).isNotNull().isEqualTo(context);
+
+		CountDownLatch latch = new CountDownLatch(1);
+		new Thread(() -> {
+			CurrentTraceContext.Scope ws = currentTraceContext().maybeScope(context);
+			Assertions.assertThat(queue.poll()).isEqualTo(2);
+			Assertions.assertThat(currentTraceContext().context()).isNotNull().isEqualTo(context);
+
+			Assertions.assertThat(queue.poll()).isNull();
+			Assertions.assertThat(currentTraceContext().context()).isNotNull().isEqualTo(context);
+
+			ws.close();
+			Assertions.assertThat(currentTraceContext().context()).isNull();
+			latch.countDown();
+		}).start();
+
+		Assertions.assertThat(latch.await(10, TimeUnit.SECONDS)).isTrue();
+
+		Assertions.assertThat(queue.poll()).isNull();
+		Assertions.assertThat(currentTraceContext().context()).isNull();
+	}
+
+	@Test
+	void checkContextIsNotCleanOnNullCleanedIfContextWasAvailableOnThreadAnotherThreadCase2()
+			throws InterruptedException {
+		springContext.registerBean(CurrentTraceContext.class, this::currentTraceContext);
+		springContext.refresh();
+
+		final Queue queue = traceQueue(this.springContext, Queues.get(128).get());
+
+		TraceContext context;
+		CurrentTraceContext.Scope ws1 = currentTraceContext().newScope(context());
+		context = currentTraceContext().context();
+		queue.offer(1);
+		queue.offer(2);
+
+		Assertions.assertThat(queue.poll()).isEqualTo(1);
+		Assertions.assertThat(currentTraceContext().context()).isNotNull().isEqualTo(context);
+
+		ws1.close();
+		Assertions.assertThat(currentTraceContext().context()).isNull();
+
+		CountDownLatch latch = new CountDownLatch(1);
+		new Thread(() -> {
+			Assertions.assertThat(queue.poll()).isEqualTo(2);
+			Assertions.assertThat(currentTraceContext().context()).isNotNull().isEqualTo(context);
+
+			Assertions.assertThat(queue.poll()).isNull();
+			Assertions.assertThat(currentTraceContext().context()).isNull();
+
+			latch.countDown();
+		}).start();
+
+		Assertions.assertThat(latch.await(10, TimeUnit.SECONDS)).isTrue();
+
+		Assertions.assertThat(queue.poll()).isNull();
+		Assertions.assertThat(currentTraceContext().context()).isNull();
+	}
+
+}