Skip to content
This repository was archived by the owner on Jul 1, 2022. It is now read-only.

Commit 74f4aa2

Browse files
avimasavim
andauthored
Add MDCScopeManager for correlating logs with trace context (#718)
* Support MDC initial Signed-off-by: avim <[email protected]> * add license Signed-off-by: avim <[email protected]> * add license Signed-off-by: avim <[email protected]> * PR feedback add tests Signed-off-by: avim <[email protected]> * fix style Signed-off-by: avim <[email protected]> * PR test feedback Signed-off-by: avim <[email protected]> * test feedback Signed-off-by: avim <[email protected]> * remove empty line Signed-off-by: avim <[email protected]> * add license expand imports Signed-off-by: avim <[email protected]> * allow 4 consecutive capital letters Signed-off-by: avim <[email protected]> * fix style errors Signed-off-by: avim <[email protected]> * revert global stylecheck config escape only MDC classes Signed-off-by: avim <[email protected]> Co-authored-by: avim <[email protected]>
1 parent 760fdda commit 74f4aa2

File tree

3 files changed

+274
-0
lines changed

3 files changed

+274
-0
lines changed

config/checkstyle/suppressions.xml

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,4 +6,8 @@
66
<suppressions>
77
<suppress checks="AbbreviationAsWordInName"
88
files="SpanContext.java"/>
9+
<suppress checks="AbbreviationAsWordInName"
10+
files="MDCScopeManagerTest.java"/>
11+
<suppress checks="AbbreviationAsWordInName"
12+
files="MDCScopeManager.java"/>
913
</suppressions>
Lines changed: 128 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,128 @@
1+
/*
2+
* Copyright (c) 2020, The Jaeger Authors
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
5+
* in compliance with the License. You may obtain a copy of the License at
6+
*
7+
* http://www.apache.org/licenses/LICENSE-2.0
8+
*
9+
* Unless required by applicable law or agreed to in writing, software distributed under the License
10+
* is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
11+
* or implied. See the License for the specific language governing permissions and limitations under
12+
* the License.
13+
*/
14+
15+
package io.jaegertracing.internal;
16+
17+
import io.opentracing.Scope;
18+
import io.opentracing.ScopeManager;
19+
import io.opentracing.Span;
20+
import io.opentracing.util.ThreadLocalScopeManager;
21+
import org.slf4j.MDC;
22+
23+
public class MDCScopeManager implements ScopeManager {
24+
private final ScopeManager wrappedScopeManager;
25+
private final String mdcTraceIdKey;
26+
private final String mdcSpanIdKey;
27+
private final String mdcSampledKey;
28+
29+
private MDCScopeManager(Builder builder) {
30+
this.wrappedScopeManager = builder.scopeManager;
31+
this.mdcTraceIdKey = builder.mdcTraceIdKey;
32+
this.mdcSpanIdKey = builder.mdcSpanIdKey;
33+
this.mdcSampledKey = builder.mdcSampledKey;
34+
}
35+
36+
@Override
37+
public Scope activate(Span span) {
38+
return new MDCScope(wrappedScopeManager.activate(span), span);
39+
}
40+
41+
@Override
42+
public Span activeSpan() {
43+
return wrappedScopeManager.activeSpan();
44+
}
45+
46+
/**
47+
* Builds an {@link MDCScopeManager} with options.
48+
* Calling {@code new MDCScopeManager.Builder().build()}
49+
* Builds an {@link MDCScopeManager} with configuration as follows:
50+
* mdcTraceIDKey set to "traceId"
51+
* mdcSpanIdKey set to "spanId"
52+
* mdcSampledKey set to "sampled"
53+
*/
54+
public static class Builder {
55+
private ScopeManager scopeManager = new ThreadLocalScopeManager();
56+
private String mdcTraceIdKey = "traceId";
57+
private String mdcSpanIdKey = "spanId";
58+
private String mdcSampledKey = "sampled";
59+
60+
public Builder withScopeManager(ScopeManager scopeManager) {
61+
this.scopeManager = scopeManager;
62+
return this;
63+
}
64+
65+
public Builder withMDCTraceIdKey(String mdcTraceIdKey) {
66+
this.mdcTraceIdKey = mdcTraceIdKey;
67+
return this;
68+
}
69+
70+
public Builder withMDCSpanIdKey(String mdcSpanIdKey) {
71+
this.mdcSpanIdKey = mdcSpanIdKey;
72+
return this;
73+
}
74+
75+
public Builder withMDCSampledKey(String mdcSampledKey) {
76+
this.mdcSampledKey = mdcSampledKey;
77+
return this;
78+
}
79+
80+
public MDCScopeManager build() {
81+
return new MDCScopeManager(this);
82+
}
83+
84+
}
85+
86+
private class MDCScope implements Scope {
87+
private final Scope wrappedScope;
88+
private final String previousTraceId;
89+
private final String previousSpanId;
90+
private final String previousSampled;
91+
92+
/**
93+
* mdcScope.
94+
*/
95+
MDCScope(Scope scope, Span span) {
96+
this.wrappedScope = scope;
97+
this.previousTraceId = MDC.get(mdcTraceIdKey);
98+
this.previousSpanId = MDC.get(mdcSpanIdKey);
99+
this.previousSampled = MDC.get(mdcSampledKey);
100+
101+
if (span.context() instanceof JaegerSpanContext) {
102+
putContext((JaegerSpanContext) span.context());
103+
}
104+
}
105+
106+
protected void putContext(JaegerSpanContext spanContext) {
107+
replace(mdcTraceIdKey, spanContext.toTraceId());
108+
replace(mdcSpanIdKey, spanContext.toSpanId());
109+
replace(mdcSampledKey, String.valueOf(spanContext.isSampled()));
110+
}
111+
112+
private void replace(String key, String value) {
113+
if (value == null) {
114+
MDC.remove(key);
115+
} else {
116+
MDC.put(key, value);
117+
}
118+
}
119+
120+
@Override
121+
public void close() {
122+
wrappedScope.close();
123+
replace(mdcTraceIdKey, previousTraceId);
124+
replace(mdcSpanIdKey, previousSpanId);
125+
replace(mdcSampledKey, previousSampled);
126+
}
127+
}
128+
}
Lines changed: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,142 @@
1+
/*
2+
* Copyright (c) 2020, The Jaeger Authors
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
5+
* in compliance with the License. You may obtain a copy of the License at
6+
*
7+
* http://www.apache.org/licenses/LICENSE-2.0
8+
*
9+
* Unless required by applicable law or agreed to in writing, software distributed under the License
10+
* is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
11+
* or implied. See the License for the specific language governing permissions and limitations under
12+
* the License.
13+
*/
14+
15+
package io.jaegertracing.internal;
16+
17+
import static org.junit.Assert.assertEquals;
18+
import static org.junit.Assert.assertNull;
19+
20+
import io.jaegertracing.internal.reporters.InMemoryReporter;
21+
import io.jaegertracing.internal.samplers.ConstSampler;
22+
import io.opentracing.Scope;
23+
import io.opentracing.ScopeManager;
24+
import io.opentracing.Span;
25+
import io.opentracing.Tracer;
26+
import org.junit.Before;
27+
import org.junit.Test;
28+
import org.mockito.Mockito;
29+
import org.slf4j.MDC;
30+
31+
public class MDCScopeManagerTest {
32+
33+
private InMemoryReporter reporter;
34+
private JaegerTracer defaultTracer;
35+
private static final String TRACE_ID = "traceId";
36+
private static final String SPAN_ID = "spanId";
37+
private static final String SAMPLED = "sampled";
38+
39+
40+
@Before
41+
public void setUp() {
42+
reporter = new InMemoryReporter();
43+
defaultTracer = createTracer(new MDCScopeManager.Builder().build());
44+
}
45+
46+
@Test
47+
public void testActiveSpan() {
48+
Span mockSpan = Mockito.mock(JaegerSpan.class);
49+
try (Scope scope = defaultTracer.activateSpan(mockSpan)) {
50+
assertEquals(mockSpan, defaultTracer.activeSpan());
51+
}
52+
}
53+
54+
@Test
55+
public void testNestedSpans() {
56+
Span parentSpan = defaultTracer.buildSpan("parent").start();
57+
try (Scope scope = defaultTracer.activateSpan(parentSpan)) {
58+
assertSpanContextEqualsToMDC((JaegerSpanContext) parentSpan.context(), TRACE_ID, SPAN_ID, SAMPLED);
59+
Span childSpan = defaultTracer.buildSpan("child").start();
60+
try (Scope childScope = defaultTracer.activateSpan(childSpan)) {
61+
assertSpanContextEqualsToMDC((JaegerSpanContext) childSpan.context(), TRACE_ID, SPAN_ID, SAMPLED);
62+
}
63+
assertSpanContextEqualsToMDC((JaegerSpanContext) parentSpan.context(), TRACE_ID, SPAN_ID, SAMPLED);
64+
}
65+
assertNullMDCKeys(TRACE_ID, SPAN_ID, SAMPLED);
66+
}
67+
68+
@Test
69+
public void testDefaultCreation() {
70+
Span span = defaultTracer.buildSpan("test Default").start();
71+
Scope scope = defaultTracer.activateSpan(span);
72+
73+
assertSpanContextEqualsToMDC((JaegerSpanContext) span.context(), TRACE_ID, SPAN_ID, SAMPLED);
74+
75+
scope.close();
76+
assertNullMDCKeys(TRACE_ID, SPAN_ID, SAMPLED);
77+
}
78+
79+
@Test
80+
public void testCustomKeysCreation() {
81+
ScopeManager mdcScopeManager = new MDCScopeManager
82+
.Builder()
83+
.withMDCTraceIdKey("CustomTraceId")
84+
.withMDCSampledKey("customSampled")
85+
.withMDCSpanIdKey("customSpanId")
86+
.build();
87+
88+
Tracer tracer = createTracer(mdcScopeManager);
89+
Span span = tracer.buildSpan("testCustomKeysCreation").start();
90+
Scope scope = tracer.activateSpan(span);
91+
92+
assertSpanContextEqualsToMDC((JaegerSpanContext) span.context(), "CustomTraceId", "customSpanId", "customSampled");
93+
94+
scope.close();
95+
96+
assertNullMDCKeys("CustomTraceId", "customSampled", "customSpanId");
97+
}
98+
99+
@Test
100+
public void testCustomAndDefaultKeysCreation() {
101+
ScopeManager mdcScopeManager = new MDCScopeManager
102+
.Builder()
103+
.withMDCSampledKey("customSampled")
104+
.withMDCSpanIdKey("customSpanId")
105+
.build();
106+
107+
Tracer tracer = createTracer(mdcScopeManager);
108+
Span span = tracer.buildSpan("testCustomAndDefaultKeysCreation").start();
109+
Scope scope = tracer.activateSpan(span);
110+
111+
assertSpanContextEqualsToMDC((JaegerSpanContext) span.context(), TRACE_ID, "customSpanId", "customSampled");
112+
113+
scope.close();
114+
115+
assertNullMDCKeys(TRACE_ID, "customSpanId", "customSampled");
116+
}
117+
118+
private JaegerTracer createTracer(ScopeManager scopeManager) {
119+
return new JaegerTracer.Builder("MDCScopeManagerTest")
120+
.withReporter(reporter)
121+
.withSampler(new ConstSampler(true))
122+
.withScopeManager(scopeManager)
123+
.build();
124+
}
125+
126+
private void assertSpanContextEqualsToMDC(JaegerSpanContext context,
127+
String traceIDKey,
128+
String spanIdKey,
129+
String sampledKey) {
130+
131+
assertEquals(context.toTraceId(), MDC.get(traceIDKey));
132+
assertEquals(context.toSpanId(), MDC.get(spanIdKey));
133+
assertEquals(String.valueOf(context.isSampled()), MDC.get(sampledKey));
134+
}
135+
136+
private void assertNullMDCKeys(String traceIDKey, String spanIdKey, String sampleKey) {
137+
assertNull(MDC.get(traceIDKey));
138+
assertNull(MDC.get(spanIdKey));
139+
assertNull(MDC.get(sampleKey));
140+
}
141+
142+
}

0 commit comments

Comments
 (0)