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}