Disallowing null values in ContextSnapshot

context-propagation/src/main/java/io/micrometer/context/ContextAccessor.java

@@ -42,6 +42,10 @@ public interface ContextAccessor<READ, WRITE> {
      * @param sourceContext the context to read from; the context type should be
      * {@link Class#isAssignableFrom(Class) assignable} from the type returned by
      * {@link #readableType()}.
+     * <p>
+     * When an {@link ContextAccessor} is used to populate a {@link ContextSnapshot}, the
+     * snapshot implementations are required to filter out {@code null} mappings, so it is
+     * not required to implement special handling in the accessor.
      * @param keyPredicate a predicate to decide which keys to read
      * @param readValues a map where to put read values
      */
@@ -53,7 +57,7 @@ public interface ContextAccessor<READ, WRITE> {
      * {@link Class#isAssignableFrom(Class) assignable} from the type returned by
      * {@link #readableType()}.
      * @param key the key to use to look up the context value
-     * @return the value, if any
+     * @return the value, if present
      */
     @Nullable
     <T> T readValue(READ sourceContext, Object key);
@@ -65,7 +69,7 @@ public interface ContextAccessor<READ, WRITE> {
 
     /**
      * Write values from a {@link Map} to a target context.
-     * @param valuesToWrite the values to write to the target context
+     * @param valuesToWrite the values to write to the target context.
      * @param targetContext the context to write to; the context type should be
      * {@link Class#isAssignableFrom(Class) assignable} from the type returned by
      * {@link #writeableType()}.

context-propagation/src/main/java/io/micrometer/context/ContextSnapshot.java

@@ -15,6 +15,7 @@
  */
 package io.micrometer.context;
 
+import java.util.Map;
 import java.util.concurrent.Callable;
 import java.util.concurrent.Executor;
 import java.util.function.Consumer;
@@ -25,6 +26,15 @@
  * methods to propagate those values.
  *
  * <p>
+ * Implementations are disallowed to store {@code null} values. If a {@link ThreadLocal}
+ * is not set, or it's value is {@code null}, there is no way of distinguishing one from
+ * the other. In such a case, the {@link ContextSnapshot} simply must not contain a
+ * capture for the particular {@link ThreadLocal}. Implementations should filter out any
+ * {@code null} values after reading into the storage also obtained by calling
+ * {@link ContextAccessor#readValues(Object, Predicate, Map)}, and should likewise ignore
+ * {@code null} values from {@link ContextAccessor#readValue(Object, Object)}.
+ *
+ * <p>
  * Use static factory methods on this interface to create a snapshot.
  *
  * @author Rossen Stoyanchev

context-propagation/src/main/java/io/micrometer/context/DefaultContextSnapshot.java

@@ -17,6 +17,7 @@
 
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Objects;
 import java.util.function.Predicate;
 
 /**
@@ -76,7 +77,9 @@ public Scope setThreadLocals(Predicate<Object> keyPredicate) {
         for (ThreadLocalAccessor<?> accessor : this.contextRegistry.getThreadLocalAccessors()) {
             Object key = accessor.key();
             if (keyPredicate.test(key) && this.containsKey(key)) {
-                previousValues = setThreadLocal(key, get(key), accessor, previousValues);
+                Object value = get(key);
+                assert value != null : "snapshot contains disallowed null mapping for key: " + key;
+                previousValues = setThreadLocal(key, value, accessor, previousValues);
             }
         }
         return DefaultScope.from(previousValues, this.contextRegistry);
@@ -162,6 +165,9 @@ static DefaultContextSnapshot captureFromContext(Predicate<Object> keyPredicate,
             snapshot = (snapshot != null ? snapshot : new DefaultContextSnapshot(contextRegistry));
             ((ContextAccessor<Object, ?>) accessor).readValues(context, keyPredicate, snapshot);
         }
+        if (snapshot != null) {
+            snapshot.values().removeIf(Objects::isNull);
+        }
         return (snapshot != null ? snapshot : emptyContextSnapshot);
     }
 

context-propagation/src/main/java/io/micrometer/context/ThreadLocalAccessor.java

@@ -38,13 +38,16 @@ public interface ThreadLocalAccessor<V> {
     Object key();
 
     /**
-     * Return the current {@link ThreadLocal} value, or {@code null} if not set.
+     * Return the current {@link ThreadLocal} value.
      */
     @Nullable
     V getValue();
 
     /**
      * Set the {@link ThreadLocal} value.
+     * <p>
+     * The argument will not be {@code null} when called from {@link ContextSnapshot}
+     * implementations, which are disallowed to store mappings to {@code null}.
      * @param value the value to set
      */
     void setValue(V value);
@@ -56,6 +59,9 @@ public interface ThreadLocalAccessor<V> {
 
     /**
      * Remove the current {@link ThreadLocal} value and set the previously stored one.
+     * <p>
+     * The argument will not be {@code null} when called from {@link ContextSnapshot}
+     * implementations, which are disallowed to store mappings to {@code null}.
      * @param previousValue previous value to set
      * @since 1.0.1
      */

context-propagation/src/test/java/io/micrometer/context/DefaultContextSnapshotTests.java

@@ -16,13 +16,15 @@
 package io.micrometer.context;
 
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.Map;
 
 import io.micrometer.context.ContextSnapshot.Scope;
 import org.assertj.core.api.BDDAssertions;
 import org.junit.jupiter.api.Test;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
 import static org.assertj.core.api.BDDAssertions.then;
 
 /**
@@ -136,24 +138,6 @@ void should_throw_an_exception_when_no_keys_are_passed_for_version_with_no_regis
             .hasMessage("You must provide at least one key when setting thread locals");
     }
 
-    @Test
-    void should_not_fail_on_empty_thread_local() {
-        this.registry.registerThreadLocalAccessor(new ObservationThreadLocalAccessor());
-
-        then(ObservationThreadLocalHolder.getValue()).isNull();
-
-        ContextSnapshot snapshot = ContextSnapshot.captureAllUsing(key -> true, this.registry);
-
-        ObservationThreadLocalHolder.reset();
-        then(ObservationThreadLocalHolder.getValue()).isNull();
-
-        try (Scope scope = snapshot.setThreadLocals()) {
-            then(ObservationThreadLocalHolder.getValue()).isNull();
-        }
-
-        then(ObservationThreadLocalHolder.getValue()).isNull();
-    }
-
     @Test
     void should_filter_thread_locals_on_capture() {
         ThreadLocal<String> fooThreadLocal = new ThreadLocal<>();
@@ -209,6 +193,96 @@ void should_filter_thread_locals_on_restore() {
         then(barThreadLocal.get()).isNull();
     }
 
+    @Test
+    void should_not_fail_on_empty_thread_local() {
+        this.registry.registerThreadLocalAccessor(new ObservationThreadLocalAccessor());
+
+        then(ObservationThreadLocalHolder.getValue()).isNull();
+
+        ContextSnapshot snapshot = ContextSnapshot.captureAll(this.registry);
+
+        ObservationThreadLocalHolder.reset();
+        then(ObservationThreadLocalHolder.getValue()).isNull();
+
+        try (Scope scope = snapshot.setThreadLocals()) {
+            then(ObservationThreadLocalHolder.getValue()).isNull();
+        }
+
+        then(ObservationThreadLocalHolder.getValue()).isNull();
+    }
+
+    @Test
+    @SuppressWarnings("unchecked")
+    void should_ignore_null_value_in_source_context() {
+        String key = "foo";
+        ThreadLocal<String> fooThreadLocal = new ThreadLocal<>();
+        TestThreadLocalAccessor fooThreadLocalAccessor = new TestThreadLocalAccessor(key, fooThreadLocal);
+
+        this.registry.registerContextAccessor(new TestContextAccessor());
+        this.registry.registerThreadLocalAccessor(fooThreadLocalAccessor);
+
+        // We capture null from an uninitialized ThreadLocal:
+        String emptyValue = fooThreadLocalAccessor.getValue();
+        Map<String, String> sourceContext = Collections.singletonMap(key, emptyValue);
+
+        ContextSnapshot snapshot = ContextSnapshot.captureFromContext(this.registry, sourceContext);
+
+        HashMap<Object, Object> snapshotStorage = (HashMap<Object, Object>) snapshot;
+        assertThat(snapshotStorage).isEmpty();
+
+        try (Scope scope = snapshot.setThreadLocals()) {
+            assertThat(fooThreadLocalAccessor.getValue()).isEqualTo(emptyValue);
+        }
+        assertThat(fooThreadLocalAccessor.getValue()).isEqualTo(emptyValue);
+    }
+
+    @Test
+    @SuppressWarnings("unchecked")
+    void should_ignore_null_mapping_in_source_context_when_skipping_intermediate_snapshot() {
+        String key = "foo";
+        ThreadLocal<String> fooThreadLocal = new ThreadLocal<>();
+        TestThreadLocalAccessor fooThreadLocalAccessor = new TestThreadLocalAccessor(key, fooThreadLocal);
+
+        this.registry.registerContextAccessor(new TestContextAccessor());
+        this.registry.registerThreadLocalAccessor(fooThreadLocalAccessor);
+
+        // We capture null from an uninitialized ThreadLocal:
+        String emptyValue = fooThreadLocalAccessor.getValue();
+        Map<String, String> sourceContext = Collections.singletonMap(key, emptyValue);
+
+        // Validate setting all values
+        try (Scope scope = ContextSnapshot.setAllThreadLocalsFrom(sourceContext, this.registry)) {
+            assertThat(fooThreadLocalAccessor.getValue()).isEqualTo(emptyValue);
+        }
+        assertThat(fooThreadLocalAccessor.getValue()).isEqualTo(emptyValue);
+
+        // Validate setting a subset of values
+        try (Scope scope = ContextSnapshot.setThreadLocalsFrom(sourceContext, this.registry, key)) {
+            assertThat(fooThreadLocalAccessor.getValue()).isEqualTo(emptyValue);
+        }
+        assertThat(fooThreadLocalAccessor.getValue()).isEqualTo(emptyValue);
+    }
+
+    @Test
+    @SuppressWarnings("unchecked")
+    void should_fail_assertion_if_null_value_makes_it_into_snapshot() {
+        ThreadLocal<String> fooThreadLocal = new ThreadLocal<>();
+        TestThreadLocalAccessor fooThreadLocalAccessor = new TestThreadLocalAccessor("foo", fooThreadLocal);
+        this.registry.registerThreadLocalAccessor(fooThreadLocalAccessor);
+
+        fooThreadLocal.set("present");
+
+        ContextSnapshot snapshot = ContextSnapshot.captureAll(this.registry);
+        fooThreadLocal.remove();
+
+        HashMap<Object, Object> snapshotStorage = (HashMap<Object, Object>) snapshot;
+        // Imitating a broken implementation that let mapping to null into the storage:
+        snapshotStorage.put("foo", null);
+
+        assertThatExceptionOfType(AssertionError.class).isThrownBy(snapshot::setThreadLocals)
+            .withMessage("snapshot contains disallowed null mapping for key: foo");
+    }
+
     @Test
     void toString_should_include_values() {
         ThreadLocal<String> fooThreadLocal = new ThreadLocal<>();

context-propagation/src/test/java/io/micrometer/context/ObservationThreadLocalAccessor.java

@@ -15,6 +15,8 @@
  */
 package io.micrometer.context;
 
+import java.util.Objects;
+
 /**
  * Example {@link ThreadLocalAccessor} implementation.
  */
@@ -34,6 +36,9 @@ public String getValue() {
 
     @Override
     public void setValue(String value) {
+        // ThreadLocalAccessor API is @NonNullApi by default
+        // so we don't expect null here
+        Objects.requireNonNull(value);
         ObservationThreadLocalHolder.setValue(value);
     }
 
@@ -42,4 +47,12 @@ public void reset() {
         ObservationThreadLocalHolder.reset();
     }
 
+    @Override
+    public void restore(String previousValue) {
+        // ThreadLocalAccessor API is @NonNullApi by default
+        // so we don't expect null here
+        Objects.requireNonNull(previousValue);
+        setValue(previousValue);
+    }
+
 }

context-propagation/src/test/java/io/micrometer/context/ObservationThreadLocalHolder.java

@@ -19,14 +19,11 @@ public class ObservationThreadLocalHolder {
 
     private static final ThreadLocal<String> holder = new ThreadLocal<>();
 
-    public static void resetValue() {
-        holder.remove();
-    }
-
     public static void setValue(String value) {
         holder.set(value);
     }
 
+    @Nullable
     public static String getValue() {
         return holder.get();
     }

context-propagation/src/test/java/io/micrometer/context/TestThreadLocalAccessor.java

@@ -15,6 +15,8 @@
  */
 package io.micrometer.context;
 
+import java.util.Objects;
+
 /**
  * ThreadLocalAccessor for testing purposes with a given key and {@link ThreadLocal}
  * instance.
@@ -46,6 +48,9 @@ public String getValue() {
 
     @Override
     public void setValue(String value) {
+        // ThreadLocalAccessor API is @NonNullApi by default
+        // so we don't expect null here
+        Objects.requireNonNull(value);
         this.threadLocal.set(value);
     }