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. 026 * 027 * @since 2.6 028 */ 029public class CloseableThreadContext { 030 031 private CloseableThreadContext() { 032 } 033 034 /** 035 * Pushes new diagnostic context information on to the Thread Context Stack. The information will be popped off when 036 * the instance is closed. 037 * 038 * @param message The new diagnostic context information. 039 * @return a new instance that will back out the changes when closed. 040 */ 041 public static CloseableThreadContext.Instance push(final String message) { 042 return new CloseableThreadContext.Instance().push(message); 043 } 044 045 /** 046 * Pushes new diagnostic context information on to the Thread Context Stack. The information will be popped off when 047 * the instance is closed. 048 * 049 * @param message The new diagnostic context information. 050 * @param args Parameters for the message. 051 * @return a new instance that will back out the changes when closed. 052 */ 053 public static CloseableThreadContext.Instance push(final String message, final Object... args) { 054 return new CloseableThreadContext.Instance().push(message, args); 055 } 056 057 /** 058 * Populates the Thread Context Map with the supplied key/value pairs. Any existing keys in the 059 * {@link ThreadContext} will be replaced with the supplied values, and restored back to their original values when 060 * the instance is closed. 061 * 062 * @param key The key to be added 063 * @param value The value to be added 064 * @return a new instance that will back out the changes when closed. 065 */ 066 public static CloseableThreadContext.Instance put(final String key, final String value) { 067 return new CloseableThreadContext.Instance().put(key, value); 068 } 069 070 public static class Instance implements AutoCloseable { 071 072 private int pushCount = 0; 073 private final Map<String, String> originalValues = new HashMap<>(); 074 075 private Instance() { 076 } 077 078 /** 079 * Pushes new diagnostic context information on to the Thread Context Stack. The information will be popped off when 080 * the instance is closed. 081 * 082 * @param message The new diagnostic context information. 083 * @return the instance that will back out the changes when closed. 084 */ 085 public Instance push(final String message) { 086 ThreadContext.push(message); 087 pushCount++; 088 return this; 089 } 090 091 /** 092 * Pushes new diagnostic context information on to the Thread Context Stack. The information will be popped off when 093 * the instance is closed. 094 * 095 * @param message The new diagnostic context information. 096 * @param args Parameters for the message. 097 * @return the instance that will back out the changes when closed. 098 */ 099 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}