调用链跨线程传递ThreadLocal对象

清疚 2023-07-08 13:26 117阅读 0赞

在全链路跟踪框架中，Trace信息的传递功能是基于ThreadLocal的。但实际业务中可能会使用异步调用，这样就会丢失Trace信息，破坏了链路的完整性。

在同一线程中trace信息的传递流程使用代码模拟如下：

ThreadLocal<String> traceContext = new ThreadLocal<>();
    
    String traceId = Tracer.startServer();
    traceContext.set(traceId) //生成trace信息 传入threadlocal
    ...
    Tracer.startClient(traceContext.get()); //从threadlocal获取trace信息
    Tracer.endClient();
    ...
    Tracer.endServer();

那么显然如果是异步线程的话，下一个Span拿不到上一个Span的trace信息，就会造成调用链跟踪断了。那么怎么才能在异步的情况下传递ThreadLocal对象呢。

## InheritableThreadLocal ##

如果仅仅是父子之间传递ThreadLocal对象的话，JDK自身就有实现InheritableThreadLocal。

Thread内部为InheritableThreadLocal开辟了一个单独的ThreadLocalMap。在父线程创建一个子线程的时候，会检查这个ThreadLocalMap是否为空，不为空则会浅拷贝给子线程的ThreadLocalMap。

Thread的init相关逻辑如下：

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

赋值拷贝代码如下：

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

需要注意的是，拷贝为浅拷贝。父子线程的 ThreadLocalMap 内的 key 都指向同一个 InheritableThreadLocal 对象，Value 也指向同一个 Value。子线程的Value更改可以覆盖父线程的Value。

这样一来InheritableThreadLocal让我们可以在父线程创建子线程的时候将ThreadLocal中的值传递给子线程。但在大部分场景下，业务应用都会使用线程池。而在这种复用线程的池化场景中，线程池中的线程和主线程却都不是父子线程的关系，并不能直接使用InheritableThreadLocal。

## Transmittable ThreadLocal ##

Transmittable ThreadLocal是阿里开源的库，继承了InheritableThreadLocal，优化了在使用线程池等会池化复用线程的情况下传递ThreadLocal的使用。

简单来说，有个专门的TtlRunnable和TtlCallable包装类，用于读取原Thread的ThreadLocal对象及值并存于Runnable/Callable中，在执行run或者call方法的时候再将存于Runnable/Callable中的ThreadLocal对象和值读取出来，存入调用run或者call的线程中。

以TtlRunnable为例，构造函数如下:

private final AtomicReference<Object> capturedRef;
    private final Runnable runnable;
    private final boolean releaseTtlValueReferenceAfterRun;
    
    private TtlRunnable(Runnable runnable, boolean releaseTtlValueReferenceAfterRun) {
        //从父类capture复制到本类
        this.capturedRef = new AtomicReference<>(capture());
        this.runnable = runnable; //提交的runnable对象
        this.releaseTtlValueReferenceAfterRun = releaseTtlValueReferenceAfterRun;
    }

capture函数的复制过程如下:

@Nonnull
            public static Object capture() {
                Map<TransmittableThreadLocal<?>, Object> captured = new HashMap<TransmittableThreadLocal<?>, Object>();
                for (TransmittableThreadLocal<?> threadLocal : holder.get().keySet()) {
                    captured.put(threadLocal, threadLocal.copyValue());
                }
                return captured;
            }

其中holder记录了当前 Thread 绑定了哪些 TransmittableThreadLocal 对象。captured保存了父线程ThreadLocal的值。

接着任务提交到线程池，线程开始运行时，取出保存在captured中的父线程ThreadLocal值并重新set。即将父线程值传递到了任务执行时。

@Override
    public void run() {
        Object captured = capturedRef.get();
        if (captured == null || releaseTtlValueReferenceAfterRun && !capturedRef.compareAndSet(captured, null)) {
            throw new IllegalStateException("TTL value reference is released after run!");
        }
    
        Object backup = replay(captured);
        try {
            runnable.run();
        } finally {
            restore(backup);
        }
    }

这样TransmittableThreadLocal就解决了在线程池场景下的ThreadLocal对象传递。整个流程图如下：  
![format_png][]

## 调用链跨线程传递trace信息 ##

有了TransmittableThreadLocal作为基础，调用链跨线程传递trace信息也不再困难，只需将trace信息均存于TransmittableThreadLocal中，使用异步线程池时使用Ttl相关类修饰即可。模拟代码如下：

public void testAsync() {
        ExecutorService ttlExecutorService = TtlExecutors.getTtlExecutorService(executorService);
        String traceId = Tracer.startServer(); //父线程的traceId
        ThreadLocal<String> traceContext = new TransmittableThreadLocal<>(); 
        traceContext.set(traceId); //存入TransmittableThreadLocal
        ttlExecutorService.submit(new Runnable() {
            @Override
            public void run() {
                //runnable执行中获取当前线程的traceId与父线程的traceId一致
                String childTraceId = traceContext.get();
                Assert.assertEquals(childTraceId, traceId);
                Tracer.startClient(traceId);
                Tracer.endClient();
            }
        });
        Tracer.endServer();
    }

## 使用Java Agent植入修饰代码 ##

以上所有使用需要业务代码去改动自己的线程池类，runnable或者callable类。而使用Java Agent实现线程池的传递是透明的，可以做到应用代码无侵入。

Java Agent（Instrumentation）是JDK1.5引入的技术，基于JVM TI机制，使得开发者可以构建一个独立于应用程序的代理（Agent），用来监测和协助运行在 JVM 上的程序，以及替换和修改某些类的定义。开发者可以在一个普通 Java 程序运行时，通过 – javaagent 参数指定一个特定的 jar 文件（包含 Instrumentation 代理）来启动相应的代理程序，植入自己扩展的修饰代码以实现功能。

在TransmittableThreadLocal中，相关Agent的源码分析如下:

//需要通过agent插入Executor类中的某个方法
    private static void updateMethodOfExecutorClass(final CtClass clazz, final CtMethod method) throws NotFoundException, CannotCompileException {
        if (method.getDeclaringClass() != clazz) {
            return;
        }
        //插入的方法需要Public并且非静态
        final int modifiers = method.getModifiers();
        if (!Modifier.isPublic(modifiers) || Modifier.isStatic(modifiers)) {
            return;
        } 
        //获取该方法的参数类型存入parameterTypes数组中
        CtClass[] parameterTypes = method.getParameterTypes();
        StringBuilder insertCode = new StringBuilder(); 
    
        //根据参数类型顺序，进行代码格式化插入
        for (int i = 0; i < parameterTypes.length; i++) {
            CtClass paraType = parameterTypes[i];
            
            //区分Runnable/Callable
            if (RUNNABLE_CLASS_NAME.equals(paraType.getName())) {
                String code = String.format("$%d = %s.get($%d, false, true);", i + 1, TTL_RUNNABLE_CLASS_NAME, i + 1);
                logger.info("insert code before method " + method + " of class " + method.getDeclaringClass().getName() + ": " + code);
                insertCode.append(code);
            } else if (CALLABLE_CLASS_NAME.equals(paraType.getName())) {
                String code = String.format("$%d = %s.get($%d, false, true);", i + 1, TTL_CALLABLE_CLASS_NAME, i + 1);
                logger.info("insert code before method " + method + " of class " + method.getDeclaringClass().getName() + ": " + code);
                insertCode.append(code);
            }
        }
        
        //调用insertBefore()完成代码插入
        if (insertCode.length() > 0) {
            method.insertBefore(insertCode.toString());
        }
    }

将封装好的TransmittableThreadLocal Jar包放在类目录下的某个文件夹下，例如agent，那么只需在启动参数加入：`-javaagent:agent/transmittable-thread-local-xxx.jar`即可完成修饰代码的植入。

[format_png]: https://imgconvert.csdnimg.cn/aHR0cHM6Ly9mcmVkYWwtYmxvZy5vc3MtY24taGFuZ3pob3UuYWxpeXVuY3MuY29tLzIwMTktMTEtMDQtMTI1MDU0LmpwZw?x-oss-process=image/format,png