ThreadLocalMap源码分析总结

ThreadLocalMap源码分析总结

@author：Jingdai
@date：2021.07.23

网上关于ThreadLocal的介绍非常多，但是想理解ThreadLocal，ThreadLocalMap是绕不开的，这篇文章从源码的角度分析ThreadLocalMap。作者水平有限，不免有理解错误之处，欢迎讨论。同时，读这篇文章前希望有一点点ThreadLocal的理解，否则可能不太容易懂。

简介

ThreadLocalMap 是一个Map，不同于 HashMap，它没有实现 Map 接口，也不采用链接法来解决 hash 冲突，而采用开放寻址法来解决 hash 冲突。具体来说，它使用开放寻址法的线性探查，即当前槽有元素了，就尝试下一个槽，一直到找到一个空的槽为止。

由于使用的解决冲突方法不同，它的扩容阈值，加载因子等都和 HashMap 有区别，同时考虑到 TheadLocal 弱引用的特殊性，还会考虑无效节点的擦除，下面将会从源码角度一一进行介绍。（基于JDK8源码）

主要的类结构

ThreadLocalMap 类

static class ThreadLocalMap {

    // 初始容量
    private static final int INITIAL_CAPACITY = 16;
    
    // 桶数组
    private Entry[] table;
    
    // 元素个数
    private int size = 0;
    
    // 扩容阈值 len * 2 / 3;
    private int threshold; // Default to 0    
    
}

Entry 类

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

    Entry(ThreadLocal<?> k, Object v) {
        // 可以看成 new WeakReference(k)
        // 更容易理解
        // 就是将 ThreadLocal 类型的key
        // 设置为弱引用
        super(k);
        value = v;
    }
}

如果还是不是很懂，继续点进去。

WeakReference 类

public class WeakReference<T> extends Reference<T> {

	// xx
    public WeakReference(T referent) {
        super(referent);
    }
    
    //xx
}

Reference 类

public abstract class Reference<T> {
    // xx
    private T referent;         /* Treated specially by GC */
    
    // 构造
    Reference(T referent) {
        this(referent, null);
    }

    Reference(T referent, ReferenceQueue<? super T> queue) {
        this.referent = referent;
        this.queue = (queue == null) ? ReferenceQueue.NULL : queue;
    }
    
}

可以发现最后 Entry 的 key 其实是赋值给了Reference的 referent，而如果调用 Entry 的 get() 方法也就是调用从Reference类中继承过来的 get() 方法，会返回 referent，也就是类型为ThreadLocal类型的key。即Entry的结构就如下图所示。就相当于 Entry 从 Reference 继承了 referent 变量，并把它当成key，即将ThreadLocalMap的 key 设计成弱引用。其实简单的知道 ThreadLocal 为key，同时它是弱引用就可以了。

工作流程源码分析

扩容阈值

和 HashMap不同，ThreadLocalMap 的加载因子是 2 / 3，即扩容阈值是 len * 2 / 3。

/**
 * Set the resize threshold to maintain at worst a 2/3 load factor.
 */
private void setThreshold(int len) {
    threshold = len * 2 / 3;
}

但是，和 HashMap 还有一些不同，HashMap是达到了扩容阈值就扩容，而这个有点不一样。首先在 set 方法中如果 size >= 扩容阈值，会进入 rehash 方法，rehash 会清理 table 所有的 stale Entry，如果清理过后 size >= 3/4 threshold 就调用 resize 方法扩容。当然，扩容的大小和 HashMap 一样，都是将容量扩大为原来的两倍。

// set 方法片段
int sz = ++size;
if (!cleanSomeSlots(i, sz) && sz >= threshold)
    rehash();

private void rehash() {
    expungeStaleEntries();

    // Use lower threshold for doubling to avoid hysteresis
    if (size >= threshold - threshold / 4)
        resize();
}

初始容量

private static final int INITIAL_CAPACITY = 16;

和 HashMap 一致，也是16，同时每次扩容都是两倍，所以可以保证桶的个数是2的次幂，这样在计算 key 在数组中的下标时就和HashMap一样，利用 key.hash % 2^n = key.hash & 2^n-1，使用位运算加快速度。

hash值计算

HashMap 的hash 值是直接用 key 的hashCode 然后经过hash扰动将高位参与计算后得到，而ThreadLocalMap不同，看下面代码。

int i = key.threadLocalHashCode & (len-1);

这是 set 方法中的一行，set 方法后面会讲，i 是key 应在数组中的下标，可以看出我们要想得到下标，必须得到这个hashCode，即 ThreadLocal 的 hashCode。往下看，其实这已经不是 ThreadLocalMap 的内容了，是ThreadLocal 的属性。

public class ThreadLocal<T> {

    // 需要求的值
    private final int threadLocalHashCode = nextHashCode();
    
    private static AtomicInteger nextHashCode =
        new AtomicInteger();
    
    private static final int HASH_INCREMENT = 0x61c88647;
    
    private static int nextHashCode() {
        return nextHashCode.getAndAdd(HASH_INCREMENT);
    }
    
}

看上面的代码，注意只有第一个 threadLocalHashCode 是成员变量，其他都是静态变量或方法。所以在类加载完成的时候 nextHashCode 和 HASH_INCREMENT 就已经有值了。在创建对象时，会调用静态方法给 threadLocalHashCode 赋值，即将 nextHashCode 的值赋给他，之后再将这个值加 0x61c88647，为什么是加这个，我就不懂了，应该是这个值可以保证 hash 分布更加均匀。同时，threadLocalHashCode 值是 final修饰的，代表之后就不会再修改了。

擦除 stale 的 Entry

当 ThreadLocalMap 中 Entry 的 key 即ThreadLocal对象不再使用时，又没有调用 remove 方法将它删除时，就会发生一种现象，即只有 Entry 的弱引用指向它，当发生GC时，这个 key 就会被回收，但是这个 value 却不会被回收，而我们又无法访问到这个 value，就会出现内存泄露，源码中将 Entry 的 key 被回收的 Entry 称为 stale 的Entry，会在某些操作中对这些 stale Entry进行删除。

代码中注释原文：

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.

1. expungeStaleEntry 方法

// @param: 已知为 stale Entry的索引
// @return: staleSlot后面下一个为空的槽的下标
private int expungeStaleEntry(int staleSlot) {
    Entry[] tab = table;
    int len = tab.length;

    // 擦除 staleSlot 上的Entry
    tab[staleSlot].value = null;
    tab[staleSlot] = null;
    size--;

    // 循环直到遇到null
    Entry e;
    int i;
    for (i = nextIndex(staleSlot, len);
         (e = tab[i]) != null;
         i = nextIndex(i, len)) {
        // nextIndex 是取tab的下一个下标
        ThreadLocal<?> k = e.get();
        if (k == null) {
            // 擦除遇到的stale节点
            e.value = null;
            tab[i] = null;
            size--;
        } else {
            int h = k.threadLocalHashCode & (len - 1);
            // 说明这个 key 是线性探测后的下标，不在 h 上
            if (h != i) {
                tab[i] = null;
                // Unlike Knuth 6.4 Algorithm R, we must scan until
                // null because multiple entries could have been stale.
                // 对 e 进行rehash，放到更应该放的位置
                while (tab[h] != null)
                    h = nextIndex(h, len);
                tab[h] = e;
            }
        }
    }
    return i;
}

这个方法传入一个已知为 stale Entry 的节点的下标，然后从这个下标开始找下一个为 null 的位置并返回。在找的过程中，当遇到 stale Entry，就会擦除这个entry，如果遇到 entry 被线性探测过，就尝试将这个 entry rehash，使它放置在更加合适的位置上。

2.expungeStaleEntries 方法

private void expungeStaleEntries() {
    Entry[] tab = table;
    int len = tab.length;
    for (int j = 0; j < len; j++) {
        Entry e = tab[j];
        if (e != null && e.get() == null)
            expungeStaleEntry(j);
    }
}

这个方法很简单，就是擦除 table 上所有的 stale Entry 节点。

3.cleanSomeSlots 方法

// 启发式的扫描 stale Entry
// 执行对数复杂度的扫描
// 作为不扫描和全table扫描的一种折中
// @param: i 已知不是staleEntry的下标，扫描从i后一个开始
// @param: n 如果没有找到stale Entry，log2(n) 个单元格会被扫描
//         如果找到了stale Entry，将额外扫描 log2(table.length)-1
//         个单元格
//         插入方法调用时，这个参数是元素个数
//         replaceStaleEntry 方法调用时，参数是 table length
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;
}

这个方法从下标 i 开始，执行 log2n 次的扫描来发现 stale Entry，当发现 stale Entry时，会调用 expungeStaleEntry 方法来擦除stale entry。这是全 table 扫描和不扫描的一种折中。

get 方法

getEntry 方法

private Entry getEntry(ThreadLocal<?> key) {
    // 如前所述，hash & (len - 1) 取得下标
    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);
}

getEntryAfterMiss 方法

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;
        // 发现 stale Entry 就去 擦除
        if (k == null)
            expungeStaleEntry(i);
        // 往后移动一位
        else
            i = nextIndex(i, len);
        e = tab[i];
    }
    return null;
}

看到这里，不知道大家有没有一个疑问，这样会不会由于线性探查找不到存入的数据。比如有三个元素A、B、C，它们三个元素 hash 取余后都为1，那么依次放入这三个元素就会导致 A 在槽1，B在槽2，C在槽3，这时删掉A，槽1为 null，我们再去调用 getEntry 方法去获取B，由于 B 的hash 取余后得到1，而1处为 null，getEntryAfterMiss 方法就会返回null，这样 B 元素不就丢了吗？看到这里是这样，带着这个问题继续往下看。

put 方法

ThreadLocalMap 的 put 方法叫做 set 方法。

set 方法

private void set(ThreadLocal<?> key, Object value) {

    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)]) {
        ThreadLocal<?> k = e.get();

        if (k == key) {
            // 更新
            e.value = value;
            return;
        }

        if (k == null) {
            // 替换 stale Entry
            replaceStaleEntry(key, value, i);
            return;
        }
    }

    // 到这说明tab[i] = null
    tab[i] = new Entry(key, value);
    int sz = ++size;
    // cleanSomeSlots 没有删除元素 
    // 且 sz >= threshold
    if (!cleanSomeSlots(i, sz) && sz >= threshold)
        rehash();
}

这个方法就是尝试去找到一个空槽，把值放进去，然后根据情况考虑是否去rehash。在遍历找空槽的过程中，如果找到了相同的 key，说明需要覆盖，就把旧值进行覆盖；如果发现了 stale Entry，就调用 replaceStaleEntry 方法替换 stale entry。

replaceStaleEntry 方法

在看 replaceStaleEntry 方法的源码之前，先明确一下这个方法要干什么事，要不就会一脸懵逼，这个方法应该是ThreadLocalMap 里算逻辑复杂的了。看源代码注释。

this method expunges all stale entries in the "run" containing the stale entry.  (A run is a sequence of entries between two null slots.)

看上面，这个方法除了替换 stale entry ，还会将传入的下标之前第一个的 null 下标和之后的第一个 null 下标之间的所有 stale entry 全部擦除掉。注释中将两个 null 之间的序列称为一个 run。下面看具体怎么做的。

private void replaceStaleEntry(ThreadLocal<?> key, Object value,
                               int staleSlot) {
    Entry[] tab = table;
    int len = tab.length;
    Entry e;

	// 用 slotToExpunge 记录这个 run
    // 的第一个 stale entry
    // 直到遇到空槽跳出
    int slotToExpunge = staleSlot;
    for (int i = prevIndex(staleSlot, len);
         (e = tab[i]) != null;
         i = prevIndex(i, len))
        if (e.get() == null)
            slotToExpunge = i;

    // 从前向后找key，遇到空槽跳出
    for (int i = nextIndex(staleSlot, len);
         (e = tab[i]) != null;
         i = nextIndex(i, len)) {
        ThreadLocal<?> k = e.get();

        // 这里有一点点复杂，看下面的图好理解一点
        if (k == key) {
            e.value = value;

            // 就是将entry 放置到staleslot的位置上，
            // 同时向i的位置放原来staleslot 的stale entry
            // 不知道你们发现没有，整个过程都没有擦除过stale entry，
            // 因为方法要保证整个 run 的stale entry 都要被擦除，而
            // 如果擦除将会使它为null，就分割了这个run，再调用
            // expungeStaleEntry 方法就只能擦除分割的run，
            // 不是replaceStaleEntry 方法调用时的run
            tab[i] = tab[staleSlot];
            tab[staleSlot] = e;

            // Start expunge at preceding stale entry if it exists
            if (slotToExpunge == staleSlot)
                slotToExpunge = i;
            cleanSomeSlots(expungeStaleEntry(slotToExpunge), len);
            return;
        }
        
       	// key == null 说明是staleEntry
        // slotToExpunge == staleSlot 说明之前没有发现stale entry
        // 1.之后是null跳出，会擦除原始的staleSlot，再从slotToExpunge清除√
        // 2.之后找到 k == null，不会变，还是第一个  √
        // 3.之后找到 k == key，因为会交换一下位置，所以slotToExpunge 还是第一个 √
        if (k == null && slotToExpunge == staleSlot)
            slotToExpunge = i;
    }

    // 走到这里说明是遇到空槽跳出的
    // key 没有找到，则插入新的 entry
    
    // 擦除原来的 value 
    tab[staleSlot].value = null;
    // 插入新entry
    tab[staleSlot] = new Entry(key, value);

    // 说明run中有其他的stale entry，清除一下
    if (slotToExpunge != staleSlot)
        cleanSomeSlots(expungeStaleEntry(slotToExpunge), len);
}

当 k == key 时交换的流程，最开始走到这里，发现 k 和 key相同。

然后执行 e.value = value; 如下图。

接着进行交换。

从图就很容易理解了，执行的结果就是将这个新Entry插入到staleSlot位置上，并使 i 位置保存原来 staleSlot 的 stale entry。有些人可能会想这里为什么不把这个stale entry直接擦除，原因是这个方法的目的是将整个 run 上所有的 stale entry 全部擦除，之后会调用 expungeStaleEntry 方法从这个 run 的第一个 stale entry开始清除，如果在这个方法之前将 stale entry 擦除了，就会导致整个 run 断开，expungeStaleEntry 方法就只会擦除原始 run 的前一部分，所以必须等最后将擦除任务交给 expungeStaleEntry 方法。

现在总结一下 set 方法的流程。

步骤1：求出要插入 key 的下标，从这个下标依次往后遍历，直到遇到3种情况。情况1：要插入的 key 和 下标处的 k 相等，执行步骤2；情况2：遇到空槽，执行步骤3；情况3：遇到 stale entry，执行步骤4。

步骤2：要插入的 key 和此槽处的值相同，说明需要更新，将 value 值更新，整个过程结束。

步骤3：遇到空槽，就新建一个entry，插入到空槽上，然后执行 cleanSomeSlots 方法清理一些stale entry，如果方法返回 false（代表没有清理）并且 size >= threshold，执行rehash。

步骤4：遇到 stale entry，需要替换整个stale entry，执行 replaceStaleEntry 方法来完成。这个方法除了会替换stale entry，还会将整个 run（两个null之间的序列） 上所有的stale entry进行清除。

remove 方法

remove 方法

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;
        }
    }
    // 走到这里说明遇到空槽，没有对应的 key
}

这个方法没什么说的。

rehash 和扩容

rehash 方法

private void rehash() {
    // 擦除 table上所有的 stale entry
    expungeStaleEntries();

    // Use lower threshold for doubling to avoid hysteresis
    // 擦除后当 size 大于 3/4 threshold 时才会resize
    if (size >= threshold - threshold / 4)
        resize();
}

resize 方法

private void resize() {
    Entry[] oldTab = table;
    int oldLen = oldTab.length;
    // 扩容为2倍
    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();
            // 应该永远不会为null，因为整个类中
            // 只有 rehash 方法调用了 resize
            // 而 rehash函数一开始就清除了所有的
            // stale entry,所以k不会为null
            if (k == null) {
                e.value = null; // Help the GC
            } else {
                int h = k.threadLocalHashCode & (newLen - 1);
                while (newTab[h] != null)
                    h = nextIndex(h, newLen);
                newTab[h] = e;
                count++;
            }
        }
    }
	
    // 设置阈值为 len * 2 / 3;
    setThreshold(newLen);
    size = count;
    table = newTab;
}

这两个方法比较简单，rehash 方法首先会擦除所有的 stale entry，擦除过后如果 size >= 3/4 threshold，则进行resize 扩容，否则不会扩容，即 rehash 方法不一定会导致扩容。扩容时会将数组长度扩大为原来的两倍，然后依次把旧数组上的元素移动过去。

问题

到这里，整个ThreadLocalMap的源码全部分析完了， 除了个别2-3行的方法没有分析，大家一看就懂，其他所有的方法全部分析完成。现在就是回答上面的问题。

复述一下问题：比如有三个元素A、B、C，它们三个元素 hash 取余后都为1，那么依次放入这三个元素就会导致 A 在槽1，B在槽2，C在槽3，这时删掉A，槽1为 null，我们再去调用 getEntry 方法去获取B，由于 B 的hash 取余后得到1，而1处为 null，getEntryAfterMiss 方法就会返回null，这样 B 元素不就丢了吗？

真的会丢吗？上面所有的方法中删除元素，不管是 remove 方法直接删除，还是删除 stale entry，都不会直接去删除（除了expungeStaleEntry方法内部），都会去调用 expungeStaleEntry 方法去删除。这个方法会删除下一个遇到的 null 之前的所有 stale entry，同时，在遍历的过程中，如果遇到元素的下标和直接 hash 取余的下标不同的话，这个方法还会尝试将这个元素重新放置，使它置于更加合适的位置上。所以并不会发生上面问题中描述的情况，在删除元素 A 时，由于后面还有两个元素，不是null，同时元素B、C的hash值取余都不是他们现在所处的位置，所以 expungeStaleEntry 方法会尝试重新放置B、C元素，方法执行完成后 B 会位于槽1，C会位于槽2，所以再次调用 getEntry 方法会得到 B，不会出现元素丢失的情况。

参考

• Java8-API