diff --git a/log4j-api/src/main/java/org/apache/logging/log4j/CloseableThreadContext.java b/log4j-api/src/main/java/org/apache/logging/log4j/CloseableThreadContext.java index d84f395..d3b569d 100644 --- a/log4j-api/src/main/java/org/apache/logging/log4j/CloseableThreadContext.java +++ b/log4j-api/src/main/java/org/apache/logging/log4j/CloseableThreadContext.java @@ -1,156 +1,130 @@ package org.apache.logging.log4j; import java.util.HashMap; +import java.util.Iterator; import java.util.Map; -import org.apache.logging.log4j.util.Strings; - /** * Adds entries to the {@link ThreadContext stack or map} and them removes them when the object is closed, e.g. as part * of a try-with-resources. - * + * * @since 2.6 */ public class CloseableThreadContext implements AutoCloseable { + private int pushCount = 0; + private final Map originalValues = new HashMap<>(); + + private CloseableThreadContext() { + } + /** * Pushes new diagnostic context information on to the Thread Context Stack. The information will be popped off when * the instance is closed. - * - * @param message - * The new diagnostic context information. + * + * @param message The new diagnostic context information. * @return a new instance that will back out the changes when closed. */ public static CloseableThreadContext push(final String message) { - return new CloseableThreadContext(message); + return new CloseableThreadContext().andPush(message); } /** * Pushes new diagnostic context information on to the Thread Context Stack. The information will be popped off when * the instance is closed. - * - * @param message - * The new diagnostic context information. - * @param args - * Parameters for the message. - * @return a new instance that will back out the changes when closed. + * + * @param message The new diagnostic context information. + * @return the CloseableThreadContext instance. */ - public static CloseableThreadContext push(final String message, final Object... args) { - return new CloseableThreadContext(message, args); + public CloseableThreadContext andPush(final String message) { + ThreadContext.push(message); + pushCount++; + return this; } /** - * Populates the Thread Context Map with the supplied key/value pairs. Any existing keys in the - * {@link ThreadContext} will be replaced with the supplied values, and restored back to their original values when + * Pushes new diagnostic context information on to the Thread Context Stack. The information will be popped off when * the instance is closed. * - * @param firstKey - * The first key to be added - * @param firstValue - * The first value to be added - * @param subsequentKeyValuePairs - * Any subsequent key/value pairs to be added. Note: If the last key does not have a corresponding value - * then an empty String will be used as a value. + * @param message The new diagnostic context information. + * @param args Parameters for the message. * @return a new instance that will back out the changes when closed. */ - public static CloseableThreadContext put(final String firstKey, final String firstValue, - final String... subsequentKeyValuePairs) { - return new CloseableThreadContext(firstKey, firstValue, subsequentKeyValuePairs); - } - - private final boolean isStack; - - private final Map oldValues = new HashMap<>(); - - /** - * Creates an instance of a {@code CloseableThreadContext} that pushes new diagnostic context information on to the - * Thread Context Stack. The information will be popped off when the instance is closed. - * - * @param message - * The new diagnostic context information. - */ - public CloseableThreadContext(final String message) { - this.isStack = true; - ThreadContext.push(message); + public static CloseableThreadContext push(final String message, final Object... args) { + return new CloseableThreadContext().andPush(message, args); } /** - * Creates an instance of a {@code CloseableThreadContext} that pushes new diagnostic context information on to the - * Thread Context Stack. The information will be popped off when the instance is closed. - * - * @param message - * The new diagnostic context information. - * @param args - * Parameters for the message. + * Pushes new diagnostic context information on to the Thread Context Stack. The information will be popped off when + * the instance is closed. + * + * @param message The new diagnostic context information. + * @param args Parameters for the message. + * @return the CloseableThreadContext instance. */ - public CloseableThreadContext(final String message, final Object... args) { - this.isStack = true; + public CloseableThreadContext andPush(final String message, final Object... args) { ThreadContext.push(message, args); + pushCount++; + return this; } /** - * Creates an instance of a {@code CloseableThreadContext} that populates the Thread Context Map with the supplied - * key/value pairs. Any existing keys in the ThreadContext will be replaced with the supplied values, and restored - * back to their original values when the instance is closed. + * Populates the Thread Context Map with the supplied key/value pairs. Any existing keys in the + * {@link ThreadContext} will be replaced with the supplied values, and restored back to their original values when + * the instance is closed. * - * @param firstKey - * The first key to be added - * @param firstValue - * The first value to be added - * @param subsequentKeyValuePairs - * Any subsequent key/value pairs to be added. Note: If the last key does not have a corresponding value - * then an empty String will be used as a value. + * @param key The key to be added + * @param value The value to be added + * @return a new instance that will back out the changes when closed. */ - public CloseableThreadContext(final String firstKey, final String firstValue, - final String... subsequentKeyValuePairs) { - this.isStack = false; - storeAndSet(firstKey, firstValue); - for (int i = 0; i < subsequentKeyValuePairs.length; i += 2) { - final String key = subsequentKeyValuePairs[i]; - final String value = (i + 1) < subsequentKeyValuePairs.length ? subsequentKeyValuePairs[i + 1] : Strings.EMPTY; - storeAndSet(key, value); + public static CloseableThreadContext put(final String key, final String value) { + return new CloseableThreadContext().andPut(key, value); + } + + public CloseableThreadContext andPut(final String key, final String value) { + // If there are no existing values, a null will be stored as an old value + if (!originalValues.containsKey(key)) { + originalValues.put(key, ThreadContext.get(key)); } + ThreadContext.put(key, value); + return this; } + /** * Removes the values from the {@link ThreadContext}. *

- * If this {@code CloseableThreadContext} was added to the {@link ThreadContext} stack, then this will pop - * the diagnostic information off the stack. + * Values pushed to the {@link ThreadContext} stack will be popped off. *

*

- * If the {@code CloseableThreadContext} was added to the {@link ThreadContext} map, then this will either - * remove the values that were added, or restore them to their original values it they already existed. + * Values put on the {@link ThreadContext} map will be removed, or restored to their original values it they already existed. *

*/ @Override public void close() { - if (this.isStack) { - closeStack(); - } else { - closeMap(); - } + closeStack(); + closeMap(); } private void closeMap() { - for (final Map.Entry entry : oldValues.entrySet()) { - // If the old value was null, remove it from the ThreadContext - if (null == entry.getValue()) { - ThreadContext.remove(entry.getKey()); + for (Iterator> it = originalValues.entrySet().iterator(); it.hasNext(); ) { + final Map.Entry entry = it.next(); + final String key = entry.getKey(); + final String originalValue = entry.getValue(); + if (null == originalValue) { + ThreadContext.remove(key); } else { - ThreadContext.put(entry.getKey(), entry.getValue()); + ThreadContext.put(key, originalValue); } + it.remove(); } } - private String closeStack() { - return ThreadContext.pop(); - } - - private void storeAndSet(final String key, final String value) { - // If there are no existing values, a null will be stored as an old value - oldValues.put(key, ThreadContext.get(key)); - ThreadContext.put(key, value); + private void closeStack() { + for (int i = 0; i < pushCount; i++) { + ThreadContext.pop(); + } + pushCount = 0; } } diff --git a/log4j-api/src/test/java/org/apache/logging/log4j/CloseableThreadContextTest.java b/log4j-api/src/test/java/org/apache/logging/log4j/CloseableThreadContextTest.java index 42ed624..b808c26 100644 --- a/log4j-api/src/test/java/org/apache/logging/log4j/CloseableThreadContextTest.java +++ b/log4j-api/src/test/java/org/apache/logging/log4j/CloseableThreadContextTest.java @@ -32,7 +32,7 @@ public class CloseableThreadContextTest { public void shouldAddTwoEntriesToTheMap() throws Exception { final String key2 = "key2"; final String value2 = "value2"; - try (final CloseableThreadContext ignored = CloseableThreadContext.put(key, value, key2, value2)) { + try (final CloseableThreadContext ignored = CloseableThreadContext.put(key, value).andPut(key2, value2)) { assertThat(ThreadContext.get(key), is(value)); assertThat(ThreadContext.get(key2), is(value2)); } @@ -64,6 +64,17 @@ public class CloseableThreadContextTest { } @Test + public void ifTheSameKeyIsAddedTwiceTheOriginalShouldBeUsed() throws Exception { + final String oldValue = "oldValue"; + final String secondValue = "innerValue"; + ThreadContext.put(key, oldValue); + try (final CloseableThreadContext ignored = CloseableThreadContext.put(key, value).andPut(key, secondValue)) { + assertThat(ThreadContext.get(key), is(secondValue)); + } + assertThat(ThreadContext.get(key), is(oldValue)); + } + + @Test public void shouldPushAndPopAnEntryToTheStack() throws Exception { final String message = "message"; try (final CloseableThreadContext ignored = CloseableThreadContext.push(message)) { @@ -73,6 +84,16 @@ public class CloseableThreadContextTest { } @Test + public void shouldPushAndPopTwoEntriesToTheStack() throws Exception { + final String message1 = "message1"; + final String message2 = "message2"; + try (final CloseableThreadContext ignored = CloseableThreadContext.push(message1).andPush(message2)) { + assertThat(ThreadContext.peek(), is(message2)); + } + assertThat(ThreadContext.peek(), is("")); + } + + @Test public void shouldPushAndPopAParameterizedEntryToTheStack() throws Exception { final String parameterizedMessage = "message {}"; final String parameterizedMessageParameter = "param"; @@ -93,12 +114,63 @@ public class CloseableThreadContextTest { } @Test - public void shouldUseAnEmptyStringIfNoValueIsSupplied() throws Exception { - final String key2 = "key2"; - try (final CloseableThreadContext ignored = CloseableThreadContext.put(key, value, key2)) { + public void shouldAddEntriesToBothStackAndMap() throws Exception { + final String stackValue = "something"; + try (final CloseableThreadContext ignored = CloseableThreadContext.put(key, value).andPush(stackValue)) { assertThat(ThreadContext.get(key), is(value)); - assertThat(ThreadContext.get(key2), is("")); + assertThat(ThreadContext.peek(), is(stackValue)); } + assertThat(ThreadContext.containsKey(key), is(false)); + assertThat(ThreadContext.peek(), is("")); + } + + @Test + public void canReuseCloseableThreadContext() throws Exception { + final String stackValue = "something"; + // Create a ctc and close it + final CloseableThreadContext ctc = CloseableThreadContext.push(stackValue).andPut(key, value); + assertThat(ThreadContext.get(key), is(value)); + assertThat(ThreadContext.peek(), is(stackValue)); + ctc.close(); + + assertThat(ThreadContext.containsKey(key), is(false)); + assertThat(ThreadContext.peek(), is("")); + + final String anotherKey = "key2"; + final String anotherValue = "value2"; + final String anotherStackValue = "something else"; + // Use it again + ctc.andPush(anotherStackValue).andPut(anotherKey, anotherValue); + assertThat(ThreadContext.get(anotherKey), is(anotherValue)); + assertThat(ThreadContext.peek(), is(anotherStackValue)); + ctc.close(); + + assertThat(ThreadContext.containsKey(anotherKey), is(false)); + assertThat(ThreadContext.peek(), is("")); + } + + @Test + public void closeIsIdempotent() throws Exception { + + final String originalMapValue = "map to keep"; + final String originalStackValue = "stack to keep"; + ThreadContext.put(key, originalMapValue); + ThreadContext.push(originalStackValue); + + final String newMapValue = "temp map value"; + final String newStackValue = "temp stack to keep"; + final CloseableThreadContext ctc = CloseableThreadContext.push(newStackValue).andPut(key, newMapValue); + + ctc.close(); + assertThat(ThreadContext.get(key), is(originalMapValue)); + assertThat(ThreadContext.peek(), is(originalStackValue)); + + ctc.close(); + assertThat(ThreadContext.get(key), is(originalMapValue)); + assertThat(ThreadContext.peek(), is(originalStackValue)); + + + } } \ No newline at end of file diff --git a/src/site/xdoc/manual/thread-context.xml b/src/site/xdoc/manual/thread-context.xml index a073dae..8cdff56 100644 --- a/src/site/xdoc/manual/thread-context.xml +++ b/src/site/xdoc/manual/thread-context.xml @@ -120,7 +120,8 @@ try (final CloseableThreadContext ctc = CloseableThreadContext.push(UUID.randomU 

 // Add to the ThreadContext map for this try block only;
-try (final CloseableThreadContext ctc = CloseableThreadContext.put("id", UUID.randomUUID().toString(), "loginId", session.getAttribute("loginId"))) { 
+try (final CloseableThreadContext ctc = CloseableThreadContext.put("id", UUID.randomUUID().toString())
+                                                                .andPut("loginId", session.getAttribute("loginId"))) {
 
     logger.debug("Message 1");
 .