io.pcp.parfait.timing.ThreadContext Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of parfait-core Show documentation
Java performance monitoring framework, including PCP bridge
The newest version!
/*
 * Copyright 2009-2017 Aconex
 *
 * Licensed under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License. You may obtain a copy of the License at:
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
 * implied.  See the License for the specific language governing
 * permissions and limitations under the License.
 */

package io.pcp.parfait.timing;

import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

import com.google.common.cache.CacheBuilder;
import com.google.common.cache.CacheLoader;
import com.google.common.cache.LoadingCache;
import net.jcip.annotations.ThreadSafe;

/**
 * 
 * Map-like functions to keep track of key/value pairs for application threads.
 * Keys are Strings, with values of any arbitrary object. Optionally keeps a log
 * framework's MDC weakly in sync with changes -- that is, there is no atomicity
 * guarantee so it's plausible that this class' context information and the
 * logger's will not be in a consistent state; however, log4j and logback's MDC
 * expose data only to the current thread via a thread-local so this is
 * unlikely to be a problem in practice.
 * 
 * 
 * Most methods operate on the context of the calling thread; only
 * {@link #forThread(Thread)} allows cross-thread information retrieval.
 * 
 * 
 * @author Cowan
 */
@ThreadSafe
public class ThreadContext {
    private static final CacheLoader> NEW_CONTEXT_CREATOR = new CacheLoader>() {
        @Override
        public Map load(Thread thread) throws Exception {
            return new ConcurrentHashMap<>();
        }
    };
    private final LoadingCache> PER_THREAD_CONTEXTS = CacheBuilder.newBuilder().weakKeys().build(NEW_CONTEXT_CREATOR);

    private volatile MdcBridge mdcBridge = new NullMdcBridge();

    public ThreadContext() {
        this(new NullMdcBridge());
    }

    public ThreadContext(MdcBridge mdcBridge) {
        // TODO should that be a static variable..?
        this.mdcBridge=mdcBridge;
    }

    /**
     * Adds the given key/value pair to the current thread's context, and updates {@link MdcBridge} with
     * same.
     */
    public void put(String key, Object value) {
        PER_THREAD_CONTEXTS.getUnchecked(Thread.currentThread()).put(key, value);
        mdcBridge.put(key, value);
    }

    /**
     * Removes the given key from the current thread's context and {@link MdcBridge}.
     */
    public void remove(String key) {
        PER_THREAD_CONTEXTS.getUnchecked(Thread.currentThread()).remove(key);
        mdcBridge.remove(key);
    }

    /**
     * Retrieves the value corresponding to the supplied key for the current thread (null if no such
     * value exists)
     */
    public Object get(String key) {
        return PER_THREAD_CONTEXTS.getUnchecked(Thread.currentThread()).get(key);
    }

    /**
     * Clears all values for the current thread.
     */
    public void clear() {
        for (String key : allKeys()) {
           mdcBridge.remove(key);
        }

        PER_THREAD_CONTEXTS.getUnchecked(Thread.currentThread()).clear();
    }

    /**
     * Retrieves a copy of the thread context for the given thread
     */
    public Map forThread(Thread t) {
        return new HashMap(PER_THREAD_CONTEXTS.getUnchecked(t));
    }

    public Collection allKeys() {
        Set keys = new HashSet();
        for (Map threadMdc : PER_THREAD_CONTEXTS.asMap().values()) {
            keys.addAll(threadMdc.keySet());
        }
        return keys;
    }

    public Object getForThread(Thread thread, String key) {
        return PER_THREAD_CONTEXTS.getUnchecked(thread).get(key);
    }

    /**
     * Factory methods that create a new ThreadContext initialised to also update SLF4J's MDC
     */
    public static ThreadContext newMDCEnabledContext() {
        return newSLF4JEnabledContext();
    }

    public static ThreadContext newSLF4JEnabledContext() {
        return new ThreadContext(new Slf4jMDCBridge());
    }

    public interface MdcBridge {
		void put(String key, Object object);

		void remove(String key);
    }
    
    public static class NullMdcBridge implements MdcBridge {
		@Override
		public void put(String key, Object object) {
			// no-op
		}

		@Override
		public void remove(String key) {
			// no-op
		}
    }
    
    public static class Slf4jMDCBridge implements MdcBridge {
		@Override
		public void put(String key, Object object) {
			org.slf4j.MDC.put(key, String.valueOf(object));
		}

		@Override
		public void remove(String key) {
			org.slf4j.MDC.remove(key);
		}
	}

}