1 /*
2 * Licensed to the Apache Software Foundation (ASF) under one or more
3 * contributor license agreements. See the NOTICE file distributed with
4 * this work for additional information regarding copyright ownership.
5 * The ASF licenses this file to You under the Apache license, Version 2.0
6 * (the "License"); you may not use this file except in compliance with
7 * the License. You may obtain a copy of the License at
8 *
9 * http://www.apache.org/licenses/LICENSE-2.0
10 *
11 * Unless required by applicable law or agreed to in writing, software
12 * distributed under the License is distributed on an "AS IS" BASIS,
13 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 * See the license for the specific language governing permissions and
15 * limitations under the license.
16 */
17 package org.apache.logging.log4j;
18
19 import java.util.HashMap;
20 import java.util.Iterator;
21 import java.util.Map;
22
23 /**
24 * Adds entries to the {@link ThreadContext stack or map} and them removes them when the object is closed, e.g. as part
25 * of a try-with-resources. User code can now look like this:
26 * <pre>
27 * try (final CloseableThreadContext.Instance ignored = CloseableThreadContext.put(key1, value1).put(key2, value2)) {
28 * callSomeMethodThatLogsALot();
29 *
30 * // Entries for key1 and key2 are automatically removed from the ThreadContext map when done.
31 * }
32 * </pre>
33 *
34 * @since 2.6
35 */
36 public class CloseableThreadContext {
37
38 private CloseableThreadContext() {
39 }
40
41 /**
42 * Pushes new diagnostic context information on to the Thread Context Stack. The information will be popped off when
43 * the instance is closed.
44 *
45 * @param message The new diagnostic context information.
46 * @return a new instance that will back out the changes when closed.
47 */
48 public static CloseableThreadContext.Instance push(final String message) {
49 return new CloseableThreadContext.Instance().push(message);
50 }
51
52 /**
53 * Pushes new diagnostic context information on to the Thread Context Stack. The information will be popped off when
54 * the instance is closed.
55 *
56 * @param message The new diagnostic context information.
57 * @param args Parameters for the message.
58 * @return a new instance that will back out the changes when closed.
59 */
60 public static CloseableThreadContext.Instance push(final String message, final Object... args) {
61 return new CloseableThreadContext.Instance().push(message, args);
62 }
63
64 /**
65 * Populates the Thread Context Map with the supplied key/value pairs. Any existing keys in the
66 * {@link ThreadContext} will be replaced with the supplied values, and restored back to their original values when
67 * the instance is closed.
68 *
69 * @param key The key to be added
70 * @param value The value to be added
71 * @return a new instance that will back out the changes when closed.
72 */
73 public static CloseableThreadContext.Instance put(final String key, final String value) {
74 return new CloseableThreadContext.Instance().put(key, value);
75 }
76
77 public static class Instance implements AutoCloseable {
78
79 private int pushCount = 0;
80 private final Map<String, String> originalValues = new HashMap<>();
81
82 private Instance() {
83 }
84
85 /**
86 * Pushes new diagnostic context information on to the Thread Context Stack. The information will be popped off when
87 * the instance is closed.
88 *
89 * @param message The new diagnostic context information.
90 * @return the instance that will back out the changes when closed.
91 */
92 public Instance push(final String message) {
93 ThreadContext.push(message);
94 pushCount++;
95 return this;
96 }
97
98 /**
99 * Pushes new diagnostic context information on to the Thread Context Stack. The information will be popped off when
100 * the instance is closed.
101 *
102 * @param message The new diagnostic context information.
103 * @param args Parameters for the message.
104 * @return the instance that will back out the changes when closed.
105 */
106 public Instance push(final String message, final Object[] args) {
107 ThreadContext.push(message, args);
108 pushCount++;
109 return this;
110 }
111
112 /**
113 * Populates the Thread Context Map with the supplied key/value pairs. Any existing keys in the
114 * {@link ThreadContext} will be replaced with the supplied values, and restored back to their original values when
115 * the instance is closed.
116 *
117 * @param key The key to be added
118 * @param value The value to be added
119 * @return the instance that will back out the changes when closed.
120 */
121 public Instance put(final String key, final String value) {
122 // If there are no existing values, a null will be stored as an old value
123 if (!originalValues.containsKey(key)) {
124 originalValues.put(key, ThreadContext.get(key));
125 }
126 ThreadContext.put(key, value);
127 return this;
128 }
129
130 /**
131 * Removes the values from the {@link ThreadContext}.
132 * <p>
133 * Values pushed to the {@link ThreadContext} <em>stack</em> will be popped off.
134 * </p>
135 * <p>
136 * Values put on the {@link ThreadContext} <em>map</em> will be removed, or restored to their original values it they already existed.
137 * </p>
138 */
139 @Override
140 public void close() {
141 closeStack();
142 closeMap();
143 }
144
145 private void closeMap() {
146 for (final Iterator<Map.Entry<String, String>> it = originalValues.entrySet().iterator(); it.hasNext(); ) {
147 final Map.Entry<String, String> entry = it.next();
148 final String key = entry.getKey();
149 final String originalValue = entry.getValue();
150 if (null == originalValue) {
151 ThreadContext.remove(key);
152 } else {
153 ThreadContext.put(key, originalValue);
154 }
155 it.remove();
156 }
157 }
158
159 private void closeStack() {
160 for (int i = 0; i < pushCount; i++) {
161 ThreadContext.pop();
162 }
163 pushCount = 0;
164 }
165 }
166 }