diff --git a/config/checkstyle/suppressions.xml b/config/checkstyle/suppressions.xml index 5edab5501..5dc4da151 100644 --- a/config/checkstyle/suppressions.xml +++ b/config/checkstyle/suppressions.xml @@ -6,4 +6,8 @@ + + diff --git a/jaeger-core/src/main/java/io/jaegertracing/internal/MDCScopeManager.java b/jaeger-core/src/main/java/io/jaegertracing/internal/MDCScopeManager.java new file mode 100644 index 000000000..b5754e640 --- /dev/null +++ b/jaeger-core/src/main/java/io/jaegertracing/internal/MDCScopeManager.java @@ -0,0 +1,128 @@ +/* + * Copyright (c) 2020, The Jaeger 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 + * + * 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.jaegertracing.internal; + +import io.opentracing.Scope; +import io.opentracing.ScopeManager; +import io.opentracing.Span; +import io.opentracing.util.ThreadLocalScopeManager; +import org.slf4j.MDC; + +public class MDCScopeManager implements ScopeManager { + private final ScopeManager wrappedScopeManager; + private final String mdcTraceIdKey; + private final String mdcSpanIdKey; + private final String mdcSampledKey; + + private MDCScopeManager(Builder builder) { + this.wrappedScopeManager = builder.scopeManager; + this.mdcTraceIdKey = builder.mdcTraceIdKey; + this.mdcSpanIdKey = builder.mdcSpanIdKey; + this.mdcSampledKey = builder.mdcSampledKey; + } + + @Override + public Scope activate(Span span) { + return new MDCScope(wrappedScopeManager.activate(span), span); + } + + @Override + public Span activeSpan() { + return wrappedScopeManager.activeSpan(); + } + + /** + * Builds an {@link MDCScopeManager} with options. + * Calling {@code new MDCScopeManager.Builder().build()} + * Builds an {@link MDCScopeManager} with configuration as follows: + * mdcTraceIDKey set to "traceId" + * mdcSpanIdKey set to "spanId" + * mdcSampledKey set to "sampled" + */ + public static class Builder { + private ScopeManager scopeManager = new ThreadLocalScopeManager(); + private String mdcTraceIdKey = "traceId"; + private String mdcSpanIdKey = "spanId"; + private String mdcSampledKey = "sampled"; + + public Builder withScopeManager(ScopeManager scopeManager) { + this.scopeManager = scopeManager; + return this; + } + + public Builder withMDCTraceIdKey(String mdcTraceIdKey) { + this.mdcTraceIdKey = mdcTraceIdKey; + return this; + } + + public Builder withMDCSpanIdKey(String mdcSpanIdKey) { + this.mdcSpanIdKey = mdcSpanIdKey; + return this; + } + + public Builder withMDCSampledKey(String mdcSampledKey) { + this.mdcSampledKey = mdcSampledKey; + return this; + } + + public MDCScopeManager build() { + return new MDCScopeManager(this); + } + + } + + private class MDCScope implements Scope { + private final Scope wrappedScope; + private final String previousTraceId; + private final String previousSpanId; + private final String previousSampled; + + /** + * mdcScope. + */ + MDCScope(Scope scope, Span span) { + this.wrappedScope = scope; + this.previousTraceId = MDC.get(mdcTraceIdKey); + this.previousSpanId = MDC.get(mdcSpanIdKey); + this.previousSampled = MDC.get(mdcSampledKey); + + if (span.context() instanceof JaegerSpanContext) { + putContext((JaegerSpanContext) span.context()); + } + } + + protected void putContext(JaegerSpanContext spanContext) { + replace(mdcTraceIdKey, spanContext.toTraceId()); + replace(mdcSpanIdKey, spanContext.toSpanId()); + replace(mdcSampledKey, String.valueOf(spanContext.isSampled())); + } + + private void replace(String key, String value) { + if (value == null) { + MDC.remove(key); + } else { + MDC.put(key, value); + } + } + + @Override + public void close() { + wrappedScope.close(); + replace(mdcTraceIdKey, previousTraceId); + replace(mdcSpanIdKey, previousSpanId); + replace(mdcSampledKey, previousSampled); + } + } +} diff --git a/jaeger-core/src/test/java/io/jaegertracing/internal/MDCScopeManagerTest.java b/jaeger-core/src/test/java/io/jaegertracing/internal/MDCScopeManagerTest.java new file mode 100644 index 000000000..da38a20d0 --- /dev/null +++ b/jaeger-core/src/test/java/io/jaegertracing/internal/MDCScopeManagerTest.java @@ -0,0 +1,142 @@ +/* + * Copyright (c) 2020, The Jaeger 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 + * + * 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.jaegertracing.internal; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNull; + +import io.jaegertracing.internal.reporters.InMemoryReporter; +import io.jaegertracing.internal.samplers.ConstSampler; +import io.opentracing.Scope; +import io.opentracing.ScopeManager; +import io.opentracing.Span; +import io.opentracing.Tracer; +import org.junit.Before; +import org.junit.Test; +import org.mockito.Mockito; +import org.slf4j.MDC; + +public class MDCScopeManagerTest { + + private InMemoryReporter reporter; + private JaegerTracer defaultTracer; + private static final String TRACE_ID = "traceId"; + private static final String SPAN_ID = "spanId"; + private static final String SAMPLED = "sampled"; + + + @Before + public void setUp() { + reporter = new InMemoryReporter(); + defaultTracer = createTracer(new MDCScopeManager.Builder().build()); + } + + @Test + public void testActiveSpan() { + Span mockSpan = Mockito.mock(JaegerSpan.class); + try (Scope scope = defaultTracer.activateSpan(mockSpan)) { + assertEquals(mockSpan, defaultTracer.activeSpan()); + } + } + + @Test + public void testNestedSpans() { + Span parentSpan = defaultTracer.buildSpan("parent").start(); + try (Scope scope = defaultTracer.activateSpan(parentSpan)) { + assertSpanContextEqualsToMDC((JaegerSpanContext) parentSpan.context(), TRACE_ID, SPAN_ID, SAMPLED); + Span childSpan = defaultTracer.buildSpan("child").start(); + try (Scope childScope = defaultTracer.activateSpan(childSpan)) { + assertSpanContextEqualsToMDC((JaegerSpanContext) childSpan.context(), TRACE_ID, SPAN_ID, SAMPLED); + } + assertSpanContextEqualsToMDC((JaegerSpanContext) parentSpan.context(), TRACE_ID, SPAN_ID, SAMPLED); + } + assertNullMDCKeys(TRACE_ID, SPAN_ID, SAMPLED); + } + + @Test + public void testDefaultCreation() { + Span span = defaultTracer.buildSpan("test Default").start(); + Scope scope = defaultTracer.activateSpan(span); + + assertSpanContextEqualsToMDC((JaegerSpanContext) span.context(), TRACE_ID, SPAN_ID, SAMPLED); + + scope.close(); + assertNullMDCKeys(TRACE_ID, SPAN_ID, SAMPLED); + } + + @Test + public void testCustomKeysCreation() { + ScopeManager mdcScopeManager = new MDCScopeManager + .Builder() + .withMDCTraceIdKey("CustomTraceId") + .withMDCSampledKey("customSampled") + .withMDCSpanIdKey("customSpanId") + .build(); + + Tracer tracer = createTracer(mdcScopeManager); + Span span = tracer.buildSpan("testCustomKeysCreation").start(); + Scope scope = tracer.activateSpan(span); + + assertSpanContextEqualsToMDC((JaegerSpanContext) span.context(), "CustomTraceId", "customSpanId", "customSampled"); + + scope.close(); + + assertNullMDCKeys("CustomTraceId", "customSampled", "customSpanId"); + } + + @Test + public void testCustomAndDefaultKeysCreation() { + ScopeManager mdcScopeManager = new MDCScopeManager + .Builder() + .withMDCSampledKey("customSampled") + .withMDCSpanIdKey("customSpanId") + .build(); + + Tracer tracer = createTracer(mdcScopeManager); + Span span = tracer.buildSpan("testCustomAndDefaultKeysCreation").start(); + Scope scope = tracer.activateSpan(span); + + assertSpanContextEqualsToMDC((JaegerSpanContext) span.context(), TRACE_ID, "customSpanId", "customSampled"); + + scope.close(); + + assertNullMDCKeys(TRACE_ID, "customSpanId", "customSampled"); + } + + private JaegerTracer createTracer(ScopeManager scopeManager) { + return new JaegerTracer.Builder("MDCScopeManagerTest") + .withReporter(reporter) + .withSampler(new ConstSampler(true)) + .withScopeManager(scopeManager) + .build(); + } + + private void assertSpanContextEqualsToMDC(JaegerSpanContext context, + String traceIDKey, + String spanIdKey, + String sampledKey) { + + assertEquals(context.toTraceId(), MDC.get(traceIDKey)); + assertEquals(context.toSpanId(), MDC.get(spanIdKey)); + assertEquals(String.valueOf(context.isSampled()), MDC.get(sampledKey)); + } + + private void assertNullMDCKeys(String traceIDKey, String spanIdKey, String sampleKey) { + assertNull(MDC.get(traceIDKey)); + assertNull(MDC.get(spanIdKey)); + assertNull(MDC.get(sampleKey)); + } + +} \ No newline at end of file