Fix nullability and mutability problems

log4j-api-test/src/main/java/org/apache/logging/log4j/test/spi/ThreadContextMapSuite.java

@@ -17,6 +17,7 @@
 package org.apache.logging.log4j.test.spi;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 
 import java.time.Duration;
 import java.util.concurrent.ExecutorService;
@@ -107,24 +108,19 @@ protected static void assertContextDataCanBeTransferred(final ThreadContextMap t
     }
 
     /**
-     * Ensures the {@code restore(null)} can be used to clear the context data.
-     *
-     * @param threadContext The {@link ThreadContextMap} to test.
+     * Ensures that the saved value is always not {@code null}, even if there are no context data.
+     * <p>
+     *     This requirement allows third-party libraries to store the saved value as value of a map, even if the map
+     *     does not allow nulls.
+     * </p>
      */
-    protected static void assertNullValueClearsTheContextData(final ThreadContextMap threadContext) {
-        final String value = "restoreAcceptsNull";
-        final RuntimeException throwable = new RuntimeException();
-        threadContext.put(KEY, value);
-        final Object saved = threadContext.restore(null);
-        try {
-            assertThat(threadContext.getImmutableMapOrNull()).isNullOrEmpty();
-            throw throwable;
-        } catch (final RuntimeException e) {
-            assertThat(e).isSameAs(throwable);
-            assertThat(threadContext.getImmutableMapOrNull()).isNullOrEmpty();
-        } finally {
-            threadContext.restore(saved);
-        }
-        assertThat(threadContext.get(KEY)).isEqualTo(value);
+    protected static void assertSavedValueNotNullIfMapEmpty(final ThreadContextMap threadContext) {
+        threadContext.clear();
+        final Object saved = threadContext.save();
+        assertThat(saved).as("Saved empty context data.").isNotNull();
+    }
+
+    protected static void assertRestoreDoesNotAcceptNull(final ThreadContextMap threadContext) {
+        assertThrows(NullPointerException.class, () -> threadContext.restore(null));
     }
 }

log4j-api-test/src/test/java/org/apache/logging/log4j/spi/ThreadContextMapTest.java

@@ -94,7 +94,13 @@ void saveAndRestoreMapOnAnotherThread(final ThreadContextMap threadContext) {
 
     @ParameterizedTest
     @MethodSource("localMaps")
-    void restoreAcceptsNull(final ThreadContextMap threadContext) {
-        assertNullValueClearsTheContextData(threadContext);
+    void savedValueNotNullIfMapEmpty(final ThreadContextMap threadContext) {
+        assertSavedValueNotNullIfMapEmpty(threadContext);
+    }
+
+    @ParameterizedTest
+    @MethodSource("localMaps")
+    void restoreDoesNotAcceptNull(final ThreadContextMap threadContext) {
+        assertRestoreDoesNotAcceptNull(threadContext);
     }
 }

log4j-api/pom.xml

@@ -38,7 +38,9 @@
     <bnd-module-name>org.apache.logging.log4j</bnd-module-name>
     <bnd-extra-package-options>
       <!-- Not exported by most OSGi system bundles, hence we use the system classloader to load `sun.reflect.Reflection` -->
-      !sun.reflect
+      !sun.reflect,
+      <!-- Annotations only -->
+      org.jspecify.*;resolution:=optional
     </bnd-extra-package-options>
     <bnd-extra-module-options>
       <!-- Used in StringBuilders through reflection -->
@@ -49,11 +51,19 @@
 
   </properties>
   <dependencies>
+
+    <dependency>
+      <groupId>org.jspecify</groupId>
+      <artifactId>jspecify</artifactId>
+      <scope>provided</scope>
+    </dependency>
+
     <dependency>
       <groupId>org.osgi</groupId>
       <artifactId>org.osgi.core</artifactId>
       <scope>provided</scope>
     </dependency>
+
   </dependencies>
   <build>
     <plugins>

log4j-api/src/main/java/org/apache/logging/log4j/internal/map/StringArrayThreadContextMap.java

@@ -23,6 +23,8 @@
 import org.apache.logging.log4j.util.BiConsumer;
 import org.apache.logging.log4j.util.ReadOnlyStringMap;
 import org.apache.logging.log4j.util.TriConsumer;
+import org.jspecify.annotations.NullMarked;
+import org.jspecify.annotations.Nullable;
 
 /**
  * An equivalent for DefaultThreadContxtMap, except that it's backed by
@@ -34,23 +36,26 @@
  * linearly with the current map size, and callers are advised to minimize this
  * work.
  */
+@NullMarked
 public class StringArrayThreadContextMap implements ThreadContextMap, ReadOnlyStringMap {
     private static final long serialVersionUID = -2635197170958057849L;
 
+    private static final Object[] EMPTY_STATE = {0};
+
     /**
      * Property name ({@value} ) for selecting {@code InheritableThreadLocal} (value "true") or plain
      * {@code ThreadLocal} (value is not "true") in the implementation.
      */
     public static final String INHERITABLE_MAP = "isThreadContextMapInheritable";
 
-    private final ThreadLocal<Object[]> threadLocalMapState;
+    private final ThreadLocal<Object @Nullable []> threadLocalMapState;
 
     public StringArrayThreadContextMap() {
         threadLocalMapState = new ThreadLocal<>();
     }
 
     @Override
-    public void put(final String key, final String value) {
+    public void put(final String key, final @Nullable String value) {
         final Object[] state = threadLocalMapState.get();
         final UnmodifiableArrayBackedMap modifiedMap =
                 UnmodifiableArrayBackedMap.getInstance(state).copyAndPut(key, value);
@@ -65,12 +70,11 @@ public void putAll(final Map<String, String> m) {
     }
 
     @Override
-    public String get(final String key) {
+    public @Nullable String get(final String key) {
         final Object[] state = threadLocalMapState.get();
-        if (state == null) {
-            return null;
-        }
-        return UnmodifiableArrayBackedMap.getInstance(state).get(key);
+        return state == null
+                ? null
+                : UnmodifiableArrayBackedMap.getInstance(state).get(key);
     }
 
     @Override
@@ -130,7 +134,7 @@ public <V, S> void forEach(final TriConsumer<String, ? super V, S> action, final
 
     @SuppressWarnings("unchecked")
     @Override
-    public <V> V getValue(final String key) {
+    public <V> @Nullable V getValue(final String key) {
         return (V) get(key);
     }
 
@@ -144,7 +148,7 @@ public Map<String, String> getCopy() {
     }
 
     @Override
-    public Map<String, String> getImmutableMapOrNull() {
+    public @Nullable Map<String, String> getImmutableMapOrNull() {
         final Object[] state = threadLocalMapState.get();
         return (state == null ? null : UnmodifiableArrayBackedMap.getInstance(state));
     }
@@ -162,17 +166,19 @@ public int size() {
 
     @Override
     public Object save() {
-        return threadLocalMapState.get();
+        final Object[] state = threadLocalMapState.get();
+        return state != null ? state : EMPTY_STATE;
     }
 
     @Override
     @SuppressWarnings("unchecked")
     public Object restore(final Object contextMap) {
         final Object current = save();
-        if (contextMap == null) {
+        final Object[] state = Objects.requireNonNull((Object[]) contextMap);
+        if (UnmodifiableArrayBackedMap.getInstance(state).isEmpty()) {
             clear();
         } else {
-            threadLocalMapState.set((Object[]) contextMap);
+            threadLocalMapState.set(state);
         }
         return current;
     }

log4j-api/src/main/java/org/apache/logging/log4j/spi/CopyOnWriteSortedArrayThreadContextMap.java

@@ -24,6 +24,8 @@
 import org.apache.logging.log4j.util.ReadOnlyStringMap;
 import org.apache.logging.log4j.util.SortedArrayStringMap;
 import org.apache.logging.log4j.util.StringMap;
+import org.jspecify.annotations.NullMarked;
+import org.jspecify.annotations.Nullable;
 
 /**
  * {@code SortedArrayStringMap}-based implementation of the {@code ThreadContextMap} interface that creates a copy of
@@ -34,6 +36,7 @@
  *
  * @since 2.7
  */
+@NullMarked
 class CopyOnWriteSortedArrayThreadContextMap implements ReadOnlyThreadContextMap, ObjectThreadContextMap, CopyOnWrite {
 
     /**
@@ -52,14 +55,14 @@ class CopyOnWriteSortedArrayThreadContextMap implements ReadOnlyThreadContextMap
      */
     protected static final String PROPERTY_NAME_INITIAL_CAPACITY = "log4j2.ThreadContext.initial.capacity";
 
-    private static final StringMap EMPTY_CONTEXT_DATA = new SortedArrayStringMap(1);
+    static final StringMap EMPTY_CONTEXT_DATA = new SortedArrayStringMap(1);
 
     static {
         EMPTY_CONTEXT_DATA.freeze();
     }
 
     private final int initialCapacity;
-    private final ThreadLocal<StringMap> localMap;
+    private final ThreadLocal<@Nullable StringMap> localMap;
 
     public CopyOnWriteSortedArrayThreadContextMap() {
         this(PropertiesUtil.getProperties());
@@ -79,7 +82,7 @@ protected StringMap childValue(final StringMap parentValue) {
                         return stringMap;
                     }
                 }
-                : new ThreadLocal<StringMap>();
+                : new ThreadLocal<>();
     }
 
     /**
@@ -107,12 +110,12 @@ protected StringMap createStringMap(final ReadOnlyStringMap original) {
     }
 
     @Override
-    public void put(final String key, final String value) {
+    public void put(final String key, final @Nullable String value) {
         putValue(key, value);
     }
 
     @Override
-    public void putValue(final String key, final Object value) {
+    public void putValue(final String key, final @Nullable Object value) {
         StringMap map = localMap.get();
         map = map == null ? createStringMap() : createStringMap(map);
         map.putValue(key, value);
@@ -149,12 +152,12 @@ public <V> void putAllValues(final Map<String, V> values) {
     }
 
     @Override
-    public String get(final String key) {
+    public @Nullable String get(final String key) {
         return (String) getValue(key);
     }
 
     @Override
-    public <V> V getValue(final String key) {
+    public <V> @Nullable V getValue(final String key) {
         final StringMap map = localMap.get();
         return map == null ? null : map.<V>getValue(key);
     }
@@ -223,17 +226,19 @@ public boolean isEmpty() {
 
     @Override
     public Object save() {
-        return localMap.get();
+        final StringMap map = localMap.get();
+        return map != null ? map : EMPTY_CONTEXT_DATA;
     }
 
     @Override
     @SuppressWarnings("unchecked")
     public Object restore(final Object contextMap) {
         final Object current = save();
-        if (contextMap == null) {
+        final StringMap map = (StringMap) contextMap;
+        if (map.isEmpty()) {
             clear();
         } else {
-            localMap.set((StringMap) contextMap);
+            localMap.set(map);
         }
         return current;
     }

log4j-api/src/main/java/org/apache/logging/log4j/spi/DefaultThreadContextMap.java

@@ -24,13 +24,16 @@
 import org.apache.logging.log4j.util.PropertiesUtil;
 import org.apache.logging.log4j.util.ReadOnlyStringMap;
 import org.apache.logging.log4j.util.TriConsumer;
+import org.jspecify.annotations.NullMarked;
+import org.jspecify.annotations.Nullable;
 
 /**
  * The actual ThreadContext Map. A new ThreadContext Map is created each time it is updated and the Map stored is always
  * immutable. This means the Map can be passed to other threads without concern that it will be updated. Since it is
  * expected that the Map will be passed to many more log events than the number of keys it contains the performance
  * should be much better than if the Map was copied for each event.
  */
+@NullMarked
 public class DefaultThreadContextMap implements ThreadContextMap, ReadOnlyStringMap {
     private static final long serialVersionUID = 8218007901108944053L;
 
@@ -41,7 +44,7 @@ public class DefaultThreadContextMap implements ThreadContextMap, ReadOnlyString
     public static final String INHERITABLE_MAP = "isThreadContextMapInheritable";
 
     private final boolean useMap;
-    private final ThreadLocal<Map<String, String>> localMap;
+    private final ThreadLocal<@Nullable Map<String, String>> localMap;
 
     public DefaultThreadContextMap() {
         this(true);
@@ -60,17 +63,18 @@ public DefaultThreadContextMap(final boolean useMap) {
         localMap = properties.getBooleanProperty(INHERITABLE_MAP)
                 ? new InheritableThreadLocal<Map<String, String>>() {
                     @Override
-                    protected Map<String, String> childValue(final Map<String, String> parentValue) {
+                    protected @Nullable Map<String, String> childValue(
+                            final @Nullable Map<String, String> parentValue) {
                         return parentValue != null && useMap
                                 ? Collections.unmodifiableMap(new HashMap<>(parentValue))
                                 : null;
                     }
                 }
-                : new ThreadLocal<Map<String, String>>();
+                : new ThreadLocal<>();
     }
 
     @Override
-    public void put(final String key, final String value) {
+    public void put(final String key, final @Nullable String value) {
         if (!useMap) {
             return;
         }
@@ -177,7 +181,7 @@ public Map<String, String> getCopy() {
     }
 
     @Override
-    public Map<String, String> getImmutableMapOrNull() {
+    public @Nullable Map<String, String> getImmutableMapOrNull() {
         return localMap.get();
     }
 
@@ -234,14 +238,16 @@ public boolean equals(final Object obj) {
 
     @Override
     public Object save() {
-        return localMap.get();
+        final Map<String, String> map = localMap.get();
+        return map != null ? map : Collections.emptyMap();
     }
 
     @Override
     @SuppressWarnings("unchecked")
     public Object restore(final Object contextMap) {
         final Object current = save();
-        if (contextMap == null) {
+        final Map<String, String> map = (Map<String, String>) contextMap;
+        if (map.isEmpty()) {
             clear();
         } else {
             localMap.set((Map<String, String>) contextMap);