From 6a3612ce94e6f3b47e5a4b855b897c0c0581bbb9 Mon Sep 17 00:00:00 2001 From: Oleh Dokuka <5380167+OlegDokuka@users.noreply.github.com> Date: Wed, 4 Jan 2023 00:38:37 +0200 Subject: [PATCH] improves traceQueue cleaning logic (#2241) --- .../TraceReactorAutoConfiguration.java | 100 +------ .../instrument/reactor/ReactorSleuth.java | 127 ++++++++- .../instrument/reactor/QueueWrapperTests.java | 45 ++++ .../instrument/reactor/QueueWrapperTests.java | 247 ++++++++++++++++++ 4 files changed, 420 insertions(+), 99 deletions(-) create mode 100644 tests/brave/spring-cloud-sleuth-instrumentation-reactor-tests/src/test/java/org/springframework/cloud/sleuth/brave/instrument/reactor/QueueWrapperTests.java create mode 100644 tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/QueueWrapperTests.java diff --git a/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/reactor/TraceReactorAutoConfiguration.java b/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/reactor/TraceReactorAutoConfiguration.java index 01aeb8cb4..5c250e6d9 100644 --- a/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/reactor/TraceReactorAutoConfiguration.java +++ b/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() { - - @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 iterator() { - Iterator iterator = queue.iterator(); - return new Iterator() { - @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; - } - - } - } diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ReactorSleuth.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ReactorSleuth.java index 2c4ecff5c..4463d8715 100644 --- a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ReactorSleuth.java +++ b/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() { + + 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 iterator() { + Iterator iterator = queue.iterator(); + return new Iterator() { + @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 implements Subscription, CoreSubscriber, Scannable { diff --git a/tests/brave/spring-cloud-sleuth-instrumentation-reactor-tests/src/test/java/org/springframework/cloud/sleuth/brave/instrument/reactor/QueueWrapperTests.java b/tests/brave/spring-cloud-sleuth-instrumentation-reactor-tests/src/test/java/org/springframework/cloud/sleuth/brave/instrument/reactor/QueueWrapperTests.java new file mode 100644 index 000000000..24515628f --- /dev/null +++ b/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; + } + +} diff --git a/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/QueueWrapperTests.java b/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/QueueWrapperTests.java new file mode 100644 index 000000000..8b4f4c750 --- /dev/null +++ b/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(); + } + +}