什么是 Java 的 TransmittableThreadLocal?

Sherwin.Wei Lv7

什么是 Java 的 TransmittableThreadLocal?

这玩意是阿里开源的一个组件,原生的 ThreadLocal 不支持在线程池中传递本地变量,所以弄了个这个来实现下这个需求

68747470733a2f2f63646e2e6a7364656c6976722e6e65742f67682f79657373696d6964612f63646e5f696d6167652f696d672f696d6167652d32303232303132333136353834313538332e706e67.png

接下来我们来详细探讨一下

缘由

任何一个组件的出现必有其缘由,知其缘由背景才能更深刻地理解它。

我们知道 ThreadLocal 的出现就是为了本地化线程资源,防止不必要的多线程之间的竞争。

在有些场景,当父线程 new 一个子线程的时候,希望把它的 ThreadLocal 继承给子线程。

这时候 InheritableThreadLocal 就来了,它就是为了父子线程传递本地化资源而提出的。

具体的实现是在子线程对象被 new 的时候,即 Thread.init 的时,如果查看到父线程内部有 InheritableThreadLocal 的数据。

那就在子 Thread 初始化的时,把父线程的 InheritableThreadLocal 拷贝给子线程。

Snipaste_2024-05-02_20-18-46.jpg

就这样简单地把父线程的 ThreadLocal 数据传递给子线程了。

但是,这个场景只能发生在 new Thread 的时候!也就是手动创建线程之时!那就有个问题了,在平时我们使用的时候基本用的都是线程池。

那就麻了啊,线程池里面的线程都预创建好了,调用的时候就没法直接用 InheritableThreadLocal 了。

所以就产生了一个需求,如何往线程池内的线程传递 ThreadLocal?,JDK 的类库没这个功能,所以怎么搞?

只能我们自己造轮子了。

如何设计?

需求已经明确了,但是怎么实现呢?

平时我们用线程池的话,比如你要提交任务,则使用代码如下:

1
2
Runnable task = new Runnable....;
executorService.submit(task);

小贴士:以下的 ThreadLocal 泛指线程本地数据,不是指 ThreadLocal 这个类

这时候,我们想着把当前线程的 ThreadLocal 传递给线程池内部将要执行这个 task 的线程。

但此时我们哪知道线程池里面的哪个线程会来执行这个任务?

所以,我们得先把当前线程的 ThreadLocal 保存到这个 task 中。

然后当线程池里的某个线程,比如线程 A 获取这个任务要执行的时候,看看 task 里面是否有存储着的 ThreadLocal 。

如果存着那就把这个 ThreadLocal 放到线程 A 的本地变量里,这样就完成了传递。

然后还有一步,也挺关键的,就是恢复线程池内部执行线程的上下文,也就是该任务执行完毕之后,把任务带来的本地数据给删了,把线程以前的本地数据复原。

68747470733a2f2f63646e2e6a7364656c6976722e6e65742f67682f79657373696d6964612f63646e5f696d6167652f696d672f696d6167652d32303232303132333136353834313538332e706e67.png

设计思路应该已经很明确了吧?来看看具体需要如何实现吧!

如何实现?

把上面的设计简单地、直白地翻译成代码如下:

Snipaste_2024-05-02_20-20-17.jpg

如果你读过我之前分析 ThreadLocal 的文章,应该可以很容易的理解上面的操作。

这样虽然可以实现,但是可操作性太差,耦合性太高。

所以我们得想想怎么优化一下,其实有个设计模式就很合适,那就是装饰器模式。

我们可以自己搞一个 Runnable 类,比如 YesRunnable,然后在 new YesRunnable 的时候,在构造器里面把当前线程的 threadlocal 赋值进去。

然后 run 方法那里也修饰一下,我们直接看看伪代码:

1
2
3
4
5
6
7
8
9
10
11
12
13
public YesRunnable(Runnable runable) {
this.threadlocalCopy = copyFatherThreadlocal();
this.runable = runable;
}
public void run() {
//塞入父threadlocal,并返回当前线程原先threadlocal
Object backup = setThreadlocal(threadlocalCopy);
try {
runable.run();//执行被装饰的任务逻辑
} finally {
restore(backup) ; //复原当前线程的上下文
}
}

使用方式如下:

1
2
3
Runnable task = () -> {...};
YesRunnable yesRunnable = new YesRunnable(task);
executorService.submit(yesRunnable);

你看,这不就实现我们上面的设计了嘛!

不过还有一个点没有揭秘,就是如何实现 copyFatherThreadlocal。

我们如何得知父线程现在到底有哪些 Threadlocal?并且哪些是需要上下文传递的?

所以我们还需要创建一个类来继承 Threadlocal。

比如叫 YesThreadlocal,用它声明的变量就表明需要父子传递的!

1
public class YesThreadlocal<T>  extends ThreadLocal<T>

然后我们需要搞个地方来存储当前父线程上下文用到的所有 YesThreadlocal,这样在 copyFatherThreadlocal的时候我们才好遍历复制对吧?

我们可以搞个 holder 来保存这些 YesThreadlocal ,不过 holder 变量也得线程隔离。

毕竟每个线程所要使用的 YesThreadlocal 都不一样,所以需要用 ThreadLocal 来修饰 holder 。

然后 YesThreadlocal 可能会有很多,我们可以用 set 来保存。

但是为了防止我们搞的这个 holder 造成内存泄漏的风险,我们需要弱引用它,不过没有 WeakHashSet,那我们就用 WeakHashMap 来替代存储。

1
private static final ThreadLocal<WeakHashMap<YesThreadLocal<Object>, ?>> holder = new .....

这样我们就打造了一个变量,它是线程独有的,且又能拿来存储当前线程用到的所有 YesThreadLocal ,便于后面的复制,且又不会造成内存泄漏(弱引用)。

是不是感觉有点暂时理不清?没事,我们继续来看看具体怎么用上这个 hold ,可能会清晰些。

首先我们将需要传递给线程池的本地变量从 ThreadLocal 替换成 YesThreadLocal。

然后重写 set 方法,实现如下:

1
2
3
4
5
6
7
8
9
10
@Override
public final void set(T value) {
super.set(value); //调用 ThreadLocal 的 set
addThisToHolder(); // 把当前的 YesThreadLocal 对象塞入 hold 中。
}
private void addThisToHolder() {
if (!holder.get().containsKey(this)) {
holder.get().put((YesThreadLocal <Object>) this, null);
}
}

你看这样就把所有用到的 YesThreadLocal 塞到 holder 中了,然后再来看看 copyFatherThreadlocal 应该如何实现。

1
2
3
4
5
6
7
private static HashMap<YesThreadLocal <Object>, Object> copyFatherThreadlocal () {
HashMap<YesThreadLocal <Object>, Object> fatherMap = new HashMap<YesThreadLocal <Object>, Object>();
for (YesThreadLocal <Object> threadLocal : YesThreadLocal.holder.get().keySet()) {
fatherMap .put(threadLocal, threadLocal.copyValue());
}
return fatherMap ;
}

逻辑很简单,就是一个 map 遍历拷贝。

我现在用一段话来小结一下,把上面的全部操作联合起来理解,应该会清晰很多。

实现思路小结
1.新建一个 YesThreadLocal 类继承自 ThreadLocal ,用于标识这个修饰的变量需要父子线程拷贝
2.新建一个 YesRunnable 类继承自 Runnable,采用装饰器模式,这样就不用修改原有的 Runnable。在构造阶段复制父线程的 YesThreadLocal 变量赋值给 YesRunnable 的一个成员变量 threadlocalCopy 保存。
3.并修饰 YesRunnable#run 方法,在真正逻辑执行前将 threadlocalCopy 赋值给当前执行线程的上下文,且保存当前线程之前的上下文,在执行完毕之后,再复原此线程的上下文。
4.由于需要在构造的时候复制所有父线程用到的 YesThreadLocal ,因此需要有个 holder 变量来保存所有用到的 YesThreadLocal ,这样在构造的时候才好遍历赋值。
5.并且 holder 变量也需要线程隔离,所以用 ThreadLocal 修饰,并且为了防止 holder 强引用导致内存泄漏,所以用 WeakHashMap 存储。
6.往 holder 添加 YesThreadLocal 的时机就在 YesThreadLocal#set 之时

TransmittableThreadLocal 的实现

这篇只讲 TTL 核心思想(关键路径),由于篇幅原因其它的不作展开,之后再写一篇详细的。

我上面的实现其实就是 TTL 的复制版,如果你理解了上面的实现,那么接下来对 TTL 介绍理解起来应该很简单,相当于复习了。

我们先简单看一下 TTL 的使用方式。

Snipaste_2024-05-02_20-23-26.jpg

使用起来很简单对吧?

TTL 对标上面的 YesThreadLocal ,差别在于它继承的是 InheritableThreadLocal,因为这样直接 new TTL 也会拥有父子线程本地变量的传递能力。

Snipaste_2024-05-02_20-23-52.jpg

我们再来看看 TTL 的 get 和 set 这两个核心操作:

Snipaste_2024-05-02_20-24-08.jpg

可以看到 get 和 set 其实就是复用父类 ThreadLocal 的方法,关键就在于 addThisToHolder,就是我上面分析的将当前使用的 TTL 对象加到 holder 里面。

Snipaste_2024-05-02_20-24-38.jpg

所以,在父线程赋值即执行 set 操作之后,父线程里的 holder 就存储了当前的 TTL 对象了,即上面演示代码的 ttl.set() 操作。

然后重点就移到了TtlRunnable.get 上了,根据上面的理解我们知道这里是要进行一个装饰的操作,这个 get 代码也比较简单,核心就是 new 一个 TtlRunnable 包装了原始的 task。

Snipaste_2024-05-02_20-25-04.jpg

那我们来看一下它的构造方法:

Snipaste_2024-05-02_20-25-35.jpg

这个 capturedRef 其实就是父线程本地变量的拷贝,然后 capture() 其实就等同于copyFatherThreadlocal()

再来看一下 TtlRunnable 装饰的 run 方法:

Snipaste_2024-05-02_20-26-03.jpg

逻辑很清晰的四步骤:

  1. 拿到父类本地变量拷贝
  2. 赋值给当前线程(线程池内的某线程),并保存之前的本地变量
  3. 执行逻辑
  4. 复原当前线程之前的本地变量

我们再来分析一下 capture() 方法,即如何拷贝的。

在 TTL 中是专门定义了一个静态工具类 Transmitter 来实现上面的 capture、 replay、restore 操作。

Snipaste_2024-05-02_20-26-51.jpg

可以看到 capture 的逻辑其实就是返回一个快照,而这个快照就是遍历 holder 获取所有存储在 holder 里面的 TTL ,返回一个新的 map,还是很简单的吧!

这里还有个 captureThreadLocalValues ,这个是为兼容那些无法将 ThreadLocal 类变更至 TTL ,但是又想复制传递 ThreadLocal 的值而使用的,可以先忽略。

我们再来看看 replay,即如何将父类的本地变量赋值给当前线程的。

Snipaste_2024-05-02_20-27-41.jpg

逻辑还是很清晰的,先备份,再拷贝覆盖,最后会返回备份,拷贝覆盖的代码 setTtlValuesTo 很简单:

Snipaste_2024-05-02_20-28-02.jpg

就是 for 循环进行了一波 set ,从这里也可以得知为什么上面需要移除父线程没有的 TTL,因为这里只是进行了 set。如果不 remove 当前线程的本地变量,那就不是完全继承自父线程的本地变量了,可能掺杂着之前的本地变量,也就是不干净了,防止这种干扰,所以还是 remove 了为妙。

最后我们看下 restore 操作:

Snipaste_2024-05-02_20-28-28.jpg

至此想必对 TTL 的原理应该都很清晰了吧!

一些用法

上面我们展示的只是其中一个用法也就是利用 TtlRunnable.get 来包装 Runnable。

TTL 还提供了线程池的修饰方法,即 TtlExecutors,比如可以这样使用:

1
ExecutorService executorService = TtlExecutors.getTtlExecutorService(Executors.newFixedThreadPool(1));

其实原理也很简单,装饰了一下线程池提交任务的方法,里面实现了 TtlRunnable.get 的包装

Snipaste_2024-05-02_20-29-34.jpg
还有一种使用方式更加透明,即利用 Java Agent 来修饰 JDK 的线程池实现类,这种方式在使用上基本就是无感知了。

在 Java 的启动参数加上:-javaagent:path/to/transmittable-thread-local-2.x.y.jar 即可,然后就正常的使用就行,原生的线程池实现类已经悄悄的被改了!

1
2
3
4
TransmittableThreadLocal<String> ttl= new TransmittableThreadLocal<>();
ExecutorService executorService = Executors.newFixedThreadPool(1);
Runnable task = new RunnableTask();
executorService.submit(task);

最后

好了,有关 TTL 的原理和用法解释的都差不多了。

总结下来的核心操作就是 CRR(Capture/Replay/Restore),拷贝快照、重放快照、复原上下文。

可能有些人会疑惑为什么需要复原,线程池的线程每次执行的时候,如果用了 TTL 那执行的线程都会被覆盖上下文,没必要复原对吧?

其实也有人向作者提了这个疑问,回答是:

线程池满了且线程池拒绝策略使用的是『CallerRunsPolicy』,这样执行的线程就变成当前线程了,那肯定是要复原的,不然上下文就没了。
使用ForkJoinPool(包含并行执行Stream与CompletableFuture,底层使用ForkJoinPool)的场景,展开的ForkJoinTask会在调用线程中直接执行。
其实关于 TTL 还有很多细节可以说,不过篇幅有限,细节要说的话得再开一章。不过今天这篇也算把 TTL 的核心思想讲完了。

假设现在有个面试官问你,我要向线程池里面传递 ThreadLocal 怎么实现呀?想必你肯定可以回答出来了~

Comments