ThreadLocal源码分析

0x00 ThreadLocal

ThreadLocal提供了线程局部的变量，但和普通局部变量不同，同一个ThreadLocal变量可以被多个线程共享，而不是线程私有的。

在ThreadLocal源代码中有一个使用例子，代码如下：

import java.util.concurrent.atomic.AtomicInteger;

public class ThreadId {
    // Atomic integer containing the next thread ID to be assigned
    private static final AtomicInteger nextId = new AtomicInteger(0);

    // Thread local variable containing each thread's ID
    private static final ThreadLocal<Integer> threadId =
    new ThreadLocal<Integer>() {
        @Override protected Integer initialValue() {
            return nextId.getAndIncrement();
        }
    };

    // Returns the current thread's unique ID, assigning it if necessary
    public static int get() {
        return threadId.get();
    }
}

我们以这段代码为例，简单探索一下ThreadLocal的内部原理。

假设我们现在需要生成线程ID，则可以这么写：

public static void main(String[] args) {
        new Thread(() -> {
            System.out.println(ThreadId.get());
        }).start();

        new Thread(() -> {
            System.out.println(ThreadId.get());
        }).start();
}

// 输出
// 0
// 1

为什么不同线程调用ThreadId.get()会有不同的行为呢？我们看一下ThreadLocal类的get()方法，

public T get() {
    Thread t = Thread.currentThread();
    ThreadLocalMap map = getMap(t);
    if (map != null) {
        ThreadLocalMap.Entry e = map.getEntry(this);
        if (e != null) {
            @SuppressWarnings("unchecked")
            T result = (T)e.value;
            return result;
        }
    }
    return setInitialValue();
}

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

可以发现当一个线程调用get()时，会取出该线程内部的threadLocals变量，并以ThreadLocal对象为key取出map中对应的值。如果threadLocals为null，则再调用setInitialValue()方法，创建一个新的threadLocals变量。

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);
    }
    if (this instanceof TerminatingThreadLocal) {
        TerminatingThreadLocal.register((TerminatingThreadLocal<?>) this);
    }
    return value;
}

void createMap(Thread t, T firstValue) {
    t.threadLocals = new ThreadLocalMap(this, firstValue);
}

protected T initialValue() {
    return null;
}

在调用setInitialValue()方法时需要调用initialValue()获得一个初值，默认情况下是null。因此在开头的例子中也重写了这个initialValue()来为ThreadLocal变量提供初值。

既然有get()方法，同样也有set()方法，内部仅仅是调用了map的set方法而已。

public void set(T value) {
    Thread t = Thread.currentThread();
    ThreadLocalMap map = getMap(t);
    if (map != null) {
        map.set(this, value);
    } else {
        createMap(t, value);
    }
}

所以到目前为止，我们发现：每个Thread对象内部都会有一个成员变量threadLocals，这实质上是一个哈希表，它以ThreadLocal为key，而value的类型是不固定的，取决于ThreadLocal对象声明时的泛型。
当同一个ThreadLocal对象被多个线程访问时，每个线程都会创建一个键值对并放入自己的threadLocals变量内，这个键值对的key是该ThreadLocal对象，而值则是由ThreadLocal对象的initialValue()提供。除了一个ThreadLocal对象可以被多个线程访问，一个线程内的threadLocals也可以同时持有多个ThreadLocal对象，因此线程和ThreadLocal是多对多的关系。

回顾开头提到的，同一个ThreadLocal对象被多个线程共享但仍然线程安全的原因在于，这个ThreadLocal对象仅仅是作为一个只读的key存放在每个线程的threadLocals内部，而它对应的值则是以副本的形式作为键值对的value，每个线程的threadLocals以相同的ThreadLocal作为key，但是值是不共享的。

0x01 ThreadLocalMap

在前一节中，我们发现ThreadLocal仅仅是存在每个线程的threadLocals里，threadLocals是一个ThreadLocalMap对象，我们简单探索一下它的内部。

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;
        }
    }

    private Entry[] table;
    
    // ...
}

可以发现ThreadLocalMap是一个哈希表，并且采用线性探测法来解决冲突。其内部类Entry继承了弱引用，因此每个entry的key是对ThreadLocal对象的弱引用，而value则是强引用。

这里其实有些令人疑惑，为什么Entry类要对key使用弱引用而不是强引用呢？在源代码中我发现这样一段注释：
To help deal with very large and long-lived usages, the hash table entries use WeakReferences for keys. However, since reference queues are not used, stale entries are guaranteed to be removed only when the table starts running out of space.
原意大概是说，为了帮助处理那些大且长时间存活的对象，因此对key进行弱引用，但是由于没有使用到引用队列，因此这些过期的entry只有在哈希表快用完时才保证一定删除。实际上，在ThreadLocalMap的getEntry或setEntry方法中，只要检测到某个key为null，就会将对应的value引用置为null，让GC能够及时回收。

因此我们可以想象，当某一个ThreadLocal对象不再使用了，那么所有线程中关于该ThreadLocal对象的value也会被回收。这里涉及到弱引用的概念，如果某个对象仅仅存在弱引用，那么GC一旦发现该对象，无论在内存是否充足的情况下都会立刻回收它。

0x02 ThreadLocal内存泄漏

前面我们提高，ThreadLocal类对象一旦不再使用，其对应的value在各个线程中也会被回收。然而，实际使用中我们通常将ThreadLocal类对象设为static，这就意味着整个程序运行过程中它都会一直存活，那么只要该线程不结束，即使某一个ThreadLocal对象不再使用了，线程内部的threadLocals始终不会被回收。这就引起了内存泄漏问题。

要解决这个问题也很简单，就是一旦ThreadLocal对象在某个线程中不再使用了，就调用remove方法去手动删除对应的entry。

/**
 * Remove the entry for key.
 */
private void remove(ThreadLocal<?> key) {
    Entry[] tab = table;
    int len = tab.length;
    int i = key.threadLocalHashCode & (len-1);
    for (Entry e = tab[i];
         e != null;
         e = tab[i = nextIndex(i, len)]) {
        if (e.refersTo(key)) {
            e.clear();
            expungeStaleEntry(i);
            return;
        }
    }
}