Dubbo中InternalThreadLocal源碼分析

在Java中，ThreadLocal是實現線程安全的一種手段，它的作用是對於同一個ThreadLocal變量，在每一個線程中都有一個副本，當修改任何一個線程的變量時，不會影響到其他線程。它通過在每一個Thread中存儲一個類似於map的結構，以ThreadLocal變量爲key，變量值爲value。
Dubbo在RPC調用的上下文中，需要藉助ThreadLocal保存上下文。它從Netty中借鑑了InternalThreadLocal的實現，與Java官方的ThreadLocal不同的是，它內部沒有采用哈希表的結構，而是採用了數組作爲內部實現。下面就來看一下具體是如何實現的。
InternalThreadLocal模塊主要有以下幾個類：

使用方式

參見RpcContext類中，定義了

	private static final InternalThreadLocal<RpcContext> LOCAL = new InternalThreadLocal<RpcContext>() {
        @Override
        protected RpcContext initialValue() {
            return new RpcContext();
        }
    };

獲取線程本地變量，只需調用

	public static RpcContext getContext() {
        return LOCAL.get();
    }

重置

	public static void restoreContext(RpcContext oldContext) {
        LOCAL.set(oldContext);
    }

移除

	public static void removeContext() {
        LOCAL.remove();
    }

可以看出使用方式和官方的ThreadLocal是很類似的。

InternalThread

首先是InternalThread這個類，看它的類頭和主要的field

/**
 * InternalThread
 */
public class InternalThread extends Thread {

    private InternalThreadLocalMap threadLocalMap;

以及獲取和設置這個InternalThreadLocalMap變量的方法

	/**
     * Returns the internal data structure that keeps the threadLocal variables bound to this thread.
     * Note that this method is for internal use only, and thus is subject to change at any time.
     */
    public final InternalThreadLocalMap threadLocalMap() {
        return threadLocalMap;
    }

    /**
     * Sets the internal data structure that keeps the threadLocal variables bound to this thread.
     * Note that this method is for internal use only, and thus is subject to change at any time.
     */
    public final void setThreadLocalMap(InternalThreadLocalMap threadLocalMap) {
        this.threadLocalMap = threadLocalMap;
    }

從這個類可以看出，它跟官方的ThreaLocal的實現方式很像，也是通過一個內部的“map”來實現的。那麼這個“InternalThreadLocalMap”真的是一個map嗎？我們接着看。

remove

	public static void remove() {
	    // 獲取當前線程
        Thread thread = Thread.currentThread();
        // 若是InternalThread則設置爲null，否則調用ThreadLocal的remove方法
        if (thread instanceof InternalThread) {
            ((InternalThread) thread).setThreadLocalMap(null);
        } else {
            slowThreadLocalMap.remove();
        }
    }

destroy

	public static void destroy() {
        slowThreadLocalMap = null;
    }

nextVariableIndex

	public static int nextVariableIndex() {
		// 先返回值再加1，通過Atomic變量讓該操作爲原子操作
        int index = NEXT_INDEX.getAndIncrement();
        // 溢出處理
        if (index < 0) {
            NEXT_INDEX.decrementAndGet();
            throw new IllegalStateException("Too many thread-local indexed variables");
        }
        return index;
    }

這個方法就是在獲取變量時返回變量下標的方法，它通過一個AtomicInteger變量自增來獲得，若溢出，則拋異常

構造方法

	private InternalThreadLocalMap() {
        indexedVariables = newIndexedVariableTable();
    }

構造方法就是初始化了存儲數據的數組。

newIndexedVariableTable

	private static Object[] newIndexedVariableTable() {
        Object[] array = new Object[32];
        Arrays.fill(array, UNSET);
        return array;
    }

初始化數組爲32個元素大小，並用UNSET填充。

對InternalThreadLocalMap的類的介紹暫且到這裏。下面從功能的角度介紹一下InternalThreadLocal的源碼。看一下get方法：

get

InternalThreadLocal.get()

	public final V get() {
		// 獲取map
        InternalThreadLocalMap threadLocalMap = InternalThreadLocalMap.get();
        // 通過下標獲取變量值
        Object v = threadLocalMap.indexedVariable(index);
        // 若不爲UNSET則直接返回
        if (v != InternalThreadLocalMap.UNSET) {
            return (V) v;
        }
        // 初始化
        return initialize(threadLocalMap);
    }

邏輯很清晰

InternalThreadLocalMap.get()

	public static InternalThreadLocalMap get() {
		// 獲取當前線程
        Thread thread = Thread.currentThread();
        // 如果是InternalThread線程，則調用fastGet獲取
        if (thread instanceof InternalThread) {
            return fastGet((InternalThread) thread);
        }
        // 如果不是InternalThread，則調用slowGet獲取
        return slowGet();
    }

區分了當前線程是否是InternalThread來進行獲取

fastGet

	private static InternalThreadLocalMap fastGet(InternalThread thread) {
		// 獲取InternalThreadLocalMap變量
        InternalThreadLocalMap threadLocalMap = thread.threadLocalMap();
        // 如果爲null，則初始化這個map
        if (threadLocalMap == null) {
            thread.setThreadLocalMap(threadLocalMap = new InternalThreadLocalMap());
        }
        // 返回結果
        return threadLocalMap;
    }

直接從InternalThread中獲取InternalThreadLocalMap的field，若爲空則進行初始化。

slowGet

	private static InternalThreadLocalMap slowGet() {
		// 獲取存儲InternalThreadLocalMap的ThreadLocal變量
        ThreadLocal<InternalThreadLocalMap> slowThreadLocalMap = InternalThreadLocalMap.slowThreadLocalMap;
        // 從ThreadLocal中獲取InternalThreadLocalMap
        InternalThreadLocalMap ret = slowThreadLocalMap.get();
        // 若爲null則初始化並設置到ThreadLocal中
        if (ret == null) {
            ret = new InternalThreadLocalMap();
            slowThreadLocalMap.set(ret);
        }
        return ret;
    }

如果不是InternalThread，則需要通過官方的ThreadLocal獲取，需要進行哈希等計算，所以性能較差。

以上是獲取InternalThreadLocalMap的方式，再回到InternalThreadLocal的get方法，如果獲取到的map不爲空，則會調用
InternalThreadLocalMap.indexedVariable(int index)方法獲取數據。注意這裏傳入的index，它就是當前線程在實際保存數據的數組中的下標。它是什麼時候初始化的呢？

	public InternalThreadLocal() {
        index = InternalThreadLocalMap.nextVariableIndex();
    }

答案是構造方法中，通過調用InternalThreadLocalMap的nextVariableIndex方法進行初始化。因爲InternalThreadLocalMap是所有線程公用的。所以只要保證這個方法是原子的，就能保證不同線程獲取到不同的下標，因此也就不會衝突。

InternalThreadLocalMap.nextVariableIndex()

	public static int nextVariableIndex() {
		// 先返回值再加1，通過Atomic變量讓該操作爲原子操作
        int index = NEXT_INDEX.getAndIncrement();
        // 溢出處理
        if (index < 0) {
            NEXT_INDEX.decrementAndGet();
            throw new IllegalStateException("Too many thread-local indexed variables");
        }
        return index;
    }

這個方法就是在獲取變量時返回變量下標的方法，它通過一個AtomicInteger變量自增來獲得，若溢出，則拋異常

回到

InternalThreadLocalMap.indexedVariable(int index)

	public Object indexedVariable(int index) {
        Object[] lookup = indexedVariables;
        return index < lookup.length ? lookup[index] : UNSET;
    }

indexedVariables是一個field，用於保存實際的數據。對於超過下標範圍的index，返回UNSET。

再次回到InternalThreadLocal.get()方法，如果獲取的數據爲UNSET，則需要初始化。

InternalThreadLocal.initialize()

	private V initialize(InternalThreadLocalMap threadLocalMap) {
        V v = null;
        try {
        	// 調用初始化方法
            v = initialValue();
        } catch (Exception e) {
            throw new RuntimeException(e);
        }

		// 設置初始值到指定位置
        threadLocalMap.setIndexedVariable(index, v);
        addToVariablesToRemove(threadLocalMap, this);
        return v;
    }

默認的初始化變量爲null，子類可繼承這個方法。

InternalThreadLocalMap.setIndexedVariable(int index, Object value)

	public boolean setIndexedVariable(int index, Object value) {
        Object[] lookup = indexedVariables;
        // index未超限則直接獲取
        if (index < lookup.length) {
            Object oldValue = lookup[index];
            lookup[index] = value;
            // 當且僅當一個新的線程本地變量被創建才返回true
            return oldValue == UNSET;
        } else {
        	// index超限則擴容數組，並設置值
            expandIndexedVariableTableAndSet(index, value);
            return true;
        }
    }

InternalThreadLocalMap.expandIndexedVariableTableAndSet

	private void expandIndexedVariableTableAndSet(int index, Object value) {
        Object[] oldArray = indexedVariables;
        final int oldCapacity = oldArray.length;
        int newCapacity = index;
        newCapacity |= newCapacity >>> 1;
        newCapacity |= newCapacity >>> 2;
        newCapacity |= newCapacity >>> 4;
        newCapacity |= newCapacity >>> 8;
        newCapacity |= newCapacity >>> 16;
        newCapacity++;

        Object[] newArray = Arrays.copyOf(oldArray, newCapacity);
        Arrays.fill(newArray, oldCapacity, newArray.length, UNSET);
        newArray[index] = value;
        indexedVariables = newArray;
    }

這個擴展數組實際上就是將數組大小擴展到大於等於index的2的最小次冪。之後先把擴展出來的位置全設置爲UNSET，之後再設置index位置爲value。最後將indexedVariables指向新的數組。

addToVariablesToRemove這個方法後面說remove的時候再看

set

以上就是InternalThreadLocal的get方法，下面看一下set方法

InternalThreadLocal.set(V value)

	/**
     * Sets the value for the current thread.
     */
    public final void set(V value) {
        if (value == null || value == InternalThreadLocalMap.UNSET) {
            remove();
        } else {
            InternalThreadLocalMap threadLocalMap = InternalThreadLocalMap.get();
            if (threadLocalMap.setIndexedVariable(index, value)) {
                addToVariablesToRemove(threadLocalMap, this);
            }
        }
    }

InternalThreadLocalMap.setIndexedVariable(int index, Object value)方法上面已經介紹過，不再贅述。

下面看一下remove。我們知道，有時候我們是需要清除InternalThreadLocal中的變量的，比如某一次RPC調用結束以後。那麼是如何清除的呢，首先就看一下上面略過的

InternalThreadLocal.addToVariablesToRemove(InternalThreadLocalMap threadLocalMap, InternalThreadLocal<?> variable)方法

	private static void addToVariablesToRemove(InternalThreadLocalMap threadLocalMap, InternalThreadLocal<?> variable) {
		// 獲取保存待移除變量的集合
        Object v = threadLocalMap.indexedVariable(VARIABLES_TO_REMOVE_INDEX);
        Set<InternalThreadLocal<?>> variablesToRemove;
        // 如果還沒有設置，則初始化這個集合
        if (v == InternalThreadLocalMap.UNSET || v == null) {
        	// 使用IdentityHashMap是說比較兩個元素是否相等，通過==而不是equals來比較
            variablesToRemove = Collections.newSetFromMap(new IdentityHashMap<InternalThreadLocal<?>, Boolean>());
            
            threadLocalMap.setIndexedVariable(VARIABLES_TO_REMOVE_INDEX, variablesToRemove);
        } else {
            variablesToRemove = (Set<InternalThreadLocal<?>>) v;
        }

		// 添加到集合中
        variablesToRemove.add(variable);
    }

就是從保存數據的數組中拿出一個位置來保存這個待移除變量的集合。這個方法在intialize或者set值的時候都會調用。

再看

InternalThreadLocal.remove()

	/**
     * Sets the value to uninitialized; a proceeding call to get() will trigger a call to initialValue().
     */
	public final void remove() {
        remove(InternalThreadLocalMap.getIfSet());
    }

這個方法要將對應位置的value設置爲未初始化狀態，注意這裏調用的是getIfSet方法，也就是說如果map沒初始化，返回的會是null

InternalThreadLocal.remove(InternalThreadLocalMap threadLocalMap)

	/**
     * Sets the value to uninitialized for the specified thread local map;
     * a proceeding call to get() will trigger a call to initialValue().
     * The specified thread local map must be for the current thread.
     */
    @SuppressWarnings("unchecked")
    public final void remove(InternalThreadLocalMap threadLocalMap) {
        if (threadLocalMap == null) {
            return;
        }
        // 刪除數組中對應位置的數據
        Object v = threadLocalMap.removeIndexedVariable(index);
        // 從待刪除數據集合中刪除
        removeFromVariablesToRemove(threadLocalMap, this);

		// 若刪除的數據不是UNSET，則進行一些額外的處理
        if (v != InternalThreadLocalMap.UNSET) {
            try {
            	// 擴展點
                onRemoval((V) v);
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        }
    }

InternalThreadLocal.removeFromVariablesToRemove(InternalThreadLocalMap threadLocalMap, InternalThreadLocal<?> variable)

	private static void removeFromVariablesToRemove(InternalThreadLocalMap threadLocalMap, InternalThreadLocal<?> variable) {
		// 通過特定下標獲取集合
        Object v = threadLocalMap.indexedVariable(VARIABLES_TO_REMOVE_INDEX);

		// 若未初始化直接返回
        if (v == InternalThreadLocalMap.UNSET || v == null) {
            return;
        }
        
        Set<InternalThreadLocal<?>> variablesToRemove = (Set<InternalThreadLocal<?>>) v;
        // 從集合中刪除該InternalThreadLocal變量
        variablesToRemove.remove(variable);
    }

InternalThreadLocal.removeAll()

	/**
     * Removes all {@link InternalThreadLocal} variables bound to the current thread.  This operation is useful when you
     * are in a container environment, and you don't want to leave the thread local variables in the threads you do not
     * manage.
     */
    @SuppressWarnings("unchecked")
    public static void removeAll() {
        InternalThreadLocalMap threadLocalMap = InternalThreadLocalMap.getIfSet();
        if (threadLocalMap == null) {
            return;
        }

        try {
        	// 獲取集合
            Object v = threadLocalMap.indexedVariable(VARIABLES_TO_REMOVE_INDEX);
            if (v != null && v != InternalThreadLocalMap.UNSET) {
                Set<InternalThreadLocal<?>> variablesToRemove = (Set<InternalThreadLocal<?>>) v;
                // 將集合轉化爲數組，提高遍歷效率
                InternalThreadLocal<?>[] variablesToRemoveArray =
                        variablesToRemove.toArray(new InternalThreadLocal[variablesToRemove.size()]);
                for (InternalThreadLocal<?> tlv : variablesToRemoveArray) {
                    tlv.remove(threadLocalMap);
                }
            }
        } finally {
            InternalThreadLocalMap.remove();
        }
    }

根據註釋，刪除所有和當前線程綁定的InternalThreadLocal，這個操作在某些容器環境中是有用的，因爲我們不想將這些變量留在線程中，當我們不再擁有這些線程的控制權的時候。（實際上這句話的意思就是說，容器中的線程池的線程是複用的，當我的某一次請求結束之後，如果不清理掉這些線程本地變量，就會污染下一次分配到使用這個線程的請求）

這裏爲什麼要使用一個集合來保存這個線程的所有ThreadLocal變量呢？主要是出於性能考慮。我們知道我們初始化的數組，可能只有很少一部分被使用了，如果沒有這個集合，那麼要清理就只能遍歷這個數組，那樣是非常划不來的。

什麼情況下使用InternalThreadLocal

從代碼分析中可以看出，只有當前線程是InternalThread的時候，使用InternalThreadLocal纔是合適的。因爲如果不是InternalThread，那麼要先從Java官方的ThreadLocal去獲取InternalThreadLocal，又多了一層基於哈希表的封裝，性能自然會慢。

使用InternalThreadLocal主要是出於性能考慮，因爲每一個InternalThreadLocal變量在初始化的時候index就已經確定了，獲取或設置變量時可以直接使用，無需進行哈希運算。此外，ThreadLocal的Entry的key（ThreadLocal類型）使用了弱引用，以避免內存泄漏（也許這裏你會問，爲什麼不把ThreadLocalMap中的Entry整個搞成弱引用，這樣做是爲了易用性，對ThreadLocal的使用者只暴露了ThreadLocal，而沒有職級暴露Entry，所以結合弱引用被回收的規律，不能把整個Entry設置爲弱引用），但是Entry中的value不是弱引用，還是存在內存泄漏的風險。而InternalThread沒有這個問題。