Java ThreadLocal 实现

引言

前面我们已经介绍了 JDK 中常用的并发库(JUC)的使用方式, 本文我们着重介绍 JUC 中 ThreadLocal 的实现方式已经 Netty 对 ThreadLocal 的增强实现。所有关于 Java 并发的文章均收录于<Java并发系列文章>。

ThreadLocal实现

这里我们简单地介绍一下 ThreadLocal 的实现原理，在每个 Thread 对象中都保存了一个 ThreadLocal Map，其中 key 为 ThreadLocal 对象，Value 为 ThreadLocal 的值。这里 ThreadLocalMap 中的 Entry 使用了弱引用，是为了帮助 GC。当一个 ThreadLocal 对象不再被引用时，就会被 GC，这时候 ThreadLocalMap 中就会出现 key 为 null 的情况。但是因为 value 是强引用，所以如果 key 为 null 的数据不加管理的话，就会出现内存泄漏问题。ThreadLocalMap实现中已经考虑了这种情况，在调用 set()、get()、remove() 方法的时候，会清理掉 key 为 null 的记录。使用完 ThreadLocal方法后最好手动调用remove()方法，来帮助 GC。

public class Thread implements Runnable {
    /* ThreadLocal values pertaining to this thread. This map is maintained
     * by the ThreadLocal class. */
    ThreadLocal.ThreadLocalMap threadLocals = null;
}

static class ThreadLocalMap {

    /**
     * The entries in this hash map extend WeakReference, using
     * its main ref field as the key (which is always a
     * ThreadLocal object).  Note that null keys (i.e. entry.get()
     * == null) mean that the key is no longer referenced, so the
     * entry can be expunged from table.  Such entries are referred to
     * as "stale entries" in the code that follows.
     */
    static class Entry extends WeakReference<ThreadLocal<?>> {
        /** The value associated with this ThreadLocal. */
        Object value;

        Entry(ThreadLocal<?> k, Object v) {
            super(k);
            value = v;
        }
    }
}

如果一个对象只具有弱引用，那么下次 GC 时该对象就会被清理。
弱引用与软引用的区别在于：只含有软引用的对象只有在内存不足（即将发生OOM）时才会清除，而只含有弱引用的对象下次 GC 时就会清除无论内存是否紧张。

当我们想要获取 ThreadLocal 值的时候，会从当前 Thread 的 ThreadLocalMap 中查找，如果没有找到时，它会将初始值塞入该 Map 并返回。

// ThreadLocal.java
 /**
 * Returns the value in the current thread's copy of this
 * thread-local variable.  If the variable has no value for the
 * current thread, it is first initialized to the value returned
 * by an invocation of the {@link #initialValue} method.
 *
 * @return the current thread's value of this thread-local
 */
public T get() {
    Thread t = Thread.currentThread();
    ThreadLocalMap map = getMap(t);
    if (map != null) {
        ThreadLocalMap.Entry e = map.getEntry(this);
        if (e != null) {
            T result = (T)e.value;
            return result;
        }
    }
    return setInitialValue();
}

/**
 * Get the map associated with a ThreadLocal.
 *
 * @param  t the current thread
 * @return the map
 */
ThreadLocalMap getMap(Thread t) {
    return t.threadLocals;
}
/**
 * Returns the current thread's "initial value" for this
 * thread-local variable.  This method will be invoked the first
 * time a thread accesses the variable with the {@link #get}
 * method, unless the thread previously invoked the {@link #set}
 * method, in which case the {@code initialValue} method will not
 * be invoked for the thread.  Normally, this method is invoked at
 * most once per thread, but it may be invoked again in case of
 * subsequent invocations of {@link #remove} followed by {@link #get}.
 *
 * <p>This implementation simply returns {@code null}; if the
 * programmer desires thread-local variables to have an initial
 * value other than {@code null}, {@code ThreadLocal} must be
 * subclassed, and this method overridden.  Typically, an
 * anonymous inner class will be used.
 *
 * @return the initial value for this thread-local
 */
private T setInitialValue() {
    T value = initialValue();
    Thread t = Thread.currentThread();
    ThreadLocalMap map = getMap(t);
    if (map != null)
        map.set(this, value);
    else
        createMap(t, value);
    return value;
}

当我们进行 ThreadLocal 的设置时，道理相同，也是从 Thread 对象中获取 ThreadLocalMap 然后计算该 ThreadLocal 对象所处的槽位，并设置值。

/**
 * Sets the current thread's copy of this thread-local variable
 * to the specified value.  Most subclasses will have no need to
 * override this method, relying solely on the {@link #initialValue}
 * method to set the values of thread-locals.
 *
 * @param value the value to be stored in the current thread's copy of
 *        this thread-local.
 */
public void set(T value) {
    Thread t = Thread.currentThread();
    ThreadLocalMap map = getMap(t);
    if (map != null)
        map.set(this, value);
    else
        createMap(t, value);
}

/**
 * Set the value associated with key.
 *
 * @param key the thread local object
 * @param value the value to be set
 */
private void set(ThreadLocal<?> key, Object value) {

    // We don't use a fast path as with get() because it is at
    // least as common to use set() to create new entries as
    // it is to replace existing ones, in which case, a fast
    // path would fail more often than not.

    Entry[] tab = table;
    int len = tab.length;
    // 计算 hash 值对应的槽位
    int i = key.threadLocalHashCode & (len-1);
    for (Entry e = tab[i];
         e != null;
         e = tab[i = nextIndex(i, len)]) {
        ThreadLocal<?> k = e.get();
        // 如果该槽位存储的 ThreadLocal 对象就是自己，就返回
        if (k == key) {
            e.value = value;
            return;
        }
        // 如果遍历找到了一个空的槽位，就占用它
        if (k == null) {
            replaceStaleEntry(key, value, i);
            return;
        }
    }

    tab[i] = new Entry(key, value);
    int sz = ++size;
    // 清除 key 为 null的槽位，如果size过大就扩容，扩容阈值是哈希表长度的 2 / 3
    if (!cleanSomeSlots(i, sz) && sz >= threshold)
        rehash();
}

ThreadLocalMap内部使用一个数组来保存数据，类似HashMap；每个ThreadLocal在初始化的时候会分配一个threadLocalHashCode，然后和数组的长度进行取模来计算当前 ThreadLocal 变量所处的槽位。但是这样也会出现hash冲突的情况，在HashMap中处理冲突是使用链表+红黑树的方式。而在ThreadLocalMap中，我们可以看到它直接使用nextIndex，进行遍历操作，期间如果找到了自己之前使用到的槽位，就直接返回，否则占用一个没有被使用的槽位。很明显当 ThreadLocal 很多时这样效率很低。

获取 ThreadLocal 值的过程也类似，先通过 hash 找到槽位，如果该槽位保存的不是我们要的 ThreadLocal 对象，则进行遍历查找。

private Entry getEntry(ThreadLocal<?> key) {
    int i = key.threadLocalHashCode & (table.length - 1);
    Entry e = table[i];
    // 直接 hash 槽位找到了目标对象，直接返回
    if (e != null && e.get() == key)
        return e;
    else
        // 否则，遍历查找
        return getEntryAfterMiss(key, i, e);
}

private Entry getEntryAfterMiss(ThreadLocal<?> key, int i, Entry e) {
    Entry[] tab = table;
    int len = tab.length;

    while (e != null) {
        ThreadLocal<?> k = e.get();
        if (k == key)
            return e;
        if (k == null)
            // 这里如果发现了空的槽位，要进行重新 hash，来提升效率
            expungeStaleEntry(i);
        else
            i = nextIndex(i, len);
        e = tab[i];
    }
    return null;
}

private int expungeStaleEntry(int staleSlot) {
    Entry[] tab = table;
    int len = tab.length;

    // expunge entry at staleSlot
    tab[staleSlot].value = null;
    tab[staleSlot] = null;
    size--;

    // Rehash until we encounter null
    Entry e;
    int i;
    for (i = nextIndex(staleSlot, len);
         (e = tab[i]) != null;
         i = nextIndex(i, len)) {
        ThreadLocal<?> k = e.get();
        // 同样如果找到了 key 为 null 的槽位，就把他清空来帮助 GC
        if (k == null) {
            e.value = null;
            tab[i] = null;
            size--;
        } else {
            // 如果一个节点通过hash 计算的槽位，和实际保存的槽位不一样时，从计算所得的槽位出发，找到一个为 null 的槽位，并将该节点存进去
            int h = k.threadLocalHashCode & (len - 1);
            if (h != i) {
                tab[i] = null;
                // Unlike Knuth 6.4 Algorithm R, we must scan until
                // null because multiple entries could have been stale.
                while (tab[h] != null)
                    h = nextIndex(h, len);
                tab[h] = e;
            }
        }
    }
    return i;
}

FastThreadLocal

正是因为 JDK 提供的 ThreadLocal 存在性能问题，所以在 Netty 中，对 ThreadLocal 进行了改写，Netty 配套的提供了 FastThreadLocal ， FastThreadLocalThread ， FastThreadLocalRunnable ，其中 FastThreadLocalRunnable 比较简单，就是在原始 Runnable 接口之上的装饰者。达到了执行完毕后自动清除 FastThreadLocal 的效果。

final class FastThreadLocalRunnable implements Runnable {
    private final Runnable runnable;

    private FastThreadLocalRunnable(Runnable runnable) {
        this.runnable = ObjectUtil.checkNotNull(runnable, "runnable");
    }

    @Override
    public void run() {
        try {
            runnable.run();
        } finally {
            FastThreadLocal.removeAll();
        }
    }

    static Runnable wrap(Runnable runnable) {
        return runnable instanceof FastThreadLocalRunnable ? runnable : new FastThreadLocalRunnable(runnable);
    }
}

FastThreadLocal 必须配合 FastThreadLocalThread 一起使用，才能达到性能更优的效果，否则可能还不如直接使用 JDK 自身的 ThreadLocal。

FastThreadLocal 为什么能比 JDK 的实现更快呢，原因就在于 FastThreadLocal 不是使用 hash 表来保存 ThreadLocal 的值，而是直接使用了数组。让我们看看它的具体实现方案吧。

FastThreadLocal 所有的数据都保存在了 InternalThreadLocalMap 中, 这里我们先明白它是一个保存数据的容器就行，它是如何保存数据的我们后面介绍。那么，InternalThreadLocalMap 这个容器存在哪了呢？是不是像 JDK 提供的 ThreadLocal 一样把容器（JDK 中用到的Hash 表）存在了Thread 对象中呢？没错，如果我们使用的是 FastThreadLocalThread 的话，InternalThreadLocalMap 就是 FastThreadLocalThread 的一个成员变量。

public class FastThreadLocalThread extends Thread {
    // This will be set to true if we have a chance to wrap the Runnable.
    private final boolean cleanupFastThreadLocals;
    // 保存了所有的 FastThreadLocal 值
    private InternalThreadLocalMap threadLocalMap;
    // ...
}

而如果我们使用的不是 FastThreadLocalThread，而是 JDK 提供的 Thread 时，InternalThreadLocalMap 保存在 JDK 提供的 ThreadLocal 对象中。

下面给大家展示的是 UnpaddedInternalThreadLocalMap 它是 InternalThreadLocalMap 的父类，大部分重要的数据都是保存在 UnpaddedInternalThreadLocalMap 中的。

class UnpaddedInternalThreadLocalMap {
    // 当没有使用 FastThreadLocalThread 时，通过 JDK ThreadLocal 来保存
    static final ThreadLocal<InternalThreadLocalMap> slowThreadLocalMap = new ThreadLocal<InternalThreadLocalMap>();
    static final AtomicInteger nextIndex = new AtomicInteger();

    /** Used by {@link FastThreadLocal} */
    Object[] indexedVariables;
    // ...
}

在 UnpaddedInternalThreadLocalMap 中，通过 JDK 提供的 ThreadLocal 来保存一个 InternalThreadLocalMap，以应对没有使用 FastThreadLocalThread 的情况。

正因如此，当 FastThreadLocal 需要获取 InternalThreadLocalMap 对象时会根据当前运行的线程是不是 FastThreadLocalThread 来决定到底从哪里提取 InternalThreadLocalMap。

// InternalThreadLocalMap.java
public static InternalThreadLocalMap get() {
    Thread thread = Thread.currentThread();
    if (thread instanceof FastThreadLocalThread) {
        // 当前使用的是 FastThreadLocalThread
        return fastGet((FastThreadLocalThread) thread);
    } else {
        // 当前使用的是 Thread
        return slowGet();
    }
}
// 使用 FastThreadLocalThread 时，直接从成员变量获取
private static InternalThreadLocalMap fastGet(FastThreadLocalThread thread) {
    InternalThreadLocalMap threadLocalMap = thread.threadLocalMap();
    if (threadLocalMap == null) {
        thread.setThreadLocalMap(threadLocalMap = new InternalThreadLocalMap());
    }
    return threadLocalMap;
}
// 使用 Thread 时，通过 InternalThreadLocalMap 内部的 ThreadLocal 获取
private static InternalThreadLocalMap slowGet() {
    ThreadLocal<InternalThreadLocalMap> slowThreadLocalMap = UnpaddedInternalThreadLocalMap.slowThreadLocalMap;
    InternalThreadLocalMap ret = slowThreadLocalMap.get();
    if (ret == null) {
        ret = new InternalThreadLocalMap();
        slowThreadLocalMap.set(ret);
    }
    return ret;
}

明确了 InternalThreadLocalMap 这个容器存在哪了之后，我们再来看看这个容器内部是如何保存数据的。再回到 UnpaddedInternalThreadLocalMap 的代码中。我们可以看到这里有一个 indexedVariables 数组，它就是保存所有数据的地方，但是如果我们要使用数组就得有明确的数据下标对应关系。而那个 nextIndex 原子变量就是维护下标对应关系的关键。

class UnpaddedInternalThreadLocalMap {
    static final AtomicInteger nextIndex = new AtomicInteger();

    /** Used by {@link FastThreadLocal} */
    Object[] indexedVariables;
    // ...
    public static int nextVariableIndex() {
        int index = nextIndex.getAndIncrement();
        if (index < 0) {
            nextIndex.decrementAndGet();
            throw new IllegalStateException("too many thread-local indexed variables");
        }
        return index;
    }
}

回到 FastThreadLocal 对象中，我们可以看到这里面有一个静态属性 variablesToRemoveIndex，它的值通过 UnpaddedInternalThreadLocalMap 的静态字段 nextIndex 计算所得。因为 variablesToRemoveIndex 是一个静态属性，而且也是唯一一个静态调用 nextIndex 的地方，所以它的值恒为 0。而每个 FastThreadLocal 对象，都会在构造函数中申请一个新的 index 槽位。

public class FastThreadLocal<V> {

    private static final int variablesToRemoveIndex = InternalThreadLocalMap.nextVariableIndex();
    private final int index;
    public FastThreadLocal() {
        index = InternalThreadLocalMap.nextVariableIndex();
    }

当我们使用 FastThreadLocal 时，比如调用 set 函数，它会先获取 InternalThreadLocalMap 对象，然后根据当前 FastThreadLocal 分配的下标 index，直接设置数组 indexedVariables 中的值。

// FastThreadLocal.java
/**
 * Set the value for the current thread.
 */
public final void set(V value) {
    if (value != InternalThreadLocalMap.UNSET) {
        // 获取 Map
        InternalThreadLocalMap threadLocalMap = InternalThreadLocalMap.get();
        setKnownNotUnset(threadLocalMap, value);
    } else {
        remove();
    }
}

/**
 * @return see {@link InternalThreadLocalMap#setIndexedVariable(int, Object)}.
 */
private void setKnownNotUnset(InternalThreadLocalMap threadLocalMap, V value) {
    if (threadLocalMap.setIndexedVariable(index, value)) {
        // 如果设置成功，则将该 FastThreadLocal 保存起来，方便后续清理
        addToVariablesToRemove(threadLocalMap, this);
    }
}

// InternalThreadLocalMap.java
/**
 * @return {@code true} if and only if a new thread-local variable has been created
 */
public boolean setIndexedVariable(int index, Object value) {
    Object[] lookup = indexedVariables;
    // 如果长度足够，直接通过 index 进行修改
    if (index < lookup.length) {
        Object oldValue = lookup[index];
        lookup[index] = value;
        return oldValue == UNSET;
    } else {
        // 否则进行扩容，扩容后的大小是比 index 大的最小的2的幂
        expandIndexedVariableTableAndSet(index, value);
        return true;
    }
}

设置成功后，会将当前 FastThreadLocal 保存在一个集合中，之后在进行清理工作时，能够快速的进行清除。

// FastThreadLocal.java
@SuppressWarnings("unchecked")
private static void addToVariablesToRemove(InternalThreadLocalMap threadLocalMap, FastThreadLocal<?> variable) {
    Object v = threadLocalMap.indexedVariable(variablesToRemoveIndex);
    Set<FastThreadLocal<?>> variablesToRemove;
    // 保存 FastThreadLocal 对象的集合存在数组的 index 0 位置，因为 variablesToRemoveIndex 恒等于 0
    if (v == InternalThreadLocalMap.UNSET || v == null) {
        variablesToRemove = Collections.newSetFromMap(new IdentityHashMap<FastThreadLocal<?>, Boolean>());
        threadLocalMap.setIndexedVariable(variablesToRemoveIndex, variablesToRemove);
    } else {
        variablesToRemove = (Set<FastThreadLocal<?>>) v;
    }

    variablesToRemove.add(variable);
}

/**
 * Removes all {@link FastThreadLocal} 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.
 */
public static void removeAll() {
    // 在进行清除时，如果InternalThreadLocalMap为空，则说明没有使用 FastThreadLocal
    InternalThreadLocalMap threadLocalMap = InternalThreadLocalMap.getIfSet();
    if (threadLocalMap == null) {
        return;
    }

    try {
        Object v = threadLocalMap.indexedVariable(variablesToRemoveIndex);
        // 如果 index0 不存在集合，说明没有使用 FastThreadLocal
        if (v != null && v != InternalThreadLocalMap.UNSET) {
            @SuppressWarnings("unchecked")
            Set<FastThreadLocal<?>> variablesToRemove = (Set<FastThreadLocal<?>>) v;
            FastThreadLocal<?>[] variablesToRemoveArray =
                    variablesToRemove.toArray(new FastThreadLocal[0]);
            // 挨个删除集合中的所有 FastThreadLocal，其中会调用 FastThreadLocal 的 onRemoval 回调函数
            for (FastThreadLocal<?> tlv: variablesToRemoveArray) {
                tlv.remove(threadLocalMap);
            }
        }
    } finally {
        // 将整个 InternalThreadLocalMap 删除
        InternalThreadLocalMap.remove();
    }
}

/**
 * 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;
    }
    // 将对应槽位的值设为 null，并返回之前的值
    Object v = threadLocalMap.removeIndexedVariable(index);
    // 将该 FastThreadLocal 从待删除结合（index0 保存的集合）中删除
    removeFromVariablesToRemove(threadLocalMap, this);
    // 如果之前的值不是空，就调用 onRemoval 回调函数
    if (v != InternalThreadLocalMap.UNSET) {
        try {
            onRemoval((V) v);
        } catch (Exception e) {
            PlatformDependent.throwException(e);
        }
    }
}

至此，FastThreadLocal 的核心内容就介绍完了。因为 InternalThreadLocalMap 的下标分配器 nextIndex 是一个静态变量，所以如果当前线程没有使用所有的 FastThreadLocal 的话，那么该线程的 InternalThreadLocalMap 中的数组中都会有一些空洞。

最后，还有一点要提一下，就是 InternalThreadLocalMap 为了进一步的提高效率在成员变量中添加了几个填充字段。这是为了防止伪共享。

// InternalThreadLocalMap.java
// Cache line padding (must be public)
// With CompressedOops enabled, an instance of this class should occupy at least 128 bytes.
public long rp1, rp2, rp3, rp4, rp5, rp6, rp7, rp8, rp9;

通常 CPU 的缓存行一般是 64 或 128 字节，为了防止InternalThreadLocalMap的不同实例被加载到同一个缓存行，我们需要多余填充一些字段，使得每个实例的大小超出缓存行的大小。

下图是计算的基本结构。L1、L2、L3分别表示一级缓存、二级缓存、三级缓存，越靠近CPU的缓存，速度越快，容量也越小。所以L1缓存很小但很快，并且紧靠着在使用它的CPU内核；L2大一些，也慢一些，并且仍然只能被一个单独的CPU核使用；L3更大、更慢，并且被单个插槽上的所有CPU核共享；最后是主存，由全部插槽上的所有CPU核共享。

当CPU执行运算的时候，它先去L1查找所需的数据、再去L2、然后是L3，如果最后这些缓存中都没有，所需的数据就要去主内存拿。走得越远，运算耗费的时间就越长。所以如果你在做一些很频繁的事，你要尽量确保数据在L1缓存中。另外，线程之间共享一份数据的时候，需要一个线程把数据写回主存，而另一个线程访问主存中相应的数据。

Cache是由很多个cache line组成的。每个cache line通常是64字节，对应了主内存中的一块儿地址。CPU每次从主存中拉取数据时，会把相邻的数据也存入同一个cache line。这就有可能有可能Thread1要使用FastThreadLocal1时一次性地将两个内存地址相邻的 FastThreadLocal 对象（FastThreadLocal1，FastThreadLocal2）放入自己的 cache line 中。这时候如果另一个线程Thread2只修改了FastThreadLocal2，之后如果 Thread1 要使用 FastThreadLocal1 也需要从主存中重新拉取（因为 FastThreadLocal1 和 FastThreadLocal2 在同一个缓存行中，只要缓存行内的任意位置的数据被修改，那么其他线程就需要从主存中拉取最新的缓存行数据之后才能使用）。这种无法充分使用缓存行特性的现象，称为伪共享。

参考内容

[1] linux 2.6 互斥锁的实现-源码分析
[2] 深入解析条件变量(condition variables)
[3] Linux下Condition Vairable和Mutext合用的小细节
[4] 从ReentrantLock的实现看AQS的原理及应用
[5] 不可不说的Java“锁”事
[6] 从源码层面解析yield、sleep、wait、park
[7] LockSupport中的park与unpark原理
[8] Thread.sleep、Object.wait、LockSupport.park 区别
[9] 从AQS到futex-二-JVM的Thread和Parker
[10] Java的LockSupport.park()实现分析
[11] JVM源码分析之Object.wait/notify实现
[12] Java线程源码解析之interrupt
[13] Thread.interrupt()相关源码分析
[14] Java CAS 原理剖析
[15] 源码解析 Java 的 compareAndSwapObject 到底比较的是什么
[16] 《Java并发编程的艺术》
[17] 《实战 Java 高并发程序设计》
[18] volatile关键字深入学习
[19] 为什么Netty的FastThreadLocal速度快
[20] 线程池ThreadPoolExecutor实现原理
[21] 深入理解Java线程池：ThreadPoolExecutor
[22] ConcurrentHashMap 详解一
[23] ConcurrentHashMap 详解二
[24] JUC中Atomic class之lazySet的一点疑惑
[25] The JSR-133 Cookbook for Compiler Writers
[26] 就是要你懂Java中volatile关键字实现原理