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 }