Improved handling of Jaeger spans/scopes across threads

common/context/src/main/java/io/helidon/common/context/ContextAwareExecutorImpl.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019, 2020 Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2019, 2021 Oracle and/or its affiliates.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -134,7 +134,7 @@ protected <T> Callable<T> wrap(Callable<T> task) {
                     PROVIDERS.forEach(provider -> provider.propagateData(properties.get(provider.getClass())));
                     return Contexts.runInContext(context.get(), task);
                 } finally {
-                    PROVIDERS.forEach(DataPropagationProvider::clearData);
+                    PROVIDERS.forEach(provider -> provider.clearData(properties.get(provider.getClass())));
                 }
             };
         } else {
@@ -153,7 +153,7 @@ protected Runnable wrap(Runnable command) {
                     PROVIDERS.forEach(provider -> provider.propagateData(properties.get(provider.getClass())));
                     Contexts.runInContext(context.get(), command);
                 } finally {
-                    PROVIDERS.forEach(DataPropagationProvider::clearData);
+                    PROVIDERS.forEach(provider -> provider.clearData(properties.get(provider.getClass())));
                 }
             };
         } else {

common/context/src/main/java/io/helidon/common/context/spi/DataPropagationProvider.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020 Oracle and/or its affiliates.
+ * Copyright (c) 2020, 2021 Oracle and/or its affiliates.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -42,7 +42,9 @@ public interface DataPropagationProvider<T> {
 
     /**
      * Clears the propagated date from the new thread when it finishes.
+     *
+     * @param data dataa for propagation
      */
-    void clearData();
+    void clearData(T data);
 
 }

logging/jul/src/main/java/io/helidon/logging/jul/JulMdcPropagator.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020 Oracle and/or its affiliates.
+ * Copyright (c) 2020, 2021 Oracle and/or its affiliates.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -36,7 +36,7 @@ public void propagateData(Map<String, String> data) {
     }
 
     @Override
-    public void clearData() {
+    public void clearData(Map<String, String> data) {
         JulMdc.clear();
     }
 

logging/log4j/src/main/java/io/helidon/logging/log4j/Log4jMdcPropagator.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020 Oracle and/or its affiliates.
+ * Copyright (c) 2020, 2021 Oracle and/or its affiliates.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -38,7 +38,7 @@ public void propagateData(Map<String, String> data) {
     }
 
     @Override
-    public void clearData() {
+    public void clearData(Map<String, String> data) {
         ThreadContext.clearAll();
     }
 }

logging/slf4j/src/main/java/io/helidon/logging/slf4j/Slf4jMdcPropagator.java

@@ -39,7 +39,7 @@ public void propagateData(Map<String, String> data) {
     }
 
     @Override
-    public void clearData() {
+    public void clearData(Map<String, String> data) {
         MDC.clear();
     }
 

tracing/jaeger/pom.xml

@@ -46,6 +46,10 @@
             <groupId>io.helidon.common</groupId>
             <artifactId>helidon-common</artifactId>
         </dependency>
+        <dependency>
+            <groupId>io.helidon.common</groupId>
+            <artifactId>helidon-common-context</artifactId>
+        </dependency>
         <dependency>
             <groupId>io.helidon.config</groupId>
             <artifactId>helidon-config</artifactId>

tracing/jaeger/src/main/java/io/helidon/tracing/jaeger/JaegerDataPropagationProvider.java

@@ -0,0 +1,86 @@
+/*
+ * Copyright (c) 2021 Oracle and/or its affiliates.
+ *
+ * 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
+ *
+ *     http://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 io.helidon.tracing.jaeger;
+
+import io.helidon.common.context.Contexts;
+import io.helidon.common.context.spi.DataPropagationProvider;
+
+import io.opentracing.Scope;
+import io.opentracing.Span;
+import io.opentracing.Tracer;
+import io.opentracing.util.GlobalTracer;
+
+/**
+ * A data propagation provider for Jaeger. Makes sure span are properly propagated
+ * across threads managed by {@link io.helidon.common.context.ContextAwareExecutorService}.
+ */
+public class JaegerDataPropagationProvider implements DataPropagationProvider<JaegerDataPropagationProvider.JaegerContext> {
+
+    static class JaegerContext {
+        private final Span span;
+        private final Tracer tracer;
+        private Scope scope;
+
+        JaegerContext(Tracer tracer, Span span) {
+            this.tracer = tracer;
+            this.span = span;
+        }
+
+        Scope scope() {
+            return scope;
+        }
+    }
+
+    /**
+     * Closes scope in primary thread and returns a context to activate
+     * new scope in secondary thread.
+     *
+     * @return active span.
+     */
+    @Override
+    public JaegerContext data() {
+        return Contexts.context().map(context -> {
+            context.get(Scope.class).ifPresent(Scope::close);
+            return context.get(Span.class).map(span -> {
+                Tracer tracer = context.get(Tracer.class).orElseGet(GlobalTracer::get);
+                return new JaegerContext(tracer, span);
+            }).orElse(null);
+        }).orElse(null);
+    }
+
+    /**
+     * Activates scope in secondary thread.
+     *
+     * @param context the context.
+     */
+    @Override
+    public void propagateData(JaegerContext context) {
+        if (context != null) {
+            context.scope = context.tracer.scopeManager().activate(context.span);
+        }
+    }
+
+    /**
+     * Closes scope in secondary thread.
+     */
+    @Override
+    public void clearData(JaegerContext context) {
+        if (context != null && context.scope != null) {
+            context.scope.close();
+        }
+    }
+}

tracing/jaeger/src/main/java/module-info.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019, 2020 Oracle and/or its affiliates.
+ * Copyright (c) 2019, 2021 Oracle and/or its affiliates.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -21,6 +21,7 @@
     requires io.helidon.common;
     requires io.helidon.config;
     requires io.helidon.tracing;
+    requires io.helidon.common.context;
 
     requires java.logging;
     requires io.opentracing.util;
@@ -30,8 +31,9 @@
     // need to explicitly require transitive dependency, as jaeger is not a module
     requires com.google.gson;
 
-
     exports io.helidon.tracing.jaeger;
 
     provides io.helidon.tracing.spi.TracerProvider with io.helidon.tracing.jaeger.JaegerTracerProvider;
+    provides io.helidon.common.context.spi.DataPropagationProvider with io.helidon.tracing.jaeger.JaegerDataPropagationProvider;
+
 }

tracing/jaeger/src/test/java/io/helidon/tracing/jaeger/JaegerDataPropagationProviderTest.java

@@ -0,0 +1,137 @@
+/*
+ * Copyright (c) 2021 Oracle and/or its affiliates.
+ *
+ * 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
+ *
+ *     http://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 io.helidon.tracing.jaeger;
+
+import io.helidon.common.context.Context;
+import io.helidon.common.context.Contexts;
+
+import io.opentracing.Scope;
+import io.opentracing.ScopeManager;
+import io.opentracing.Span;
+import io.opentracing.SpanContext;
+import io.opentracing.Tracer;
+import io.opentracing.propagation.Format;
+
+import org.junit.jupiter.api.Test;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.CoreMatchers.is;
+
+class JaegerDataPropagationProviderTest {
+
+    private JaegerDataPropagationProvider provider = new JaegerDataPropagationProvider();
+
+    @Test
+    void dataPropagationTest() {
+        Context context = Context.create();
+        Tracer tracer = new TestTracer();
+        context.register(tracer);
+        Span span = tracer.buildSpan("span").start();
+        context.register(span);
+        Scope scope = tracer.scopeManager().activate(span);
+        context.register(scope);
+
+        Contexts.runInContext(context, () -> {
+            assertThat(closed(scope), is(false));
+            JaegerDataPropagationProvider.JaegerContext data = provider.data();
+            assertThat(closed(scope), is(true));
+            provider.propagateData(data);
+            assertThat(closed(data.scope()), is(false));
+            provider.clearData(data);
+            assertThat(closed(data.scope()), is(true));
+        });
+    }
+
+    private boolean closed(Scope scope) {
+        return ((TestScope) scope).closed;
+    }
+
+    static class TestScope implements Scope {
+
+        private final Scope delegate;
+        private boolean closed = false;
+
+        TestScope(Scope delegate) {
+            this.delegate = delegate;
+        }
+
+        @Override
+        public void close() {
+            delegate.close();
+            closed = true;
+        }
+    }
+
+    static class TestScopeManager implements ScopeManager {
+
+        private final ScopeManager delegate;
+
+        TestScopeManager(ScopeManager delegate) {
+            this.delegate = delegate;
+        }
+
+        @Override
+        public Scope activate(Span span) {
+            return new TestScope(delegate.activate(span));
+        }
+
+        @Override
+        public Span activeSpan() {
+            return delegate.activeSpan();
+        }
+    }
+
+    static class TestTracer implements Tracer {
+
+        private final Tracer delegate = JaegerTracerBuilder.create().serviceName("test-service").build();
+
+        @Override
+        public ScopeManager scopeManager() {
+            return new TestScopeManager(delegate.scopeManager());
+        }
+
+        @Override
+        public Span activeSpan() {
+            return delegate.activeSpan();
+        }
+
+        @Override
+        public Scope activateSpan(Span span) {
+            return delegate.activateSpan(span);
+        }
+
+        @Override
+        public SpanBuilder buildSpan(String s) {
+            return delegate.buildSpan(s);
+        }
+
+        @Override
+        public <C> void inject(SpanContext spanContext, Format<C> format, C c) {
+            delegate.inject(spanContext, format, c);
+        }
+
+        @Override
+        public <C> SpanContext extract(Format<C> format, C c) {
+            return delegate.extract(format, c);
+        }
+
+        @Override
+        public void close() {
+            delegate.close();
+        }
+    }
+}

tracing/jersey/src/main/java/io/helidon/tracing/jersey/AbstractTracingFilter.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2018, 2020 Oracle and/or its affiliates.
+ * Copyright (c) 2018, 2021 Oracle and/or its affiliates.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -89,11 +89,12 @@ public void filter(ContainerRequestContext requestContext) {
             Span span = spanBuilder.start();
             Scope spanScope = tracer.scopeManager().activate(span);
 
-            requestContext.setProperty(SPAN_PROPERTY, span);
-            requestContext.setProperty(SPAN_SCOPE_PROPERTY, spanScope);
+            context.register(span);
+            context.register(spanScope);
             context.register(ClientTracingFilter.class, span.context());
+            requestContext.setProperty(SPAN_PROPERTY, span);
 
-            if (!context.get(TracingContext.class).isPresent()) {
+            if (context.get(TracingContext.class).isEmpty()) {
                 context.register(TracingContext.create(tracer, requestContext.getHeaders()));
             }
 
@@ -165,10 +166,8 @@ public void filter(ContainerRequestContext requestContext, ContainerResponseCont
         requestContext.setProperty(SPAN_FINISHED_PROPERTY, true);
         span.finish();
 
-        Scope spanScope = (Scope) requestContext.getProperty(SPAN_SCOPE_PROPERTY);
-        if (null != spanScope) {
-            spanScope.close();
-        }
+        Context context = Contexts.context().orElseThrow(() -> new IllegalStateException("Context must be available in Jersey"));
+        context.get(Scope.class).ifPresent(Scope::close);
     }
 
     /**