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.

6 files changed, 285 insertions, 63 deletions

diff --git a/buildscripts/import-control.xml b/buildscripts/import-control.xml
index b89f9f27..61c42f10 100644
--- a/buildscripts/import-control.xml
+++ b/buildscripts/import-control.xml
@@ -95,6 +95,7 @@ General guidelines on imports:
       <allow pkg="io.opencensus.trace"/>
     </subpackage>
     <subpackage name="logcorrelation.log4j2">
+      <allow pkg="io.opencensus.contrib.logcorrelation.log4j2"/>
       <allow pkg="io.opencensus.trace"/>
       <disallow pkg="org.apache.logging.log4j.core.impl"/>
       <allow pkg="org.apache.logging.log4j"/>
diff --git a/checker-framework/stubs/log4j.astub b/checker-framework/stubs/log4j.astub
new file mode 100644
index 00000000..20b3240e
--- /dev/null
+++ b/checker-framework/stubs/log4j.astub
@@ -0,0 +1,8 @@
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+package org.apache.logging.log4j;
+
+class ThreadContext {
+  @Nullable
+  static ReadOnlyThreadContextMap getThreadContextMap();
+}
diff --git a/contrib/log_correlation/log4j2/src/main/java/io/opencensus/contrib/logcorrelation/log4j2/ContextDataUtils.java b/contrib/log_correlation/log4j2/src/main/java/io/opencensus/contrib/logcorrelation/log4j2/ContextDataUtils.java
new file mode 100644
index 00000000..82663e07
--- /dev/null
+++ b/contrib/log_correlation/log4j2/src/main/java/io/opencensus/contrib/logcorrelation/log4j2/ContextDataUtils.java
@@ -0,0 +1,236 @@
+/*
+ * Copyright 2018, OpenCensus 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.opencensus.contrib.logcorrelation.log4j2;
+
+import io.opencensus.contrib.logcorrelation.log4j2.OpenCensusTraceContextDataInjector.SpanSelection;
+import io.opencensus.trace.Span;
+import io.opencensus.trace.SpanContext;
+import io.opencensus.trace.SpanId;
+import io.opencensus.trace.TraceId;
+import io.opencensus.trace.unsafe.ContextUtils;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
+import javax.annotation.Nullable;
+import javax.annotation.concurrent.Immutable;
+import org.apache.logging.log4j.ThreadContext;
+import org.apache.logging.log4j.core.config.Property;
+import org.apache.logging.log4j.spi.ReadOnlyThreadContextMap;
+import org.apache.logging.log4j.util.BiConsumer;
+import org.apache.logging.log4j.util.ReadOnlyStringMap;
+import org.apache.logging.log4j.util.SortedArrayStringMap;
+import org.apache.logging.log4j.util.StringMap;
+import org.apache.logging.log4j.util.TriConsumer;
+
+// Implementation of the methods inherited from ContextDataInjector.
+//
+// This class uses "shareable" to mean that a method's return value can be passed to another
+// thread.
+final class ContextDataUtils {
+  private ContextDataUtils() {}
+
+  // The implementation of this method is based on the example in the Javadocs for
+  // ContextDataInjector.injectContextData.
+  static StringMap injectContextData(
+      SpanSelection spanSelection, @Nullable List<Property> properties, StringMap reusable) {
+    if (properties == null || properties.isEmpty()) {
+      return shareableRawContextData(spanSelection);
+    }
+    // Context data has precedence over configuration properties.
+    putProperties(properties, reusable);
+    // TODO(sebright): The following line can be optimized. See
+    //     https://github.com/census-instrumentation/opencensus-java/pull/1422/files#r216425494.
+    reusable.putAll(nonShareableRawContextData(spanSelection));
+    return reusable;
+  }
+
+  private static void putProperties(Collection<Property> properties, StringMap stringMap) {
+    for (Property property : properties) {
+      stringMap.putValue(property.getName(), property.getValue());
+    }
+  }
+
+  private static StringMap shareableRawContextData(SpanSelection spanSelection) {
+    SpanContext spanContext = shouldAddTracingDataToLogEvent(spanSelection);
+    return spanContext == null
+        ? getShareableContextData()
+        : getShareableContextAndTracingData(spanContext);
+  }
+
+  static ReadOnlyStringMap nonShareableRawContextData(SpanSelection spanSelection) {
+    SpanContext spanContext = shouldAddTracingDataToLogEvent(spanSelection);
+    return spanContext == null
+        ? getNonShareableContextData()
+        : getShareableContextAndTracingData(spanContext);
+  }
+
+  // This method returns the current span context iff tracing data should be added to the LogEvent.
+  // It avoids getting the current span when the feature is disabled, for efficiency.
+  @Nullable
+  private static SpanContext shouldAddTracingDataToLogEvent(SpanSelection spanSelection) {
+    switch (spanSelection) {
+      case NO_SPANS:
+        return null;
+      case SAMPLED_SPANS:
+        SpanContext spanContext = getCurrentSpanContext();
+        if (spanContext.getTraceOptions().isSampled()) {
+          return spanContext;
+        } else {
+          return null;
+        }
+      case ALL_SPANS:
+        return getCurrentSpanContext();
+    }
+    throw new AssertionError("Unknown spanSelection: " + spanSelection);
+  }
+
+  private static StringMap getShareableContextData() {
+    ReadOnlyThreadContextMap context = ThreadContext.getThreadContextMap();
+
+    // Return a new object, since StringMap is modifiable.
+    return context == null
+        ? new SortedArrayStringMap(ThreadContext.getImmutableContext())
+        : new SortedArrayStringMap(context.getReadOnlyContextData());
+  }
+
+  private static ReadOnlyStringMap getNonShareableContextData() {
+    ReadOnlyThreadContextMap context = ThreadContext.getThreadContextMap();
+    if (context != null) {
+      return context.getReadOnlyContextData();
+    } else {
+      Map<String, String> contextMap = ThreadContext.getImmutableContext();
+      return contextMap.isEmpty()
+          ? UnmodifiableReadOnlyStringMap.EMPTY
+          : new UnmodifiableReadOnlyStringMap(contextMap);
+    }
+  }
+
+  private static StringMap getShareableContextAndTracingData(SpanContext spanContext) {
+    ReadOnlyThreadContextMap context = ThreadContext.getThreadContextMap();
+    SortedArrayStringMap stringMap;
+    if (context == null) {
+      stringMap = new SortedArrayStringMap(ThreadContext.getImmutableContext());
+    } else {
+      StringMap contextData = context.getReadOnlyContextData();
+      stringMap = new SortedArrayStringMap(contextData.size() + 3);
+      stringMap.putAll(contextData);
+    }
+    stringMap.putValue(
+        OpenCensusTraceContextDataInjector.TRACE_ID_CONTEXT_KEY,
+        new TraceIdToLowerBase16Formatter(spanContext.getTraceId()));
+    stringMap.putValue(
+        OpenCensusTraceContextDataInjector.SPAN_ID_CONTEXT_KEY,
+        new SpanIdToLowerBase16Formatter(spanContext.getSpanId()));
+    stringMap.putValue(
+        OpenCensusTraceContextDataInjector.TRACE_SAMPLED_CONTEXT_KEY,
+        spanContext.getTraceOptions().isSampled() ? "true" : "false");
+    return stringMap;
+  }
+
+  private static SpanContext getCurrentSpanContext() {
+    Span span = ContextUtils.CONTEXT_SPAN_KEY.get();
+    return span == null ? SpanContext.INVALID : span.getContext();
+  }
+
+  private static final class TraceIdToLowerBase16Formatter {
+    private final TraceId traceId;
+
+    private TraceIdToLowerBase16Formatter(TraceId traceId) {
+      this.traceId = traceId;
+    }
+
+    @Override
+    public String toString() {
+      return traceId.toLowerBase16();
+    }
+  }
+
+  private static final class SpanIdToLowerBase16Formatter {
+    private final SpanId spanId;
+
+    private SpanIdToLowerBase16Formatter(SpanId spanId) {
+      this.spanId = spanId;
+    }
+
+    @Override
+    public String toString() {
+      return spanId.toLowerBase16();
+    }
+  }
+
+  @Immutable
+  private static final class UnmodifiableReadOnlyStringMap implements ReadOnlyStringMap {
+    private static final long serialVersionUID = 0L;
+
+    static final ReadOnlyStringMap EMPTY =
+        new UnmodifiableReadOnlyStringMap(Collections.<String, String>emptyMap());
+
+    private final Map<String, String> map;
+
+    UnmodifiableReadOnlyStringMap(Map<String, String> map) {
+      this.map = map;
+    }
+
+    @Override
+    public boolean containsKey(String key) {
+      return map.containsKey(key);
+    }
+
+    @Override
+    @SuppressWarnings("unchecked")
+    public <V> void forEach(BiConsumer<String, ? super V> action) {
+      for (Entry<String, String> entry : map.entrySet()) {
+        action.accept(entry.getKey(), (V) entry.getValue());
+      }
+    }
+
+    @Override
+    @SuppressWarnings("unchecked")
+    public <V, S> void forEach(TriConsumer<String, ? super V, S> action, S state) {
+      for (Entry<String, String> entry : map.entrySet()) {
+        action.accept(entry.getKey(), (V) entry.getValue(), state);
+      }
+    }
+
+    @Override
+    @Nullable
+    @SuppressWarnings({
+      "unchecked",
+      "TypeParameterUnusedInFormals" // This is an overridden method.
+    })
+    public <V> V getValue(String key) {
+      return (V) map.get(key);
+    }
+
+    @Override
+    public boolean isEmpty() {
+      return map.isEmpty();
+    }
+
+    @Override
+    public int size() {
+      return map.size();
+    }
+
+    @Override
+    public Map<String, String> toMap() {
+      return map;
+    }
+  }
+}
diff --git a/contrib/log_correlation/log4j2/src/main/java/io/opencensus/contrib/logcorrelation/log4j2/OpenCensusTraceContextDataInjector.java b/contrib/log_correlation/log4j2/src/main/java/io/opencensus/contrib/logcorrelation/log4j2/OpenCensusTraceContextDataInjector.java
index 0993a8e4..c00a11a2 100644
--- a/contrib/log_correlation/log4j2/src/main/java/io/opencensus/contrib/logcorrelation/log4j2/OpenCensusTraceContextDataInjector.java
+++ b/contrib/log_correlation/log4j2/src/main/java/io/opencensus/contrib/logcorrelation/log4j2/OpenCensusTraceContextDataInjector.java
@@ -17,19 +17,13 @@
 package io.opencensus.contrib.logcorrelation.log4j2;
 
 import io.opencensus.common.ExperimentalApi;
-import io.opencensus.trace.Span;
-import io.opencensus.trace.SpanContext;
-import io.opencensus.trace.unsafe.ContextUtils;
-import java.util.Collection;
 import java.util.List;
-import java.util.Map;
 import javax.annotation.Nullable;
-import org.apache.logging.log4j.ThreadContext;
 import org.apache.logging.log4j.core.ContextDataInjector;
 import org.apache.logging.log4j.core.Layout;
 import org.apache.logging.log4j.core.LogEvent;
 import org.apache.logging.log4j.core.config.Property;
-import org.apache.logging.log4j.util.SortedArrayStringMap;
+import org.apache.logging.log4j.util.ReadOnlyStringMap;
 import org.apache.logging.log4j.util.StringMap;
 
 /**
@@ -168,61 +162,15 @@ public final class OpenCensusTraceContextDataInjector implements ContextDataInje
     }
   }
 
-  // The implementation of this method is based on the example in the Javadocs for
-  // ContextDataInjector.injectContextData.
+  // Note that this method must return an object that can be passed to another thread.
   @Override
   public StringMap injectContextData(@Nullable List<Property> properties, StringMap reusable) {
-    if (properties == null || properties.isEmpty()) {
-      return rawContextData();
-    }
-    // Context data has precedence over configuration properties.
-    putProperties(properties, reusable);
-    reusable.putAll(rawContextData());
-    return reusable;
-  }
-
-  private static void putProperties(Collection<Property> properties, StringMap stringMap) {
-    for (Property property : properties) {
-      stringMap.putValue(property.getName(), property.getValue());
-    }
+    return ContextDataUtils.injectContextData(spanSelection, properties, reusable);
   }
 
-  // This method avoids getting the current span when the feature is disabled, for efficiency.
+  // Note that this method does not need to return an object that can be passed to another thread.
   @Override
-  public StringMap rawContextData() {
-    switch (spanSelection) {
-      case NO_SPANS:
-        return getContextData();
-      case SAMPLED_SPANS:
-        SpanContext spanContext = getCurrentSpanContext();
-        if (spanContext.getTraceOptions().isSampled()) {
-          return getContextAndTracingData(spanContext);
-        } else {
-          return getContextData();
-        }
-      case ALL_SPANS:
-        return getContextAndTracingData(getCurrentSpanContext());
-    }
-    throw new AssertionError("Unknown spanSelection: " + spanSelection);
-  }
-
-  private static SpanContext getCurrentSpanContext() {
-    Span span = ContextUtils.CONTEXT_SPAN_KEY.get();
-    return span == null ? SpanContext.INVALID : span.getContext();
-  }
-
-  // TODO(sebright): Improve the implementation of this method, including handling null.
-  private static StringMap getContextData() {
-    return ThreadContext.getThreadContextMap().getReadOnlyContextData();
-  }
-
-  // TODO(sebright): Improve the implementation of this method, including handling null.
-  private static StringMap getContextAndTracingData(SpanContext spanContext) {
-    Map<String, String> map = ThreadContext.getThreadContextMap().getCopy();
-    map.put(TRACE_ID_CONTEXT_KEY, spanContext.getTraceId().toLowerBase16());
-    map.put(SPAN_ID_CONTEXT_KEY, spanContext.getSpanId().toLowerBase16());
-    map.put(
-        TRACE_SAMPLED_CONTEXT_KEY, spanContext.getTraceOptions().isSampled() ? "true" : "false");
-    return new SortedArrayStringMap(map);
+  public ReadOnlyStringMap rawContextData() {
+    return ContextDataUtils.nonShareableRawContextData(spanSelection);
   }
 }
diff --git a/contrib/log_correlation/log4j2/src/test/java/io/opencensus/contrib/logcorrelation/log4j2/OpenCensusLog4jLogCorrelationAllSpansTest.java b/contrib/log_correlation/log4j2/src/test/java/io/opencensus/contrib/logcorrelation/log4j2/OpenCensusLog4jLogCorrelationAllSpansTest.java
index 2790b1e5..355c9b67 100644
--- a/contrib/log_correlation/log4j2/src/test/java/io/opencensus/contrib/logcorrelation/log4j2/OpenCensusLog4jLogCorrelationAllSpansTest.java
+++ b/contrib/log_correlation/log4j2/src/test/java/io/opencensus/contrib/logcorrelation/log4j2/OpenCensusLog4jLogCorrelationAllSpansTest.java
@@ -114,7 +114,7 @@ public final class OpenCensusLog4jLogCorrelationAllSpansTest
   public void preserveOtherKeyValuePairs() {
     String log =
         logWithSpanAndLog4jConfiguration(
-            "%X{myTestKey} %-5level - %msg",
+            "%X{opencensusTraceId} %X{myTestKey} %-5level - %msg",
             SpanContext.create(
                 TraceId.fromLowerBase16("c95329bb6b7de41afbc51a231c128f97"),
                 SpanId.fromLowerBase16("bf22ea74d38eddad"),
@@ -133,6 +133,35 @@ public final class OpenCensusLog4jLogCorrelationAllSpansTest
                 return null;
               }
             });
-    assertThat(log).isEqualTo("myTestValue ERROR - message #4");
+    assertThat(log).isEqualTo("c95329bb6b7de41afbc51a231c128f97 myTestValue ERROR - message #4");
+  }
+
+  @Test
+  public void overwriteExistingTracingKey() {
+    String log =
+        logWithSpanAndLog4jConfiguration(
+            TEST_PATTERN,
+            SpanContext.create(
+                TraceId.fromLowerBase16("18e4ae44273a0c44e0c9ea4380792c66"),
+                SpanId.fromLowerBase16("199a7e16daa000a7"),
+                TraceOptions.builder().setIsSampled(true).build(),
+                EMPTY_TRACESTATE),
+            new Function<Logger, Void>() {
+              @Override
+              public Void apply(Logger logger) {
+                ThreadContext.put(
+                    OpenCensusTraceContextDataInjector.TRACE_ID_CONTEXT_KEY, "existingTraceId");
+                try {
+                  logger.error("message #5");
+                } finally {
+                  ThreadContext.remove(OpenCensusTraceContextDataInjector.TRACE_ID_CONTEXT_KEY);
+                }
+                return null;
+              }
+            });
+    assertThat(log)
+        .isEqualTo(
+            "traceId=18e4ae44273a0c44e0c9ea4380792c66 spanId=199a7e16daa000a7 "
+                + "sampled=true ERROR - message #5");
   }
 }
diff --git a/contrib/log_correlation/log4j2/src/test/java/io/opencensus/contrib/logcorrelation/log4j2/OpenCensusTraceContextDataInjectorTest.java b/contrib/log_correlation/log4j2/src/test/java/io/opencensus/contrib/logcorrelation/log4j2/OpenCensusTraceContextDataInjectorTest.java
index 800fdf2e..e027dae3 100644
--- a/contrib/log_correlation/log4j2/src/test/java/io/opencensus/contrib/logcorrelation/log4j2/OpenCensusTraceContextDataInjectorTest.java
+++ b/contrib/log_correlation/log4j2/src/test/java/io/opencensus/contrib/logcorrelation/log4j2/OpenCensusTraceContextDataInjectorTest.java
@@ -142,10 +142,10 @@ public final class OpenCensusTraceContextDataInjectorTest {
   private static Map<String, String> toMap(StringMap stringMap) {
     final ImmutableMap.Builder<String, String> builder = ImmutableMap.builder();
     stringMap.forEach(
-        new BiConsumer<String, String>() {
+        new BiConsumer<String, Object>() {
           @Override
-          public void accept(String key, String value) {
-            builder.put(key, value);
+          public void accept(String key, Object value) {
+            builder.put(key, String.valueOf(value));
           }
         });
     return builder.build();