From 6da18e9be24ba822fdc2c28d9e15e642ae737e41 Mon Sep 17 00:00:00 2001 From: Marcin Grzejszczak Date: Wed, 2 Feb 2022 14:42:52 +0100 Subject: [PATCH] Fixes invalid ThreadLocalSpan stacking and tracing context leaks; fixes gh-2064; fixes gh-2108 --- .../messaging/TracingChannelInterceptor.java | 107 ++++++++++++------ .../messaging/ThreadLocalSpanTests.java | 73 ++++++++++++ 2 files changed, 147 insertions(+), 33 deletions(-) create mode 100644 spring-cloud-sleuth-instrumentation/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/ThreadLocalSpanTests.java diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/TracingChannelInterceptor.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/TracingChannelInterceptor.java index 21447c6e5..8e4caa5ba 100644 --- a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/TracingChannelInterceptor.java +++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/TracingChannelInterceptor.java @@ -16,11 +16,14 @@ package org.springframework.cloud.sleuth.instrument.messaging; -import java.util.NoSuchElementException; +import java.io.Closeable; +import java.util.ArrayDeque; import java.util.Set; -import java.util.concurrent.LinkedBlockingDeque; import java.util.function.Function; +import org.apache.commons.logging.Log; +import org.apache.commons.logging.LogFactory; + import org.springframework.aop.support.AopUtils; import org.springframework.beans.BeansException; import org.springframework.cloud.sleuth.Span; @@ -90,7 +93,7 @@ public final class TracingChannelInterceptor implements ExecutorChannelIntercept private static final Class directWithAttributesChannelClass = ClassUtils.isPresent(STREAM_DIRECT_CHANNEL, null) ? ClassUtils.resolveClassName(STREAM_DIRECT_CHANNEL, null) : null; - private final ThreadLocalSpan threadLocalSpan = new ThreadLocalSpan(); + private final ThreadLocalSpan threadLocalSpan; private final Tracer tracer; @@ -115,6 +118,7 @@ public final class TracingChannelInterceptor implements ExecutorChannelIntercept this.extractor = getter; this.remoteServiceNameMapper = remoteServiceNameMapper; this.messageSpanCustomizer = messageSpanCustomizer; + this.threadLocalSpan = new ThreadLocalSpan(tracer); } @Override @@ -148,8 +152,7 @@ public final class TracingChannelInterceptor implements ExecutorChannelIntercept } private void setSpanInScope(Span span) { - Tracer.SpanInScope spanInScope = this.tracer.withSpan(span); - this.threadLocalSpan.set(new SpanAndScope(span, spanInScope)); + this.threadLocalSpan.set(span); log.debug(() -> "Put span in scope " + span); } @@ -308,8 +311,8 @@ public final class TracingChannelInterceptor implements ExecutorChannelIntercept if (spanAndScope == null) { return; } - Span span = spanAndScope.span; - Tracer.SpanInScope scope = spanAndScope.scope; + Span span = spanAndScope.getSpan(); + Tracer.SpanInScope scope = spanAndScope.getScope(); if (span.isNoop()) { log.debug(() -> "Span " + span + " is noop - will stop the scope"); scope.close(); @@ -354,55 +357,93 @@ public final class TracingChannelInterceptor implements ExecutorChannelIntercept return message; } - private static class SpanAndScope { + static class SpanAndScope implements Closeable { - final Span span; + private static final Log log = LogFactory.getLog(SpanAndScope.class); - final Tracer.SpanInScope scope; + private final Span span; + + private final Tracer.SpanInScope scope; SpanAndScope(Span span, Tracer.SpanInScope scope) { this.span = span; this.scope = scope; } + public Span getSpan() { + return this.span; + } + + public Tracer.SpanInScope getScope() { + return this.scope; + } + + @Override + public String toString() { + return "SpanAndScope{" + "span=" + this.span + '}'; + } + + @Override + public void close() { + if (log.isDebugEnabled()) { + log.debug("Closing span [" + this.span + "], scope is not null [" + (this.scope != null) + "]"); + } + if (this.scope != null) { + this.scope.close(); + } + this.span.end(); + } + } - private static class ThreadLocalSpan { + static class ThreadLocalSpan { - private static final LogAccessor log = new LogAccessor(ThreadLocalSpan.class); + private final ThreadLocal> currentSpanInScopeStack = new ThreadLocal<>(); - private final ThreadLocal threadLocalSpan = new ThreadLocal<>(); + private final Tracer tracer; - private final LinkedBlockingDeque spans = new LinkedBlockingDeque<>(); - - ThreadLocalSpan() { + ThreadLocalSpan(Tracer tracer) { + this.tracer = tracer; } - void set(SpanAndScope spanAndScope) { - SpanAndScope scope = this.threadLocalSpan.get(); - if (scope != null) { - this.spans.addFirst(scope); - } - this.threadLocalSpan.set(spanAndScope); + /** + * Sets given span and scope. + * @param span - span to be put in scope + */ + public void set(Span span) { + Tracer.SpanInScope spanInScope = this.tracer.withSpan(span); + SpanAndScope newSpanAndScope = new SpanAndScope(span, spanInScope); + getCurrentSpanInScopeStack().addFirst(newSpanAndScope); } - SpanAndScope get() { - return this.threadLocalSpan.get(); + /** + * @return currently stored span and scope + */ + public SpanAndScope get() { + return getCurrentSpanInScopeStack().peekFirst(); } - void remove() { - this.threadLocalSpan.remove(); - if (this.spans.isEmpty()) { + /** + * Removes the current span from thread local and brings back the previous span to + * the current thread local. + */ + public void remove() { + SpanAndScope spanAndScope = getCurrentSpanInScopeStack().pollFirst(); + if (spanAndScope == null) { return; } - try { - SpanAndScope span = this.spans.removeFirst(); - log.debug(() -> "Took span [" + span + "] from thread local"); - this.threadLocalSpan.set(span); + if (spanAndScope.getScope() != null) { + spanAndScope.getScope().close(); } - catch (NoSuchElementException ex) { - log.trace(ex, () -> "Failed to remove a span from the queue"); + } + + private ArrayDeque getCurrentSpanInScopeStack() { + ArrayDeque stack = this.currentSpanInScopeStack.get(); + if (stack == null) { + stack = new ArrayDeque<>(); + this.currentSpanInScopeStack.set(stack); } + return stack; } } diff --git a/spring-cloud-sleuth-instrumentation/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/ThreadLocalSpanTests.java b/spring-cloud-sleuth-instrumentation/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/ThreadLocalSpanTests.java new file mode 100644 index 000000000..a4653e1fa --- /dev/null +++ b/spring-cloud-sleuth-instrumentation/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/ThreadLocalSpanTests.java @@ -0,0 +1,73 @@ +/* + * Copyright 2013-2021 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.messaging; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.BDDMockito; + +import org.springframework.cloud.sleuth.Span; +import org.springframework.cloud.sleuth.Tracer; + +import static org.assertj.core.api.BDDAssertions.then; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.BDDMockito.given; + +class ThreadLocalSpanTests { + + Tracer tracer = BDDMockito.mock(Tracer.class); + + @BeforeEach + void setup() { + given(this.tracer.withSpan(any())).willReturn(() -> { + + }); + } + + @Test + void should_properly_stack_spans() { + // given + TracingChannelInterceptor.ThreadLocalSpan threadLocalSpan = new TracingChannelInterceptor.ThreadLocalSpan( + tracer); + then(threadLocalSpan.get()).isNull(); + + // when - Span 1 + Span span = span(); + threadLocalSpan.set(span); + // then - Span 1 + then(threadLocalSpan.get().getSpan()).isSameAs(span); + + // when - Span 2 + Span secondSpan = span(); + threadLocalSpan.set(secondSpan); + // then - Span 2 + then(threadLocalSpan.get().getSpan()).isSameAs(secondSpan); + + // expect - Span 1 + threadLocalSpan.remove(); + then(threadLocalSpan.get().getSpan()).isSameAs(span); + + // expect - null + threadLocalSpan.remove(); + then(threadLocalSpan.get()).isNull(); + } + + private Span span() { + return BDDMockito.mock(Span.class); + } + +}