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没有这个问题。