TransmittableThreadLocal的原理分析

一、簡介

　　TransmittableThreadLocal是由阿里開發的一個線程變量傳遞工具包，解決了InheritableThreadLocal只能再new Thread的時候傳遞本地變量，無法應用到線程池的問題。可以應用來作鏈路追蹤，傳遞變量等用途，下面我們來了解一下原理。

二、InheritableThreadLocal

    public class InheritableThreadLocal<T> extends ThreadLocal<T> {

    protected T childValue(T parentValue) {
        return parentValue;
    }

    ThreadLocalMap getMap(Thread t) {
       return t.inheritableThreadLocals;
    }

    // set的時候調用的
    void createMap(Thread t, T firstValue) {
        t.inheritableThreadLocals = new ThreadLocalMap(this, firstValue);
    }
}

InheritableThreadLocal重寫了父類的三個方法，其中createMap最為關鍵，他是我們調用ThreadLocal的Tset方法時，調用的。這里是new了一個ThreadLocalMap賦值給了Thread的inheritableThreadLocals變量，那么我就來看一下Thread的屬性及方法。

class Thread implements Runnable {

    // 父線程使用
    ThreadLocal.ThreadLocalMap threadLocals = null;

    // 子線程使用
    ThreadLocal.ThreadLocalMap inheritableThreadLocals = null;
    
    // 構造器 創建線程
    public Thread(Runnable target) {
        init(null, target, "Thread-" + nextThreadNum(), 0);
    }

}

當我們使用ThreadLocal的時候，是賦值給threadLocals屬性，使用InheritableThreadLocal就是把值又賦給了線程的inheritableThreadLocals屬性，那么，可以猜測就是在我們new Thread()的時候觸發了線程的值傳遞，下面通過源碼驗證猜想

    private void init(ThreadGroup g, Runnable target, String name,
                      long stackSize, AccessControlContext acc,
                      boolean inheritThreadLocals) {
        ···
        
        // 獲取父線程
        Thread parent = currentThread();
        
        ···
        
        if (inheritThreadLocals && parent.inheritableThreadLocals != null)
            // 當inheritThreadLocals為true，並且父線程的inheritableThreadLocals不為null
            // 將父線程的inheritableThreadLocals賦值給子線程的inheritableThreadLocals
            this.inheritableThreadLocals =
                ThreadLocal.createInheritedMap(parent.inheritableThreadLocals);
        ···
    }

三、TransmittableThreadLocal

　　TransmittableThreadLocal繼承自InheritableThreadLocal，因此它可以在創建線程的時候將值傳遞給子線程，那么怎么確保使用線程池的時候也有效呢？我們來看一下源碼

1、構造方法

    // 構造器
    public TransmittableThreadLocal() {
        this(false);
    }
    public TransmittableThreadLocal(boolean disableIgnoreNullValueSemantics) {
        // 是否忽略null值set，默認false
        this.disableIgnoreNullValueSemantics = disableIgnoreNullValueSemantics;
    }

2、set方法

    public final void set(T value) {
        if (!disableIgnoreNullValueSemantics && null == value) {
            // 不忽略null寫入，則移除本地線程變量
            remove();
        } else {
            // 調用父類InheritableThreadLocal的set方法
            super.set(value);
            // 將自己添加到靜態線程變量holder中
            addThisToHolder();
        }
    }

先看addThisToHolder方法

    private void addThisToHolder() {
        // 判斷holder是否存在此TransmittableThreadLocal對象
        if (!holder.get().containsKey(this)) {
            // 不存則添加進holder
            holder.get().put((TransmittableThreadLocal<Object>) this, null); // WeakHashMap supports null value.
        }
    }

屬性holder又是什么呢？

    private static final InheritableThreadLocal<WeakHashMap<TransmittableThreadLocal<Object>, ?>> holder =
            new InheritableThreadLocal<WeakHashMap<TransmittableThreadLocal<Object>, ?>>() {
                @Override
                protected WeakHashMap<TransmittableThreadLocal<Object>, ?> initialValue() {
                    return new WeakHashMap<TransmittableThreadLocal<Object>, Object>();
                }

                @Override
                protected WeakHashMap<TransmittableThreadLocal<Object>, ?> childValue(WeakHashMap<TransmittableThreadLocal<Object>, ?> parentValue) {
                    return new WeakHashMap<TransmittableThreadLocal<Object>, Object>(parentValue);
                }
            };

1、final static修飾的變量，只會存在一份

2、使用了WeakHashMap，弱引用，方便垃圾回收

3、key就是TransmittableThreadLocal對象

remove方法

    public final void remove() {
        // 從holder中移除
        removeThisFromHolder();
        // 調用父類的移除方法，移除值
        super.remove();
    }

3、get方法

    public final T get() {
        // 調用父類的get
        T value = super.get();
        // 如果允許忽略null，或者value不為null，再次添加到holder
        if (disableIgnoreNullValueSemantics || null != value) addThisToHolder();
        return value;
    }

4、當我們使用線程池時，需要使用TtlRunnable.get(runnable)對runnable進行包裝，或者使用TtlExecutors.getTtlExecutor(executor)對執行器進行包裝，才能使線程池的變量傳遞起效果，那么我們就接着看一下源碼的執行流程

TtlExecutors.getTtlExecutor(executor)

    public static Executor getTtlExecutor(@Nullable Executor executor) {
        if (TtlAgent.isTtlAgentLoaded() || null == executor || executor instanceof TtlEnhanced) {
            return executor;
        }
        // 包裝執行器
        return new ExecutorTtlWrapper(executor, true);
    } 
    ExecutorTtlWrapper(@NonNull Executor executor, boolean idempotent) {
        this.executor = executor;
        this.idempotent = idempotent;
    }
    public void execute(@NonNull Runnable command) {
        // 實際上也是通過TtlRunnable對原runnable進行包裝
        executor.execute(TtlRunnable.get(command, false, idempotent));
    }

可以看到，兩種方式原理一樣，我們直接看TtlRunnable.get()

    public static TtlRunnable get(@Nullable Runnable runnable, boolean releaseTtlValueReferenceAfterRun, boolean idempotent) {
        if (null == runnable) return null;

        if (runnable instanceof TtlEnhanced) {
            if (idempotent) return (TtlRunnable) runnable;
            else throw new IllegalStateException("Already TtlRunnable!");
        }
        // 返回TtlRunnable
        return new TtlRunnable(runnable, releaseTtlValueReferenceAfterRun);
    }

構建TtlRunnable

    private TtlRunnable(@NonNull Runnable runnable, boolean releaseTtlValueReferenceAfterRun) {
        // 原子引用
        this.capturedRef = new AtomicReference<Object>(capture());
        this.runnable = runnable;
        this.releaseTtlValueReferenceAfterRun = releaseTtlValueReferenceAfterRun;
    }

capture捕獲父線程的ttl

        // 存放父線程的值
        public static Object capture() {
            return new Snapshot(captureTtlValues(), captureThreadLocalValues());
        }
        
        private static HashMap<TransmittableThreadLocal<Object>, Object> captureTtlValues() {
            HashMap<TransmittableThreadLocal<Object>, Object> ttl2Value = new HashMap<TransmittableThreadLocal<Object>, Object>();
            // 遍歷了所有holder
            for (TransmittableThreadLocal<Object> threadLocal : holder.get().keySet()) {
                // copyValue實際上調用了TransmittableThreadLocal的get方法獲取線程存儲的變量值
                ttl2Value.put(threadLocal, threadLocal.copyValue());
            }
            return ttl2Value;
        }
    
        private static HashMap<ThreadLocal<Object>, Object> captureThreadLocalValues() {
            final HashMap<ThreadLocal<Object>, Object> threadLocal2Value = new HashMap<ThreadLocal<Object>, Object>();
            // 
            for (Map.Entry<ThreadLocal<Object>, TtlCopier<Object>> entry : threadLocalHolder.entrySet()) {
                final ThreadLocal<Object> threadLocal = entry.getKey();
                final TtlCopier<Object> copier = entry.getValue();
                //
                threadLocal2Value.put(threadLocal, copier.copy(threadLocal.get()));
            }
            return threadLocal2Value;
        }

再看TtlRunnable的run方法

    public void run() {
        // 獲取Snapshot對象，里面存儲了父線程的值
        final Object captured = capturedRef.get();
        if (captured == null || releaseTtlValueReferenceAfterRun && !capturedRef.compareAndSet(captured, null)) {
            throw new IllegalStateException("TTL value reference is released after run!");
        }
        // 傳入capture方法捕獲的ttl，然后在子線程重放，也就是調用ttl的set方法，
        // 這樣就會把值設置到當前的線程中去，最后會把子線程之前存在的ttl返回
        final Object backup = replay(captured);
        try {
            // 調用原runnable的run
            runnable.run();
        } finally {
            // 
            restore(backup);
        }
    }