調用鏈跨線程傳遞ThreadLocal對象

在全鏈路跟蹤框架中,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對象傳遞。整個流程圖如下:

調用鏈跨線程傳遞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即可完成修飾代碼的植入。

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章