diff --git a/log4j-api-test/src/test/java/org/apache/logging/log4j/CloseableThreadContextTest.java b/log4j-api-test/src/test/java/org/apache/logging/log4j/CloseableThreadContextTest.java
index 56e0086a49f..bc053b080c3 100644
--- a/log4j-api-test/src/test/java/org/apache/logging/log4j/CloseableThreadContextTest.java
+++ b/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 #2946
+ */
+ @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 #2946
+ */
+ @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"));
+ }
}
diff --git a/log4j-api-test/src/test/java/org/apache/logging/log4j/internal/map/UnmodifiableArrayBackedMapTest.java b/log4j-api-test/src/test/java/org/apache/logging/log4j/internal/map/UnmodifiableArrayBackedMapTest.java
index 3f43e3e2c4d..092c082ed4e 100644
--- a/log4j-api-test/src/test/java/org/apache/logging/log4j/internal/map/UnmodifiableArrayBackedMapTest.java
+++ b/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);
+ }
}
diff --git a/log4j-api/src/main/java/org/apache/logging/log4j/internal/map/UnmodifiableArrayBackedMap.java b/log4j-api/src/main/java/org/apache/logging/log4j/internal/map/UnmodifiableArrayBackedMap.java
index eada687424f..078e8c7c70b 100644
--- a/log4j-api/src/main/java/org/apache/logging/log4j/internal/map/UnmodifiableArrayBackedMap.java
+++ b/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 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 keysToRemoveSet;
+ // Collect distinct keys to remove
+ final Set keysToRemove;
if (keysToRemoveIterable instanceof Set) {
- // we already have a set, let's cast it and reuse it
- keysToRemoveSet = (Set) keysToRemoveIterable;
+ keysToRemove = (Set) 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;
}
diff --git a/src/changelog/.2.x.x/3048_fix_ThreadContext_remove.xml b/src/changelog/.2.x.x/3048_fix_ThreadContext_remove.xml
new file mode 100644
index 00000000000..d19091dc0fa
--- /dev/null
+++ b/src/changelog/.2.x.x/3048_fix_ThreadContext_remove.xml
@@ -0,0 +1,8 @@
+
+
+
+ Fix key removal issues in Thread Context
+