Fix key removal issues in Thread Context (#3210)

This backports #3050 to the 2.24.x branch.

Co-authored-by: Volkan Yazıcı <[email protected]>

log4j-api-test/src/test/java/org/apache/logging/log4j/CloseableThreadContextTest.java

@@ -19,7 +19,9 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertNull;
 
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -244,4 +246,70 @@ public void pushAllWillPushAllValues() {
         }
         assertEquals("", ThreadContext.peek());
     }
+
+    /**
+     * User provided test stressing nesting using {@link CloseableThreadContext#put(String, String)}.
+     *
+     * @see <a href="https://github.com/apache/logging-log4j2/issues/2946#issuecomment-2382935426">#2946</a>
+     */
+    @Test
+    void testAutoCloseableThreadContextPut() {
+        try (final CloseableThreadContext.Instance ctc1 = CloseableThreadContext.put("outer", "one")) {
+            try (final CloseableThreadContext.Instance ctc2 = CloseableThreadContext.put("outer", "two")) {
+                assertEquals("two", ThreadContext.get("outer"));
+
+                try (final CloseableThreadContext.Instance ctc3 = CloseableThreadContext.put("inner", "one")) {
+                    assertEquals("one", ThreadContext.get("inner"));
+
+                    ThreadContext.put(
+                            "not-in-closeable", "true"); // Remove this line, and closing context behaves as expected
+                    assertEquals("two", ThreadContext.get("outer"));
+                }
+
+                assertEquals("two", ThreadContext.get("outer"));
+                assertNull(ThreadContext.get("inner")); // Test fails here
+            }
+
+            assertEquals("one", ThreadContext.get("outer"));
+            assertNull(ThreadContext.get("inner"));
+        }
+        assertEquals("true", ThreadContext.get("not-in-closeable"));
+
+        assertNull(ThreadContext.get("inner"));
+        assertNull(ThreadContext.get("outer"));
+    }
+
+    /**
+     * User provided test stressing nesting using {@link CloseableThreadContext#putAll(Map)}.
+     *
+     * @see <a href="https://github.com/apache/logging-log4j2/issues/2946#issuecomment-2382935426">#2946</a>
+     */
+    @Test
+    void testAutoCloseableThreadContextPutAll() {
+        try (final CloseableThreadContext.Instance ctc1 = CloseableThreadContext.put("outer", "one")) {
+            try (final CloseableThreadContext.Instance ctc2 = CloseableThreadContext.put("outer", "two")) {
+                assertEquals("two", ThreadContext.get("outer"));
+
+                try (final CloseableThreadContext.Instance ctc3 = CloseableThreadContext.put("inner", "one")) {
+                    assertEquals("one", ThreadContext.get("inner"));
+
+                    ThreadContext.put(
+                            "not-in-closeable", "true"); // Remove this line, and closing context behaves as expected
+                    ThreadContext.putAll(Collections.singletonMap("inner", "two")); // But this is not a problem
+                    assertEquals("two", ThreadContext.get("inner"));
+                    assertEquals("two", ThreadContext.get("outer"));
+                }
+
+                assertEquals("two", ThreadContext.get("outer"));
+                assertNull(ThreadContext.get("inner")); // This is where the test fails
+            }
+
+            assertEquals("one", ThreadContext.get("outer"));
+            assertNull(ThreadContext.get("inner"));
+        }
+        assertEquals("true", ThreadContext.get("not-in-closeable"));
+
+        assertNull(ThreadContext.get("inner"));
+        assertNull(ThreadContext.get("outer"));
+    }
 }

log4j-api-test/src/test/java/org/apache/logging/log4j/internal/map/UnmodifiableArrayBackedMapTest.java

@@ -34,6 +34,7 @@
 import java.util.Set;
 import org.apache.logging.log4j.util.ReadOnlyStringMap;
 import org.apache.logging.log4j.util.TriConsumer;
+import org.assertj.core.api.Assertions;
 import org.junit.jupiter.api.Test;
 
 public class UnmodifiableArrayBackedMapTest {
@@ -373,4 +374,23 @@ public void testToMap() {
         // verify same instance, not just equals()
         assertTrue(map == map.toMap());
     }
+
+    @Test
+    void copyAndRemoveAll_should_work() {
+
+        // Create the actual map
+        UnmodifiableArrayBackedMap actualMap = UnmodifiableArrayBackedMap.EMPTY_MAP;
+        actualMap = actualMap.copyAndPut("outer", "two");
+        actualMap = actualMap.copyAndPut("inner", "one");
+        actualMap = actualMap.copyAndPut("not-in-closeable", "true");
+
+        // Create the expected map
+        UnmodifiableArrayBackedMap expectedMap = UnmodifiableArrayBackedMap.EMPTY_MAP;
+        expectedMap = expectedMap.copyAndPut("outer", "two");
+        expectedMap = expectedMap.copyAndPut("not-in-closeable", "true");
+
+        // Remove the key and verify
+        actualMap = actualMap.copyAndRemoveAll(Collections.singleton("inner"));
+        Assertions.assertThat(actualMap).isEqualTo(expectedMap);
+    }
 }

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

@@ -350,82 +350,64 @@ public UnmodifiableArrayBackedMap copyAndRemove(String key) {
     }
 
     /**
-     * Creates a new instance that contains the same entries as this map, minus all
-     * of the keys passed in the arguments.
-     *
-     * @param key
-     * @param value
-     * @return
+     * Creates a new instance where the entries of provided keys are removed.
      */
     public UnmodifiableArrayBackedMap copyAndRemoveAll(Iterable<String> keysToRemoveIterable) {
+
+        // Short-circuit if the map is empty
         if (isEmpty()) {
-            // shortcut: if this map is empty, the result will continue to be empty
             return EMPTY_MAP;
         }
 
-        // now we build a Set of keys to remove
-        Set<String> keysToRemoveSet;
+        // Collect distinct keys to remove
+        final Set<String> keysToRemove;
         if (keysToRemoveIterable instanceof Set) {
-            // we already have a set, let's cast it and reuse it
-            keysToRemoveSet = (Set<String>) keysToRemoveIterable;
+            keysToRemove = (Set<String>) keysToRemoveIterable;
         } else {
-            // iterate through the keys and build a set
-            keysToRemoveSet = new HashSet<>();
-            for (String key : keysToRemoveIterable) {
-                keysToRemoveSet.add(key);
+            keysToRemove = new HashSet<>();
+            for (final String key : keysToRemoveIterable) {
+                keysToRemove.add(key);
             }
         }
 
-        int firstIndexToKeep = -1;
-        int lastIndexToKeep = -1;
-        int destinationIndex = 0;
-        int numEntriesKept = 0;
-        // build the new map
-        UnmodifiableArrayBackedMap newMap = new UnmodifiableArrayBackedMap(numEntries);
-        for (int indexInCurrentMap = 0; indexInCurrentMap < numEntries; indexInCurrentMap++) {
-            // for each key in this map, check whether it's in the set we built above
-            Object key = backingArray[getArrayIndexForKey(indexInCurrentMap)];
-            if (!keysToRemoveSet.contains(key)) {
-                // this key should be kept
-                if (firstIndexToKeep == -1) {
-                    firstIndexToKeep = indexInCurrentMap;
-                }
-                lastIndexToKeep = indexInCurrentMap;
-            } else if (lastIndexToKeep > 0) {
-                // we hit a remove, copy any keys that are known ready
-                int numEntriesToCopy = lastIndexToKeep - firstIndexToKeep + 1;
-                System.arraycopy(
-                        backingArray,
-                        getArrayIndexForKey(firstIndexToKeep),
-                        newMap.backingArray,
-                        getArrayIndexForKey(destinationIndex),
-                        numEntriesToCopy * 2);
-                firstIndexToKeep = -1;
-                lastIndexToKeep = -1;
-                destinationIndex += numEntriesToCopy;
-                numEntriesKept += numEntriesToCopy;
-            }
-        }
+        // Create the new map
+        final UnmodifiableArrayBackedMap oldMap = this;
+        final int oldMapEntryCount = oldMap.numEntries;
+        final UnmodifiableArrayBackedMap newMap = new UnmodifiableArrayBackedMap(oldMapEntryCount);
 
-        if (lastIndexToKeep > -1) {
-            // at least one key still requires copying
-            int numEntriesToCopy = lastIndexToKeep - firstIndexToKeep + 1;
-            System.arraycopy(
-                    backingArray,
-                    getArrayIndexForKey(firstIndexToKeep),
-                    newMap.backingArray,
-                    getArrayIndexForKey(destinationIndex),
-                    numEntriesToCopy * 2);
-            numEntriesKept += numEntriesToCopy;
+        // Short-circuit if there is nothing to remove
+        if (keysToRemove.isEmpty()) {
+            System.arraycopy(oldMap.backingArray, 0, newMap.backingArray, 0, oldMapEntryCount * 2);
+            newMap.numEntries = oldMapEntryCount;
+            return this;
         }
 
-        if (numEntriesKept == 0) {
-            return EMPTY_MAP;
+        // Iterate over old map entries
+        int newMapEntryIndex = 0;
+        for (int oldMapEntryIndex = 0; oldMapEntryIndex < oldMapEntryCount; oldMapEntryIndex++) {
+            final int oldMapKeyIndex = getArrayIndexForKey(oldMapEntryIndex);
+            final Object key = oldMap.backingArray[oldMapKeyIndex];
+
+            // Skip entries of removed keys
+            @SuppressWarnings("SuspiciousMethodCalls")
+            final boolean removed = keysToRemove.contains(key);
+            if (removed) {
+                continue;
+            }
+
+            // Copy the entry
+            final int oldMapValueIndex = getArrayIndexForValue(oldMapEntryIndex);
+            final Object value = oldMap.backingArray[oldMapValueIndex];
+            final int newMapKeyIndex = getArrayIndexForKey(newMapEntryIndex);
+            final int newMapValueIndex = getArrayIndexForValue(newMapEntryIndex);
+            newMap.backingArray[newMapKeyIndex] = key;
+            newMap.backingArray[newMapValueIndex] = value;
+            newMapEntryIndex++;
         }
 
-        newMap.numEntries = numEntriesKept;
+        // Cap and return the new map
+        newMap.numEntries = newMapEntryIndex;
         newMap.updateNumEntriesInArray();
-
         return newMap;
     }
 

src/changelog/.2.x.x/3048_fix_ThreadContext_remove.xml

@@ -0,0 +1,8 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<entry xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+       xmlns="https://logging.apache.org/xml/ns"
+       xsi:schemaLocation="https://logging.apache.org/xml/ns https://logging.apache.org/xml/ns/log4j-changelog-0.xsd"
+       type="fixed">
+  <issue id="3048" link="https://github.com/apache/logging-log4j2/pull/3048"/>
+  <description format="asciidoc">Fix key removal issues in Thread Context</description>
+</entry>