diff --git a/micrometer-tracing-tests/micrometer-tracing-test/src/test/java/io/micrometer/tracing/test/contextpropagation/ObservationAwareSpanThreadLocalAccessorTests.java b/micrometer-tracing-tests/micrometer-tracing-test/src/test/java/io/micrometer/tracing/test/contextpropagation/ObservationAwareSpanThreadLocalAccessorTests.java index 2c022d82..25224957 100644 --- a/micrometer-tracing-tests/micrometer-tracing-test/src/test/java/io/micrometer/tracing/test/contextpropagation/ObservationAwareSpanThreadLocalAccessorTests.java +++ b/micrometer-tracing-tests/micrometer-tracing-test/src/test/java/io/micrometer/tracing/test/contextpropagation/ObservationAwareSpanThreadLocalAccessorTests.java @@ -50,11 +50,12 @@ class ObservationAwareSpanThreadLocalAccessorTests { ExecutorService executorService = ContextExecutorService.wrap(Executors.newSingleThreadExecutor(), () -> ContextSnapshot.captureAll(contextRegistry)); + ObservationAwareSpanThreadLocalAccessor accessor = new ObservationAwareSpanThreadLocalAccessor(tracer); + @BeforeEach void setup() { observationRegistry.observationConfig().observationHandler(new DefaultTracingObservationHandler(tracer)); - contextRegistry.loadThreadLocalAccessors() - .registerThreadLocalAccessor(new ObservationAwareSpanThreadLocalAccessor(tracer)); + contextRegistry.loadThreadLocalAccessors().registerThreadLocalAccessor(accessor); } @AfterEach @@ -76,7 +77,7 @@ void asyncTracingTestWithObservationAndManualSpans() try (Tracer.SpanInScope scope2 = this.tracer.withSpan(secondSpan.start())) { logWithSpan("Async in test with span - before call"); Future future = executorService.submit(this::asyncCall); - String spanIdFromFuture = future.get(1, TimeUnit.SECONDS); + String spanIdFromFuture = future.get(100, TimeUnit.SECONDS); logWithSpan("Async in test with span - after call"); then(spanIdFromFuture).isEqualTo(secondSpan.context().spanId()); } diff --git a/micrometer-tracing/src/main/java/io/micrometer/tracing/contextpropagation/ObservationAwareSpanThreadLocalAccessor.java b/micrometer-tracing/src/main/java/io/micrometer/tracing/contextpropagation/ObservationAwareSpanThreadLocalAccessor.java index 98e1530e..4b3f5296 100644 --- a/micrometer-tracing/src/main/java/io/micrometer/tracing/contextpropagation/ObservationAwareSpanThreadLocalAccessor.java +++ b/micrometer-tracing/src/main/java/io/micrometer/tracing/contextpropagation/ObservationAwareSpanThreadLocalAccessor.java @@ -58,7 +58,7 @@ public class ObservationAwareSpanThreadLocalAccessor implements ThreadLocalAcces private static final InternalLogger log = InternalLoggerFactory .getInstance(ObservationAwareSpanThreadLocalAccessor.class); - private final Map spanActions = new ConcurrentHashMap<>(); + final Map spanActions = new ConcurrentHashMap<>(); /** * Key under which Micrometer Tracing is being registered. @@ -98,17 +98,10 @@ public Span getValue() { if (currentSpan != null && !currentSpan.equals(tracingContext.getSpan())) { // User created child spans manually and scoped them // the current span is not the same as the one from observation - spanActions.put(Thread.currentThread(), - new SpanAction(SpanSituation.OBSERVATION_AND_MANUAL_SPAN_PRESENT, spanActions)); return currentSpan; } - // Current span is same as the one from observation, we will skip this - spanActions.put(Thread.currentThread(), - new SpanAction(SpanSituation.SPAN_SAME_AS_OBSERVATION_SO_SKIP, spanActions)); return null; } - // No current observation so let's check the tracer - spanActions.put(Thread.currentThread(), new SpanAction(SpanSituation.NO_OBSERVATION_PRESENT, spanActions)); return this.tracer.currentSpan(); } @@ -117,7 +110,8 @@ public void setValue(Span value) { SpanAction spanAction = spanActions.get(Thread.currentThread()); Tracer.SpanInScope scope = this.tracer.withSpan(value); if (spanAction == null) { - spanAction = new SpanAction(SpanSituation.NO_OBSERVATION_PRESENT, spanActions); + spanAction = new SpanAction(spanActions); + spanActions.put(Thread.currentThread(), spanAction); } spanAction.setScope(scope); } @@ -125,7 +119,7 @@ public void setValue(Span value) { @Override public void setValue() { SpanAction spanAction = spanActions.get(Thread.currentThread()); - if (spanAction == null || spanAction.spanSituation == SpanSituation.SPAN_SAME_AS_OBSERVATION_SO_SKIP) { + if (spanAction == null) { return; } Tracer.SpanInScope scope = this.tracer.withSpan(null); @@ -135,7 +129,7 @@ public void setValue() { @Override public void restore(Span previousValue) { SpanAction spanAction = spanActions.get(Thread.currentThread()); - if (spanAction == null || spanAction.spanSituation == SpanSituation.SPAN_SAME_AS_OBSERVATION_SO_SKIP) { + if (spanAction == null) { return; } spanAction.close(); @@ -152,24 +146,20 @@ public void restore(Span previousValue) { @Override public void restore() { SpanAction spanAction = spanActions.get(Thread.currentThread()); - if (spanAction == null || spanAction.spanSituation == SpanSituation.SPAN_SAME_AS_OBSERVATION_SO_SKIP) { - return; + if (spanAction != null) { + spanAction.close(); } - spanAction.close(); } static class SpanAction implements Closeable { - final SpanSituation spanSituation; - final SpanAction previous; final Map todo; Closeable scope; - SpanAction(SpanSituation spanSituation, Map spanActions) { - this.spanSituation = spanSituation; + SpanAction(Map spanActions) { this.previous = spanActions.get(Thread.currentThread()); this.todo = spanActions; } @@ -198,10 +188,4 @@ public void close() { } - enum SpanSituation { - - OBSERVATION_AND_MANUAL_SPAN_PRESENT, NO_OBSERVATION_PRESENT, SPAN_SAME_AS_OBSERVATION_SO_SKIP - - } - }