View Javadoc
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.
26   *
27   * @since 2.6
28   */
29  public class CloseableThreadContext {
30  
31      private CloseableThreadContext() {
32      }
33  
34      /**
35       * Pushes new diagnostic context information on to the Thread Context Stack. The information will be popped off when
36       * the instance is closed.
37       *
38       * @param message The new diagnostic context information.
39       * @return a new instance that will back out the changes when closed.
40       */
41      public static CloseableThreadContext.Instance push(final String message) {
42          return new CloseableThreadContext.Instance().push(message);
43      }
44  
45      /**
46       * Pushes new diagnostic context information on to the Thread Context Stack. The information will be popped off when
47       * the instance is closed.
48       *
49       * @param message The new diagnostic context information.
50       * @param args    Parameters for the message.
51       * @return a new instance that will back out the changes when closed.
52       */
53      public static CloseableThreadContext.Instance push(final String message, final Object... args) {
54          return new CloseableThreadContext.Instance().push(message, args);
55      }
56  
57      /**
58       * Populates the Thread Context Map with the supplied key/value pairs. Any existing keys in the
59       * {@link ThreadContext} will be replaced with the supplied values, and restored back to their original values when
60       * the instance is closed.
61       *
62       * @param key   The  key to be added
63       * @param value The value to be added
64       * @return a new instance that will back out the changes when closed.
65       */
66      public static CloseableThreadContext.Instance put(final String key, final String value) {
67          return new CloseableThreadContext.Instance().put(key, value);
68      }
69  
70      public static class Instance implements AutoCloseable {
71  
72          private int pushCount = 0;
73          private final Map<String, String> originalValues = new HashMap<>();
74  
75          private Instance() {
76          }
77  
78          /**
79           * Pushes new diagnostic context information on to the Thread Context Stack. The information will be popped off when
80           * the instance is closed.
81           *
82           * @param message The new diagnostic context information.
83           * @return the instance that will back out the changes when closed.
84           */
85          public Instance push(final String message) {
86              ThreadContext.push(message);
87              pushCount++;
88              return this;
89          }
90  
91          /**
92           * Pushes new diagnostic context information on to the Thread Context Stack. The information will be popped off when
93           * the instance is closed.
94           *
95           * @param message The new diagnostic context information.
96           * @param args    Parameters for the message.
97           * @return the instance that will back out the changes when closed.
98           */
99          public Instance push(final String message, final Object[] args) {
100             ThreadContext.push(message, args);
101             pushCount++;
102             return this;
103         }
104 
105         /**
106          * Populates the Thread Context Map with the supplied key/value pairs. Any existing keys in the
107          * {@link ThreadContext} will be replaced with the supplied values, and restored back to their original values when
108          * the instance is closed.
109          *
110          * @param key   The  key to be added
111          * @param value The value to be added
112          * @return the instance that will back out the changes when closed.
113          */
114         public Instance put(final String key, final String value) {
115             // If there are no existing values, a null will be stored as an old value
116             if (!originalValues.containsKey(key)) {
117                 originalValues.put(key, ThreadContext.get(key));
118             }
119             ThreadContext.put(key, value);
120             return this;
121         }
122 
123         /**
124          * Removes the values from the {@link ThreadContext}.
125          * <p>
126          * Values pushed to the {@link ThreadContext} <em>stack</em> will be popped off.
127          * </p>
128          * <p>
129          * Values put on the {@link ThreadContext} <em>map</em> will be removed, or restored to their original values it they already existed.
130          * </p>
131          */
132         @Override
133         public void close() {
134             closeStack();
135             closeMap();
136         }
137 
138         private void closeMap() {
139             for (Iterator<Map.Entry<String, String>> it = originalValues.entrySet().iterator(); it.hasNext(); ) {
140                 final Map.Entry<String, String> entry = it.next();
141                 final String key = entry.getKey();
142                 final String originalValue = entry.getValue();
143                 if (null == originalValue) {
144                     ThreadContext.remove(key);
145                 } else {
146                     ThreadContext.put(key, originalValue);
147                 }
148                 it.remove();
149             }
150         }
151 
152         private void closeStack() {
153             for (int i = 0; i < pushCount; i++) {
154                 ThreadContext.pop();
155             }
156             pushCount = 0;
157         }
158     }
159 }