Java 8 ThreadLocal 源码解析

Posted 2021-01-29 Finley

Java 中的 ThreadLocal是线程内的局部变量, 它为每个线程保存变量的一个副本。ThreadLocal 对象可以在多个线程中共享, 但每个线程只能读写其中自己的副本。

目录：

• 代码示例
• 源码解析
• InheritableThreadLocal
• ThreadLocalMap
  • Get 流程
  • Set 流程
  • Remove

代码示例

我们编写一个简单的示例:

import java.util.Random;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

/**
 * @author finley
 */
public class MyThread extends Thread {

    private static final ThreadLocal<Integer> threadLocal = new ThreadLocal<>();

    private static final Random random = new Random();

    @Override
    public void run() {
        int localValue = random.nextInt();
        threadLocal.set(localValue);
        System.out.println("Thread: " + Thread.currentThread().getName() + " set thread local: " + localValue);
        try {
            Thread.sleep(1000L);
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
        System.out.println("Thread: " + Thread.currentThread().getName() + " threadLocal: " + threadLocal.get() + ", local: " + localValue);
    }

    public static void main(String[] args) {
        int concurrent = 10;
        ExecutorService service = Executors.newFixedThreadPool(concurrent);
        for (int i = 0; i < concurrent; i++) {
            service.execute(new MyThread());
        }
        service.shutdown();
    }

}

运行结果:

Thread: pool-1-thread-1 set thread local: -1262320606
Thread: pool-1-thread-2 set thread local: 1222545653
Thread: pool-1-thread-3 set thread local: 2067394038
Thread: pool-1-thread-4 set thread local: 920362206
Thread: pool-1-thread-5 set thread local: -1977931750
Thread: pool-1-thread-6 set thread local: 985735150
Thread: pool-1-thread-7 set thread local: -602752866
Thread: pool-1-thread-8 set thread local: 672437027
Thread: pool-1-thread-9 set thread local: 1063652674
Thread: pool-1-thread-10 set thread local: 1790288576
Thread: pool-1-thread-1 threadLocal: -1262320606, local: -1262320606
Thread: pool-1-thread-3 threadLocal: 2067394038, local: 2067394038
Thread: pool-1-thread-4 threadLocal: 920362206, local: 920362206
Thread: pool-1-thread-6 threadLocal: 985735150, local: 985735150
Thread: pool-1-thread-7 threadLocal: -602752866, local: -602752866
Thread: pool-1-thread-2 threadLocal: 1222545653, local: 1222545653
Thread: pool-1-thread-5 threadLocal: -1977931750, local: -1977931750
Thread: pool-1-thread-8 threadLocal: 672437027, local: 672437027
Thread: pool-1-thread-10 threadLocal: 1790288576, local: 1790288576
Thread: pool-1-thread-9 threadLocal: 1063652674, local: 1063652674

可以看到10个线程调用同一个ThreadLocal对象的set方法写入随机值, 然后读取到自己写入的副本。

源码解析

我们从ThreadLocal.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);
}

set方法将当前线程的副本写入了一个ThreadLocalMap, map的key是当前的ThreadLocal对象。

接下来通过getMap方法分析这个ThreadLocalMap是如何维护的:

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

public
class Thread implements Runnable {
    ThreadLocal.ThreadLocalMap threadLocals = null;
}

每个 Thread 对象维护了一个 ThreadLocalMap 类型的 threadLocals 字段。

ThreadLocalMap 的 key 是 ThreadLocal 对象, 值则是变量的副本, 因此允许一个线程绑定多个 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();
}

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

首先获得当前线程的ThreadLocalMap, 然后从 ThreadLocalMap 尝试获得当前 ThreadLocal 对象对应的副本。

若获取失败，则写入并返回initialValue方法定义的默认值。

Thread.threadLocals 字段是惰性初始化的。 ThreadLocal.set() 方法发现 threadLocals 为空时会调用 createMap 方法进行初始化, ThreadLocal.get()方法同样会在setInitialValue() 中调用 createMap 方法初始化 Thread.threadLocals 字段。

为了不影响读者整体了解ThreadLocal, ThreadLocalMap 的实现原理在最后一节ThreadLocalMap

InheritableThreadLocal

InheritableThreadLocal 在子线程创建时将父线程的变量副本传递给子线程。

InheritableThreadLocal 继承了 ThreadLocal 并重写了3个方法, 它使用 Thread.inheritableThreadLocals 代替了 Thread.threadLocals 字段。

public class InheritableThreadLocal<T> extends ThreadLocal<T> {

    protected T childValue(T parentValue) {
        return parentValue;
    }

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

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

ThreadLocalMap 的构造器中实现了向子线程传递的逻辑:

private ThreadLocalMap(ThreadLocalMap parentMap) {
    Entry[] parentTable = parentMap.table;
    int len = parentTable.length;
    setThreshold(len);
    table = new Entry[len];

    for (int j = 0; j < len; j++) {
        Entry e = parentTable[j];
        if (e != null) {
            @SuppressWarnings("unchecked")
            ThreadLocal<Object> key = (ThreadLocal<Object>) e.get();
            if (key != null) {
                Object value = key.childValue(e.value);
                Entry c = new Entry(key, value);
                int h = key.threadLocalHashCode & (len - 1);
                while (table[h] != null)
                    h = nextIndex(h, len);
                table[h] = c;
                size++;
            }
        }
    }
}

Thread.init 方法调用此构造器传递 InheritableThreadLocal:

if (inheritThreadLocals && parent.inheritableThreadLocals != null)
            this.inheritableThreadLocals =
                ThreadLocal.createInheritedMap(parent.inheritableThreadLocals);

ThreadLocalMap

值得一提的是, ThreadLocalMap 中使用的是 WeakReference, 当 ThreadLocal 对象不再被外部引用时, 弱引用不会阻止GC因此避免了内存泄露。

static class ThreadLocalMap {

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

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

        /**
         * The table, resized as necessary.
         * table.length MUST always be a power of two.
         */
        private Entry[] table;

        /**
         * The number of entries in the table.
         */
        private int size = 0;

        /**
         * The next size value at which to resize.
         */
        private int threshold; // Default to 0
}

Entry 的 key 始终是 ThreadLocal 对象, 值则是 ThreadLocal 对象绑定的变量副本。

Get 流程

首先来看 ThreadLocalMap.getEntry 方法:

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

利用 table 大小始终为2的整数幂的特点使用位运算找到哈希槽。

若哈希槽中为空或 key 不是当前 ThreadLocal 对象则会调用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;
        if (k == null)
            expungeStaleEntry(i);
        else
            i = nextIndex(i, len);
        e = tab[i];
    }
    return null;
}

ThreadLocalMap 使用开放定址法处理哈希冲突, nextIndex 方法会提供哈希冲突时下一个哈希槽的位置。

private static int nextIndex(int i, int len) {
    return ((i + 1 < len) ? i + 1 : 0);
}

nextIndex 返回下一个位置, 到达末尾后返回第一个位置0.

getEntryAfterMiss 方法会循环查找直到找到或遍历所有可能的哈希槽, 在循环过程中可能遇到4种情况:

• 哈希槽中是当前ThreadLocal, 说明找到了目标
• 哈希槽中为其它ThreadLocal, 需要继续查找
• 哈希槽中为null, 说明搜索结束未找到目标
• 哈希槽中存在Entry, 但是 Entry 中没有 ThreadLocal 对象。因为 Entry 使用弱引用, 这种情况说明 ThreadLocal 被GC回收。

为了处理GC造成的空洞(stale entry), 需要调用expungeStaleEntry方法进行清理。

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

    // 清理当前的空洞
    tab[staleSlot].value = null;
    tab[staleSlot] = null;
    size--;

    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 {
            // 发现不是空洞的 Entry 将其放入最靠前的哈希槽中
            int h = k.threadLocalHashCode & (len - 1);
            if (h != i) {
                tab[i] = null;
                while (tab[h] != null) // 处理移动过程中的哈希冲突
                    h = nextIndex(h, len);
                tab[h] = e;
            }
        }
        // 循环执行直到遇到空的哈希槽, 表明从 staleSlot 开始的查找序列中间不会存在空哈希槽或空Entry
    }
    return i;
}

清理分为两个部分:

1. 首先清理掉空的Entry
2. Entry被清理后可能会使 getEntryAfterMiss 方法误以为搜索已经结束，因此需要将后面的 Entry 进行 rehash 填补空洞

在执行清理时, 可能因为GC造成多个空洞因此需要循环清理。

Set 流程

首先来看 ThreadLocalMap.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;
    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) {
            replaceStaleEntry(key, value, i);
            return;
        }
    }

    tab[i] = new Entry(key, value);
    int sz = ++size;
    if (!cleanSomeSlots(i, sz) && sz >= threshold)
        rehash();
}

首先计算哈希槽的位置, 此时可能有3种情况:

1. 哈希槽为空, 直接将新 Entry 填入槽中; 此外调用 cleanSomeSlots 搜索并清理 GC 造成的空洞; 此外检查 Entry 数量是否到达阈值, 必要时调用 rehash 方法进行扩容。
2. 哈希槽中为当前 ThreadLocal, 直接进行替换
3. 哈希槽中为空 Entry, 说明原有ThreadLocal 被 GC 回收, 调用 replaceStaleEntry 将其替换。

接下来重点分析 replaceStaleEntry:

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

    int slotToExpunge = staleSlot;
    for (int i = prevIndex(staleSlot, len);
            (e = tab[i]) != null;
            i = prevIndex(i, len))
        if (e.get() == null)
            slotToExpunge = i;

    for (int i = nextIndex(staleSlot, len);
            (e = tab[i]) != null;
            i = nextIndex(i, len)) {
        ThreadLocal<?> k = e.get();
        if (k == key) {
            e.value = value;

            tab[i] = tab[staleSlot];
            tab[staleSlot] = e;

            if (slotToExpunge == staleSlot)
                slotToExpunge = i;
            cleanSomeSlots(expungeStaleEntry(slotToExpunge), len);
            return;
        }
        if (k == null && slotToExpunge == staleSlot)
            slotToExpunge = i;
    }

    tab[staleSlot].value = null;
    tab[staleSlot] = new Entry(key, value);

    if (slotToExpunge != staleSlot)
        cleanSomeSlots(expungeStaleEntry(slotToExpunge), len);
}

replaceStaleEntry 方法看上去非常复杂, 简单的说分为三部分:

1. 向前寻找空的 Entry 将其位置写入 slotToExpunge, 这是为了清理不必继续关注
2. 向后进行寻找若是找到与传入的 key 相同 Entry 则更新 Entry 的内容并将其移动到 staleSlot, 然后调用 cleanSomeSlots 进行清理
3. 若最终没有找到 key 相同的Entry, 则在 staleSlot 处写入一个新的 Entry, 调用 cleanSomeSlots 进行清理

cleanSomeSlots 调用 expungeStaleEntry 从位置 i 开始向后清理。

执行log2(n)次清理以取得清理效果(剩余空洞数量)和清理耗时之间的平衡。

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

简单看一下 rehash 的过程:

 private void rehash() {
    expungeStaleEntries();

    if (size >= threshold - threshold / 4)
        resize();
}

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();
            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++;
            }
        }
    }

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

首先进行清理，若清理后sz > thresholde * 0.75将哈希表的的大小翻倍。

Remove

remove 方法和 get 方法比较类似:

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