Skip to content
This repository was archived by the owner on Dec 23, 2023. It is now read-only.

Commit 97a1875

Browse files
sebright2sebright2
authored
Improve implementation of OpenCensusTraceContextDataInjector. (#1413)
This commit improves the thread safety of OpenCensusTraceContextDataInjector by following the thread safety requirements in the Javadocs of the overridden methods from ContextDataInjector. It also handles the possibility of ThreadContext.getThreadContextMap() returning null by adding a Nullable annotation in a Checker Framework stub file and adding null checks to the code.
1 parent 4ddf824 commit 97a1875

File tree

6 files changed

+285
-63
lines changed

6 files changed

+285
-63
lines changed

‎buildscripts/import-control.xml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,7 @@ General guidelines on imports:
9595
<allow pkg="io.opencensus.trace"/>
9696
</subpackage>
9797
<subpackage name="logcorrelation.log4j2">
98+
<allow pkg="io.opencensus.contrib.logcorrelation.log4j2"/>
9899
<allow pkg="io.opencensus.trace"/>
99100
<disallow pkg="org.apache.logging.log4j.core.impl"/>
100101
<allow pkg="org.apache.logging.log4j"/>

‎checker-framework/stubs/log4j.astub‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
import org.checkerframework.checker.nullness.qual.Nullable;
2+
3+
package org.apache.logging.log4j;
4+
5+
class ThreadContext {
6+
@Nullable
7+
static ReadOnlyThreadContextMap getThreadContextMap();
8+
}

‎contrib/log_correlation/log4j2/src/main/java/io/opencensus/contrib/logcorrelation/log4j2/ContextDataUtils.java‎

Lines changed: 236 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,236 @@
1+
/*
2+
* Copyright 2018, OpenCensus Authors
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package io.opencensus.contrib.logcorrelation.log4j2;
18+
19+
import io.opencensus.contrib.logcorrelation.log4j2.OpenCensusTraceContextDataInjector.SpanSelection;
20+
import io.opencensus.trace.Span;
21+
import io.opencensus.trace.SpanContext;
22+
import io.opencensus.trace.SpanId;
23+
import io.opencensus.trace.TraceId;
24+
import io.opencensus.trace.unsafe.ContextUtils;
25+
import java.util.Collection;
26+
import java.util.Collections;
27+
import java.util.List;
28+
import java.util.Map;
29+
import java.util.Map.Entry;
30+
import javax.annotation.Nullable;
31+
import javax.annotation.concurrent.Immutable;
32+
import org.apache.logging.log4j.ThreadContext;
33+
import org.apache.logging.log4j.core.config.Property;
34+
import org.apache.logging.log4j.spi.ReadOnlyThreadContextMap;
35+
import org.apache.logging.log4j.util.BiConsumer;
36+
import org.apache.logging.log4j.util.ReadOnlyStringMap;
37+
import org.apache.logging.log4j.util.SortedArrayStringMap;
38+
import org.apache.logging.log4j.util.StringMap;
39+
import org.apache.logging.log4j.util.TriConsumer;
40+
41+
// Implementation of the methods inherited from ContextDataInjector.
42+
//
43+
// This class uses "shareable" to mean that a method's return value can be passed to another
44+
// thread.
45+
final class ContextDataUtils {
46+
private ContextDataUtils() {}
47+
48+
// The implementation of this method is based on the example in the Javadocs for
49+
// ContextDataInjector.injectContextData.
50+
static StringMap injectContextData(
51+
SpanSelection spanSelection, @Nullable List<Property> properties, StringMap reusable) {
52+
if (properties == null || properties.isEmpty()) {
53+
return shareableRawContextData(spanSelection);
54+
}
55+
// Context data has precedence over configuration properties.
56+
putProperties(properties, reusable);
57+
// TODO(sebright): The following line can be optimized. See
58+
// https://github.com/census-instrumentation/opencensus-java/pull/1422/files#r216425494.
59+
reusable.putAll(nonShareableRawContextData(spanSelection));
60+
return reusable;
61+
}
62+
63+
private static void putProperties(Collection<Property> properties, StringMap stringMap) {
64+
for (Property property : properties) {
65+
stringMap.putValue(property.getName(), property.getValue());
66+
}
67+
}
68+
69+
private static StringMap shareableRawContextData(SpanSelection spanSelection) {
70+
SpanContext spanContext = shouldAddTracingDataToLogEvent(spanSelection);
71+
return spanContext == null
72+
? getShareableContextData()
73+
: getShareableContextAndTracingData(spanContext);
74+
}
75+
76+
static ReadOnlyStringMap nonShareableRawContextData(SpanSelection spanSelection) {
77+
SpanContext spanContext = shouldAddTracingDataToLogEvent(spanSelection);
78+
return spanContext == null
79+
? getNonShareableContextData()
80+
: getShareableContextAndTracingData(spanContext);
81+
}
82+
83+
// This method returns the current span context iff tracing data should be added to the LogEvent.
84+
// It avoids getting the current span when the feature is disabled, for efficiency.
85+
@Nullable
86+
private static SpanContext shouldAddTracingDataToLogEvent(SpanSelection spanSelection) {
87+
switch (spanSelection) {
88+
case NO_SPANS:
89+
return null;
90+
case SAMPLED_SPANS:
91+
SpanContext spanContext = getCurrentSpanContext();
92+
if (spanContext.getTraceOptions().isSampled()) {
93+
return spanContext;
94+
} else {
95+
return null;
96+
}
97+
case ALL_SPANS:
98+
return getCurrentSpanContext();
99+
}
100+
throw new AssertionError("Unknown spanSelection: " + spanSelection);
101+
}
102+
103+
private static StringMap getShareableContextData() {
104+
ReadOnlyThreadContextMap context = ThreadContext.getThreadContextMap();
105+
106+
// Return a new object, since StringMap is modifiable.
107+
return context == null
108+
? new SortedArrayStringMap(ThreadContext.getImmutableContext())
109+
: new SortedArrayStringMap(context.getReadOnlyContextData());
110+
}
111+
112+
private static ReadOnlyStringMap getNonShareableContextData() {
113+
ReadOnlyThreadContextMap context = ThreadContext.getThreadContextMap();
114+
if (context != null) {
115+
return context.getReadOnlyContextData();
116+
} else {
117+
Map<String, String> contextMap = ThreadContext.getImmutableContext();
118+
return contextMap.isEmpty()
119+
? UnmodifiableReadOnlyStringMap.EMPTY
120+
: new UnmodifiableReadOnlyStringMap(contextMap);
121+
}
122+
}
123+
124+
private static StringMap getShareableContextAndTracingData(SpanContext spanContext) {
125+
ReadOnlyThreadContextMap context = ThreadContext.getThreadContextMap();
126+
SortedArrayStringMap stringMap;
127+
if (context == null) {
128+
stringMap = new SortedArrayStringMap(ThreadContext.getImmutableContext());
129+
} else {
130+
StringMap contextData = context.getReadOnlyContextData();
131+
stringMap = new SortedArrayStringMap(contextData.size() + 3);
132+
stringMap.putAll(contextData);
133+
}
134+
stringMap.putValue(
135+
OpenCensusTraceContextDataInjector.TRACE_ID_CONTEXT_KEY,
136+
new TraceIdToLowerBase16Formatter(spanContext.getTraceId()));
137+
stringMap.putValue(
138+
OpenCensusTraceContextDataInjector.SPAN_ID_CONTEXT_KEY,
139+
new SpanIdToLowerBase16Formatter(spanContext.getSpanId()));
140+
stringMap.putValue(
141+
OpenCensusTraceContextDataInjector.TRACE_SAMPLED_CONTEXT_KEY,
142+
spanContext.getTraceOptions().isSampled() ? "true" : "false");
143+
return stringMap;
144+
}
145+
146+
private static SpanContext getCurrentSpanContext() {
147+
Span span = ContextUtils.CONTEXT_SPAN_KEY.get();
148+
return span == null ? SpanContext.INVALID : span.getContext();
149+
}
150+
151+
private static final class TraceIdToLowerBase16Formatter {
152+
private final TraceId traceId;
153+
154+
private TraceIdToLowerBase16Formatter(TraceId traceId) {
155+
this.traceId = traceId;
156+
}
157+
158+
@Override
159+
public String toString() {
160+
return traceId.toLowerBase16();
161+
}
162+
}
163+
164+
private static final class SpanIdToLowerBase16Formatter {
165+
private final SpanId spanId;
166+
167+
private SpanIdToLowerBase16Formatter(SpanId spanId) {
168+
this.spanId = spanId;
169+
}
170+
171+
@Override
172+
public String toString() {
173+
return spanId.toLowerBase16();
174+
}
175+
}
176+
177+
@Immutable
178+
private static final class UnmodifiableReadOnlyStringMap implements ReadOnlyStringMap {
179+
private static final long serialVersionUID = 0L;
180+
181+
static final ReadOnlyStringMap EMPTY =
182+
new UnmodifiableReadOnlyStringMap(Collections.<String, String>emptyMap());
183+
184+
private final Map<String, String> map;
185+
186+
UnmodifiableReadOnlyStringMap(Map<String, String> map) {
187+
this.map = map;
188+
}
189+
190+
@Override
191+
public boolean containsKey(String key) {
192+
return map.containsKey(key);
193+
}
194+
195+
@Override
196+
@SuppressWarnings("unchecked")
197+
public <V> void forEach(BiConsumer<String, ? super V> action) {
198+
for (Entry<String, String> entry : map.entrySet()) {
199+
action.accept(entry.getKey(), (V) entry.getValue());
200+
}
201+
}
202+
203+
@Override
204+
@SuppressWarnings("unchecked")
205+
public <V, S> void forEach(TriConsumer<String, ? super V, S> action, S state) {
206+
for (Entry<String, String> entry : map.entrySet()) {
207+
action.accept(entry.getKey(), (V) entry.getValue(), state);
208+
}
209+
}
210+
211+
@Override
212+
@Nullable
213+
@SuppressWarnings({
214+
"unchecked",
215+
"TypeParameterUnusedInFormals" // This is an overridden method.
216+
})
217+
public <V> V getValue(String key) {
218+
return (V) map.get(key);
219+
}
220+
221+
@Override
222+
public boolean isEmpty() {
223+
return map.isEmpty();
224+
}
225+
226+
@Override
227+
public int size() {
228+
return map.size();
229+
}
230+
231+
@Override
232+
public Map<String, String> toMap() {
233+
return map;
234+
}
235+
}
236+
}

‎contrib/log_correlation/log4j2/src/main/java/io/opencensus/contrib/logcorrelation/log4j2/OpenCensusTraceContextDataInjector.java‎

Lines changed: 6 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -17,19 +17,13 @@
1717
package io.opencensus.contrib.logcorrelation.log4j2;
1818

1919
import io.opencensus.common.ExperimentalApi;
20-
import io.opencensus.trace.Span;
21-
import io.opencensus.trace.SpanContext;
22-
import io.opencensus.trace.unsafe.ContextUtils;
23-
import java.util.Collection;
2420
import java.util.List;
25-
import java.util.Map;
2621
import javax.annotation.Nullable;
27-
import org.apache.logging.log4j.ThreadContext;
2822
import org.apache.logging.log4j.core.ContextDataInjector;
2923
import org.apache.logging.log4j.core.Layout;
3024
import org.apache.logging.log4j.core.LogEvent;
3125
import org.apache.logging.log4j.core.config.Property;
32-
import org.apache.logging.log4j.util.SortedArrayStringMap;
26+
import org.apache.logging.log4j.util.ReadOnlyStringMap;
3327
import org.apache.logging.log4j.util.StringMap;
3428

3529
/**
@@ -168,61 +162,15 @@ private static SpanSelection parseSpanSelection(String spanSelection) {
168162
}
169163
}
170164

171-
// The implementation of this method is based on the example in the Javadocs for
172-
// ContextDataInjector.injectContextData.
165+
// Note that this method must return an object that can be passed to another thread.
173166
@Override
174167
public StringMap injectContextData(@Nullable List<Property> properties, StringMap reusable) {
175-
if (properties == null || properties.isEmpty()) {
176-
return rawContextData();
177-
}
178-
// Context data has precedence over configuration properties.
179-
putProperties(properties, reusable);
180-
reusable.putAll(rawContextData());
181-
return reusable;
182-
}
183-
184-
private static void putProperties(Collection<Property> properties, StringMap stringMap) {
185-
for (Property property : properties) {
186-
stringMap.putValue(property.getName(), property.getValue());
187-
}
168+
return ContextDataUtils.injectContextData(spanSelection, properties, reusable);
188169
}
189170

190-
// This method avoids getting the current span when the feature is disabled, for efficiency.
171+
// Note that this method does not need to return an object that can be passed to another thread.
191172
@Override
192-
public StringMap rawContextData() {
193-
switch (spanSelection) {
194-
case NO_SPANS:
195-
return getContextData();
196-
case SAMPLED_SPANS:
197-
SpanContext spanContext = getCurrentSpanContext();
198-
if (spanContext.getTraceOptions().isSampled()) {
199-
return getContextAndTracingData(spanContext);
200-
} else {
201-
return getContextData();
202-
}
203-
case ALL_SPANS:
204-
return getContextAndTracingData(getCurrentSpanContext());
205-
}
206-
throw new AssertionError("Unknown spanSelection: " + spanSelection);
207-
}
208-
209-
private static SpanContext getCurrentSpanContext() {
210-
Span span = ContextUtils.CONTEXT_SPAN_KEY.get();
211-
return span == null ? SpanContext.INVALID : span.getContext();
212-
}
213-
214-
// TODO(sebright): Improve the implementation of this method, including handling null.
215-
private static StringMap getContextData() {
216-
return ThreadContext.getThreadContextMap().getReadOnlyContextData();
217-
}
218-
219-
// TODO(sebright): Improve the implementation of this method, including handling null.
220-
private static StringMap getContextAndTracingData(SpanContext spanContext) {
221-
Map<String, String> map = ThreadContext.getThreadContextMap().getCopy();
222-
map.put(TRACE_ID_CONTEXT_KEY, spanContext.getTraceId().toLowerBase16());
223-
map.put(SPAN_ID_CONTEXT_KEY, spanContext.getSpanId().toLowerBase16());
224-
map.put(
225-
TRACE_SAMPLED_CONTEXT_KEY, spanContext.getTraceOptions().isSampled() ? "true" : "false");
226-
return new SortedArrayStringMap(map);
173+
public ReadOnlyStringMap rawContextData() {
174+
return ContextDataUtils.nonShareableRawContextData(spanSelection);
227175
}
228176
}

‎contrib/log_correlation/log4j2/src/test/java/io/opencensus/contrib/logcorrelation/log4j2/OpenCensusLog4jLogCorrelationAllSpansTest.java‎

Lines changed: 31 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -114,7 +114,7 @@ public Void apply(Logger logger) {
114114
public void preserveOtherKeyValuePairs() {
115115
String log =
116116
logWithSpanAndLog4jConfiguration(
117-
"%X{myTestKey} %-5level - %msg",
117+
"%X{opencensusTraceId} %X{myTestKey} %-5level - %msg",
118118
SpanContext.create(
119119
TraceId.fromLowerBase16("c95329bb6b7de41afbc51a231c128f97"),
120120
SpanId.fromLowerBase16("bf22ea74d38eddad"),
@@ -133,6 +133,35 @@ public Void apply(Logger logger) {
133133
return null;
134134
}
135135
});
136-
assertThat(log).isEqualTo("myTestValue ERROR - message #4");
136+
assertThat(log).isEqualTo("c95329bb6b7de41afbc51a231c128f97 myTestValue ERROR - message #4");
137+
}
138+
139+
@Test
140+
public void overwriteExistingTracingKey() {
141+
String log =
142+
logWithSpanAndLog4jConfiguration(
143+
TEST_PATTERN,
144+
SpanContext.create(
145+
TraceId.fromLowerBase16("18e4ae44273a0c44e0c9ea4380792c66"),
146+
SpanId.fromLowerBase16("199a7e16daa000a7"),
147+
TraceOptions.builder().setIsSampled(true).build(),
148+
EMPTY_TRACESTATE),
149+
new Function<Logger, Void>() {
150+
@Override
151+
public Void apply(Logger logger) {
152+
ThreadContext.put(
153+
OpenCensusTraceContextDataInjector.TRACE_ID_CONTEXT_KEY, "existingTraceId");
154+
try {
155+
logger.error("message #5");
156+
} finally {
157+
ThreadContext.remove(OpenCensusTraceContextDataInjector.TRACE_ID_CONTEXT_KEY);
158+
}
159+
return null;
160+
}
161+
});
162+
assertThat(log)
163+
.isEqualTo(
164+
"traceId=18e4ae44273a0c44e0c9ea4380792c66 spanId=199a7e16daa000a7 "
165+
+ "sampled=true ERROR - message #5");
137166
}
138167
}

0 commit comments

Comments
 (0)