001/* 002 * Licensed to the Apache Software Foundation (ASF) under one or more 003 * contributor license agreements. See the NOTICE file distributed with 004 * this work for additional information regarding copyright ownership. 005 * The ASF licenses this file to You under the Apache license, Version 2.0 006 * (the "License"); you may not use this file except in compliance with 007 * the License. You may obtain a copy of the License at 008 * 009 * http://www.apache.org/licenses/LICENSE-2.0 010 * 011 * Unless required by applicable law or agreed to in writing, software 012 * distributed under the License is distributed on an "AS IS" BASIS, 013 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 014 * See the license for the specific language governing permissions and 015 * limitations under the license. 016 */ 017package org.apache.logging.log4j; 018 019import java.util.HashMap; 020import java.util.Iterator; 021import java.util.Map; 022 023/** 024 * Adds entries to the {@link ThreadContext stack or map} and them removes them when the object is closed, e.g. as part 025 * of a try-with-resources. User code can now look like this: 026 * <pre> 027 * try (CloseableThreadContext.put(key1, value1).put(key2, value2)) { 028 * callSomeMethodThatLogsALot(); 029 * 030 * // key1 and key2 are automatically removed from the ThreadContext map when done 031 * } 032 * </pre> 033 * 034 * @since 2.6 035 */ 036public class CloseableThreadContext { 037 038 private CloseableThreadContext() { 039 } 040 041 /** 042 * Pushes new diagnostic context information on to the Thread Context Stack. The information will be popped off when 043 * the instance is closed. 044 * 045 * @param message The new diagnostic context information. 046 * @return a new instance that will back out the changes when closed. 047 */ 048 public static CloseableThreadContext.Instance push(final String message) { 049 return new CloseableThreadContext.Instance().push(message); 050 } 051 052 /** 053 * Pushes new diagnostic context information on to the Thread Context Stack. The information will be popped off when 054 * the instance is closed. 055 * 056 * @param message The new diagnostic context information. 057 * @param args Parameters for the message. 058 * @return a new instance that will back out the changes when closed. 059 */ 060 public static CloseableThreadContext.Instance push(final String message, final Object... args) { 061 return new CloseableThreadContext.Instance().push(message, args); 062 } 063 064 /** 065 * Populates the Thread Context Map with the supplied key/value pairs. Any existing keys in the 066 * {@link ThreadContext} will be replaced with the supplied values, and restored back to their original values when 067 * the instance is closed. 068 * 069 * @param key The key to be added 070 * @param value The value to be added 071 * @return a new instance that will back out the changes when closed. 072 */ 073 public static CloseableThreadContext.Instance put(final String key, final String value) { 074 return new CloseableThreadContext.Instance().put(key, value); 075 } 076 077 public static class Instance implements AutoCloseable { 078 079 private int pushCount = 0; 080 private final Map<String, String> originalValues = new HashMap<>(); 081 082 private Instance() { 083 } 084 085 /** 086 * Pushes new diagnostic context information on to the Thread Context Stack. The information will be popped off when 087 * the instance is closed. 088 * 089 * @param message The new diagnostic context information. 090 * @return the instance that will back out the changes when closed. 091 */ 092 public Instance push(final String message) { 093 ThreadContext.push(message); 094 pushCount++; 095 return this; 096 } 097 098 /** 099 * 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}