ThreadLocal源码分析

一、简介

ThreadLocal提供了线程本地变量，通过get或者set操作的这些变量在每个不同线程间是不相同的，各个线程独立地初始化这些变量。ThreadLocal实例通常在类中是声明为private static域的，用于在同一个线程内关联相同的状态(e.g. 一个User ID或者Transaction ID)。ThreadLocal相当于提供了一种线程隔离，将变量与线程相绑定。

只要线程还存活并且ThreadLocal实例还能被获取到，那么每个线程会持有一个ThreadLocal变量弱引用。当线程结束生命周期时，所有的线程本地实例都会被GC。

二、简单示例

/**
 * 不同线程持有一个不同的UUID
 */
 public class ThreadLocalTest {
     private static ThreadLocal<String> uuidLocal = new ThreadLocal<String>(){
         protected String initialValue() {
             return UUID.randomUUID().toString();
         }
     };

     public static void main(String[] args) {
         UUIDThread t1 = new UUIDThread();
         UUIDThread t2 = new UUIDThread();
         t1.start();
         t2.start();
     }

     public static class UUIDThread extends Thread {
         @Override
         public void run() {
             System.out.println(Thread.currentThread().getName() + " uuid: " + uuidLocal.get());
         }
     }
 }

输出结果两个uuid不同，如下所示：

Thread-1 uuid: d8d2006f-0a8a-4999-90c0-de2648c742da
Thread-0 uuid: 5061e4bd-8f57-4ef6-8b74-e7571f9efb93

三、我司的用法

使用用户公司来做分库，不同的公司数据分在不同的业务库中，将companyID存入DataSourceContext中，查询数据库的时候从DataSourceContext获取对应的companyID，根据companyID获取对应的数据库链接。

public class DataSourceContext {
    static final ThreadLocal<String> local = new ThreadLocal<>();
    public DataSourceContext() {
    }
    
    public static String getCompany() {
        return local.get();
    }

    public static  String setCompany(String companyID) {
        return companyID == null ? null : set(companyID);
    }
}

四、成员变量

private final int threadLocalHashCode = nextHashCode(); // 初始值为0

private static AtomicInteger nextHashCode =
    new AtomicInteger();

private static final int HASH_INCREMENT = 0x61c88647;

private static int nextHashCode() {
    return nextHashCode.getAndAdd(HASH_INCREMENT);
}

ThreadLocal通过自定义threadLocalHashCode来减少线性探测的冲突，每次实例化一个ThreadLocal，threadLocalHashCode都会新增HASH_INCREMENT(0x61c88647)。

五、几个方法

1. initialValue方法

/**
 * 返回当前线程的线程本地变量初始化值，在一个线程首次调用get方法时被调用。
 * 如果调用get之前调用了set方法，就不会调用initialValue方法了。
 * 默认实现返回null，如果有需要，可以继承ThreadLocal，并覆盖该方法。
 * 一般是使用匿名内部类的形式子类化。
 */
protected T initialValue() {
    return null;
}

2. get方法

/**
 * 返回当前线程的线程本地变量
 */
public T get() {
    // 获取当前线程的引用
    Thread t = Thread.currentThread();
    // 从当前线程中获取到关联的ThreadLocalMap
    ThreadLocalMap map = getMap(t);
    if (map != null) {
        ThreadLocalMap.Entry e = map.getEntry(this);
        if (e != null)
            return (T)e.value;
    }
    // 如果当前线程没有线程本地变量，就设置初始值
    return setInitialValue();
}

/**
 * 从ThreadLocal中获取一个关联的Map
 */
ThreadLocalMap getMap(Thread t) {
    return t.threadLocals;
}

可以从代码中看出来，get()方法就是从当前线程中获取一个和当前线程相关联的ThreadLocalMap，然后以this为key，从ThreadLocalMap中取出相应的值，并返回。如果没有值，就设置一个初始值。

/* ThreadLocal values pertaining to this thread. This map is maintained
 * by the ThreadLocal class. */
ThreadLocal.ThreadLocalMap threadLocals = null;

threadLocals是Thread的成员变量，每个线程通过ThreadLocal.ThreadLocalMap与ThreadLocal相绑定，这样可以确保每个线程访问到的thread-local variable都是本线程的。

3. set方法

/**
 * 设置初始值
 */
private T setInitialValue() {
    T value = initialValue();
    Thread t = Thread.currentThread();
    ThreadLocalMap map = getMap(t);
    if (map != null)
        map.set(this, value);
    else
        // 如果线程t不存在ThreadLocalMap实例，就创建一个
        createMap(t, value);
    return value;
}

/**
 * 实例化一个ThreadLocalMap并赋值给t.threadLocals
 */
void createMap(Thread t, T firstValue) {
    t.threadLocals = new ThreadLocalMap(this, firstValue);
}

/**
 * 设置当前线程的线程本地变量值
 */
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()方法和setInitialValue()方法类似，如果当前线程存在threadLocals，那么直接把设置的值put到这个ThreadLocalMap中。否则，创建一个带有这个value的ThreadLocalMap。

4. remove方法

public void remove() {
    ThreadLocalMap m = getMap(Thread.currentThread());
    if (m != null)
        m.remove(this);
}

首先获取当前线程，并从当前线程中获取ThreadLocalMap，如果不为空，则调用ThreadLocalMap的remove方法，把以this为key的Entry移除掉。如果随后在当前线程中被调用了get方法，那么因为原先的Entry已经被移除掉了，所以还会调用一次initialValue()方法初始化值。

5. 小结

从这些方法可以看出ThreadLocal类的设计，Thread中有ThreadLocalMap成员变量，ThreadLocalMap又以ThreadLocal作为key来存放值。也就是说ThreadLocal把自身实例作为key，和需要保存的value存放到当前线程的一个Map中，来保证每个线程访问到的线程本地变量值都是各自线程的。ThreadLocal#set方法可以简单理解为Thread.currentThread().threadLocals.put(this, value)，ThreadLocal#get方法可以简单理解为Thread.currentThread().threadLocals.get(this)。

六、ThreadLocalMap

ThreadLocal实现中，核心还是ThreadLocalMap。ThreadLocal只是作为ThreadLocalMap的key, 从ThreadLocalMap中获取到相应的值。下面简单看下ThreadLocalMap的实现。

1. 成员变量

/**
 * 初始容量，必须是2^n
 */
private static final int INITIAL_CAPACITY = 16;

/**
 * 必要时会扩容，但必须是2^n
 */
private Entry[] table;

/**
 * table中entry的数量，也就是ThreadLocalMap的大小
 */
private int size = 0;

/**
 * 下一次扩容的阈值
 */
private int threshold; // Default to 0

其中INITIAL_CAPACITY代表这个ThreadLocalMap的初始容量；table是一个Entry类型的数组，用于存储数据；size代表表中的存储数目，也就是ThreadLocalMap的大小；threshold代表需要扩容时对应size的阈值。

2. 静态内部类

static class Entry extends WeakReference<ThreadLocal> {
    /** The value associated with this ThreadLocal. */
    Object value;

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

Entry继承了WeakReference，当除了Entry以外没有其它地方强引用ThreadLocal实例，那么ThreadLocal实例就会被GC回收，避免造成内存溢出。

3. 构造函数

ThreadLocalMap(ThreadLocal firstKey, Object firstValue) {
    table = new Entry[INITIAL_CAPACITY];
    // 对hashcode“取模”计算出table中索引值
    int i = firstKey.threadLocalHashCode & (INITIAL_CAPACITY - 1);
    table[i] = new Entry(firstKey, firstValue);
    size = 1;
    setThreshold(INITIAL_CAPACITY);
}

计算索引值i的时候里面采用了hashCode & (size - 1)的算法，这相当于取模运算hashCode % size的一个更高效的实现（和HashMap中的思路相同）。正是因为这种算法，要求size必须是2^n。

4. set方法

大致思路为：

1. 通过key的hashcode计算出索引值
2. 从索引值i开始，通过线性探测法从table中找到一个可以存放value的地方，然后设置值
3. 因为ThreadLocalMap的key是WeakReference，所以会存在Entry存在，但是key已经被回收的情况，这时候需要进行一些清理工作，把这些Entry清理掉。
4. 如果size大于阈值(threshold)，就要进行扩容，并rehash，从新计算映射。

private void set(ThreadLocal key, Object value) {
    Entry[] tab = table;
    int len = tab.length;
    // 计算索引值
    int i = key.threadLocalHashCode & (len-1);

    // 使用线性探测法来解决冲突，而不是HashMap中采用的拉链法
    for (Entry e = tab[i];
         e != null;
         e = tab[i = nextIndex(i, len)]) {
        ThreadLocal k = e.get();

        // 如果欲设置的key和table[i]中的相同，则更新value
        if (k == key) {
            e.value = value;
            return;
        }

        // 如果k==null，证明key(WeakReference)已经被GC回收，所以替换新的key和value
        if (k == null) {
            replaceStaleEntry(key, value, i);
            return;
        }
    }
    // 创建新的Entry
    tab[i] = new Entry(key, value);
    int sz = ++size;
    // 如果更新之后的size大于阈值threshold，则需要rehash
    if (!cleanSomeSlots(i, sz) && sz >= threshold)
        rehash();
}

/** 线性探测法的套路，找到下一个索引，如果越界了，就从0开始 */
private static int nextIndex(int i, int len) {
    return ((i + 1 < len) ? i + 1 : 0);
}

5. getEntry方法

ThreadLocal#get方法就是调用了ThreadLocalMap#getEntry方法。

大致思路为：

1. 通过key的hashcode计算出索引值i
2. 为了提高性能，直接判断索引值下的Entry是不是需要找的
3. 否则，用线性探测的方式找到相应的value

private Entry getEntry(ThreadLocal key) {
    // 计算出索引值
    int i = key.threadLocalHashCode & (table.length - 1);
    Entry e = table[i];
    // 命中，则直接返回
    if (e != null && e.get() == key)
        return e;
    else
        // 按线性探测的方式查找
        return getEntryAfterMiss(key, i, e);
}

/**
 * 和getEntry类似，用于当key的hash直接计算出的索引值上找不到Entry时
 */
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)
            // 清理过期的Entry
            expungeStaleEntry(i);
        else
            i = nextIndex(i, len);
        e = tab[i];
    }
    return null;
}

6. remove方法

remove方法和getEntry方法类似，计算索引值i，用线性探测的防止，找到Entry后，清理Entry。

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.get() == key) {
            e.clear();
            expungeStaleEntry(i);
            return;
        }
    }
}

7. 扩容方法

扩容方法也很清楚，判断目前使用的容量是否大于一定的值(size >= 3 / 4 * threshold)，如果大于，则需要resize。

resize方法的思路如下：

1. size扩大为两倍，创建一个新的table表，将oldTab上的Entry转移到newTab上。
2. 转移过程中，如果发现e.get() == null，则证明key已经被GC回收，那么这个Entry就不转移。
3. 否则，用线性探测法找到Entry在newTab存放的位置，并设置。
4. 最后设置新的threshold。

可以看出threshold的大小为len * 2 / 3，所以每次size >= 0.5 * len的时候就要进行扩容(resize)。

private void rehash() {
    expungeStaleEntries();

    // 如果size > 3/4 * threshold，则扩容
    if (size >= threshold - threshold / 4)
        resize();
}
/**
 * table容量翻倍
 */
private void resize() {
    Entry[] oldTab = table;
    int oldLen = oldTab.length;
    int newLen = oldLen * 2;
    Entry[] newTab = new Entry[newLen];
    int count = 0;

    for (int j = 0; j < oldLen; ++j) {
        Entry e = oldTab[j];
        if (e != null) {
            ThreadLocal k = e.get();
            // 如果k已经被GC回收，那么把value也设置为null，帮助GC回收，防止内存泄漏
            if (k == null) {
                e.value = null; // Help the GC
            } else {
                // 从新计算索引值，并通过线性探测的方式存放到table中
                int h = k.threadLocalHashCode & (newLen - 1);
                while (newTab[h] != null)
                    h = nextIndex(h, newLen);
                newTab[h] = e;
                count++;
            }
        }
    }

    setThreshold(newLen);
    size = count;
    table = newTab;
}

private void setThreshold(int len) {
    threshold = len * 2 / 3;
}

8. 一些清理方法

1) expungeStaleEntry

首先会清理tab[staleSlot]上过期的Entry，然后需要再散列(rehash)，中间可能还会遇到一些过期的Entry，这些也要清理掉，知道遇到table[i] == null，中间的所有Entry都要rehash。

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

    // 清理掉过期位置的Entry
    tab[staleSlot].value = null;
    tab[staleSlot] = null;
    size--;

    // 再散列，知道遇到table[i] == null
    Entry e;
    int i;
    for (i = nextIndex(staleSlot, len);
         (e = tab[i]) != null;
         i = nextIndex(i, len)) {
        ThreadLocal k = e.get();
        if (k == null) {
            e.value = null;
            tab[i] = null;
            size--;
        } else {
            int h = k.threadLocalHashCode & (len - 1);
            // 如果h==i，那么证明这个Entry就是要放在table[i]上的，就不要rehash这个Entry
            // 否则，rehash
            if (h != i) {
                // 先把当前位置tab[i]释放出来，再把Entry放到新的位置
                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;
}

2) cleanSomeSlots

/**
 * 探索式扫描寻找过期的entry，当增加新元素或者另一个过期entry被清理的时候会被调用
 */
private boolean cleanSomeSlots(int i, int n) {
    boolean removed = false;
    Entry[] tab = table;
    int len = tab.length;
    do {
        i = nextIndex(i, len);
        Entry e = tab[i];
        // 判断是否过期，如果是则处理
        if (e != null && e.get() == null) {
            n = len;
            removed = true;
            i = expungeStaleEntry(i);
        }
    } while ( (n >>>= 1) != 0);
    return removed;
}

每次新增元素的时候会进行探索式扫描，寻找过期Entry并清理。

3) 何时会清理过期Entry

处理Thread实例被GC回收，ThreadLocalMap同时被回收之外，下面这些条件下，会清理过期的Entry。

1. getEntry时，线性探索寻找Entry的时候发现Entry过期。
2. set的时候发现，key对应索引值的Entry已过期，则会清理并替换
3. 每次调用set方法的时候，会探索式扫描Entry，如果发现过期，则清理。
4. size > threshold，rehash的时候。
5. 调用remove方法的时候。

当前的应用开发过程中，出于复用的目的，常常会使用线程池的技术，线程中ThreadLocalMap可能会长期存在。因为Entry中的key被WeakReference包装，在key不存在强引用的时候，会回收key，但是Entry和value并不会被回收。所以在ThreadLocalMap中需要不时地清理过期的Entry，来保证内存不泄露。当然，如果我们在代码中每次使用完ThreadLocal，都可以remove一下，那么就可以尽早释放不需要的内存。

七、参考

并发编程 | ThreadLocal源码深入分析