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}