com.alibaba.mtc.MtContextThreadLocal Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of multithread.context Show documentation
a simple lib for transmitting context between thread even using thread pool.
There is a newer version: 1.2.1
package com.alibaba.mtc;

import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
import java.util.WeakHashMap;

/**
 * {@link MtContextThreadLocal} can transmit context from the thread of submitting task to the thread of executing task.
 * 
 * Note: this class extends {@link java.lang.InheritableThreadLocal},
 * so {@link com.alibaba.mtc.MtContextThreadLocal} first is a {@link java.lang.InheritableThreadLocal}.
 *
 * @author ding.lid
 * @see MtContextRunnable
 * @see MtContextCallable
 * @since 0.10.0
 */
public class MtContextThreadLocal extends InheritableThreadLocal {
    /**
     * Computes the context value for this multi-thread context variable
     * as a function of the source thread's value at the time the task
     * Object is created.  This method is called from {@link com.alibaba.mtc.MtContextRunnable} or
     * {@link com.alibaba.mtc.MtContextCallable} when it create, before the task is started.
     * 

     * This method merely returns reference of its source thread value, and should be overridden
     * if a different behavior is desired.
     *
     * @since 1.0.0
     */
    protected T copy(T parentValue) {
        return parentValue;
    }

    @Override
    public final T get() {
        T value = super.get();
        if (null != value) {
            addMtContextThreadLocal();
        }
        return value;
    }

    @Override
    public final void set(T value) {
        super.set(value);
        if (null == value) { // may set null to remove value
            removeMtContextThreadLocal();
        } else {
            addMtContextThreadLocal();
        }
    }

    @Override
    public final void remove() {
        removeMtContextThreadLocal();
        super.remove();
    }

    private void superRemove() {
        super.remove();
    }

    T copyMtContextValue() {
        return copy(get());
    }

    static ThreadLocal, ?>> holder =
            new ThreadLocal, ?>>() {
                @Override
                protected Map, ?> initialValue() {
                    return new WeakHashMap, Object>();
                }
            };

    void addMtContextThreadLocal() {
        if (!holder.get().containsKey(this)) {
            holder.get().put(this, null); // WeakHashMap supports null value.
        }
    }

    void removeMtContextThreadLocal() {
        holder.get().remove(this);
    }

    static Map, Object> copy() {
        Map, Object> copy = new HashMap, Object>();
        for (MtContextThreadLocal threadLocal : holder.get().keySet()) {
            copy.put(threadLocal, threadLocal.copyMtContextValue());
        }
        return copy;
    }

    static Map, Object> backupAndSet(Map, Object> copied) {
        Map, Object> backup = new HashMap, Object>();
        for (Iterator, ?>> iterator = holder.get().entrySet().iterator();
             iterator.hasNext(); ) {
            Map.Entry, ?> next = iterator.next();
            MtContextThreadLocal threadLocal = next.getKey();
            // backup 
            backup.put(threadLocal, threadLocal.get());
            // clean extra MtContext in destination thread
            if (!copied.containsKey(threadLocal)) {
                iterator.remove();
                threadLocal.superRemove();
            }
        }
        setMtContexts(copied);
        return backup;
    }

    static void restore(Map, Object> backup) {
        for (Iterator, ?>> iterator = holder.get().entrySet().iterator();
             iterator.hasNext(); ) {
            Map.Entry, ?> next = iterator.next();
            MtContextThreadLocal threadLocal = next.getKey();
            // clean extra MtContext
            if (!backup.containsKey(threadLocal)) {
                iterator.remove();
                threadLocal.superRemove();
            }
        }
        setMtContexts(backup);
    }

    static void setMtContexts(Map, Object> set) {
        for (Map.Entry, Object> entry : set.entrySet()) {
            @SuppressWarnings("unchecked")
            MtContextThreadLocal