TransmittableThreadLocal 的反覆與糾纏

TransmittableThreadLocal 相信很多人用過，一個在多執行緒情況下共享執行緒上下文的利器
名字好長，以下簡稱 ttl
本文以之前一個真實專案開發遇到的問題，還原當時從原始碼的角度分析並解決問題的過程
環境

item	version
java	8
springboot	2.2.2.RELEASE
ttl	2.11.4

程式碼如下，主執行緒並行啟複數任務丟給執行緒池處理

點選檢視程式碼

List<ProcDef> defs = createsValidate(procCreate);
List<CompletableFuture<CreateResult>> cfs = Lists.newArrayListWithCapacity(defs.size());
defs.forEach(def -> {
    AbstractTransientVariable variable = procCreate.getProcDefKeyVars().get(def.getProcDefKey());
    CompletableFuture<CreateResult> cf = CompletableFuture.supplyAsync(() -> create(def, variable), threadPoolTaskExecutor)
	    .handle((r, e) -> {
		if (e != null) {
		    expHandle(def.getProcDefKey(), procCreate.getUserId(), variable, e);
		}
		return r;
	    });
    cfs.add(cf);
});

List<CreateResult> results = cfs.stream().map(CompletableFuture::join).filter(Objects::nonNull).collect(Collectors.toList());
if (defs.size() != results.size()) {
    log.error("Process create fail exist: [{}]", JacksonUtil.toJsonString(procCreate));
}

第一行 createsValidate 方法在主執行緒設定了當前使用者的上下文
將使用者資訊放入了 ttl，方便後續子執行緒使用
測試的時候, 執行緒池在處理任務的時候，有時會獲取不到主執行緒 ttl 資訊
很奇怪，之前也是一直這樣使用，為什麼沒有問題
於是本地main方法模擬

點選檢視程式碼

UserContext.set(new ContextUser().setUserId("mycx26"));

IntStream.range(0, 10).forEach(e -> {
    Supplier<Void> supplier = () -> {
	String userId = UserContext.get() != null ? UserContext.getUserId() : null;
	System.out.println(Thread.currentThread().getName() + " get: " + userId);
	return null;
    };
    CompletableFuture.supplyAsync(supplier);
});

Thread.currentThread().join();

這裡主執行緒將使用者資訊放入 ttl，依次將非同步任務丟給執行緒池，任務執行獲取 ttl 並列印

點選檢視程式碼

ForkJoinPool.commonPool-worker-9 get: mycx26
ForkJoinPool.commonPool-worker-6 get: mycx26
ForkJoinPool.commonPool-worker-13 get: mycx26
ForkJoinPool.commonPool-worker-4 get: mycx26
ForkJoinPool.commonPool-worker-11 get: mycx26
ForkJoinPool.commonPool-worker-2 get: mycx26
ForkJoinPool.commonPool-worker-6 get: mycx26
ForkJoinPool.commonPool-worker-15 get: mycx26
ForkJoinPool.commonPool-worker-8 get: mycx26
ForkJoinPool.commonPool-worker-9 get: mycx26

從輸出結果看，各執行緒都拿到了使用者資訊，似乎又沒有問題
將相同的程式碼放到工程的單元測試方法裡跑

點選檢視程式碼

ForkJoinPool.commonPool-worker-10 get: null
ForkJoinPool.commonPool-worker-15 get: null
ForkJoinPool.commonPool-worker-9 get: null
ForkJoinPool.commonPool-worker-1 get: null
ForkJoinPool.commonPool-worker-13 get: null
ForkJoinPool.commonPool-worker-8 get: null
ForkJoinPool.commonPool-worker-3 get: null
ForkJoinPool.commonPool-worker-2 get: null
ForkJoinPool.commonPool-worker-11 get: null
ForkJoinPool.commonPool-worker-15 get: null

結果卻截然相反，到這裡我有點懷疑 ttl 對於多執行緒支援的泛用性了

找到 ttl 的 github readme 閱讀
要保證執行緒池中傳遞值，一種方式是修飾 Runnable 和 Callable，Supplier 也有類似的包裝器
於是修改程式碼重新測試，測試透過

雖然問題是解決了，但是原因卻無從得知，等於還是繞過了問題
下次遇到 ttl 的問題，不知道原理還是無從下手
找到了一個已經 closed 類似的 issue
https://github.com/alibaba/transmittable-thread-local/issues/138
但還是沒有解決我的疑問
沒有辦法，只能看原始碼了，問題還是要一個一個解決

一. main方法沒有修飾的任務為什麼能跨越執行緒池傳遞 ttl

1.1 首先看看 ttl 的 set 方法做了什麼

點選檢視程式碼

public final void set(T value) {
    if (!disableIgnoreNullValueSemantics && null == value) {
        // may set null to remove value
        remove();
    } else {
        super.set(value);
        addThisToHolder();
    }
}

else 走了父類 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);
    }
}

先看一下 ttl 的繼承體系
ttl 繼承 InheritableThreadLocal，InheritableThreadLocal 繼承 ThreadLocal
InheritableThreadLocal 可以讓子執行緒訪問父執行緒設定的本地變數

點選檢視程式碼

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

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

透過重寫 getMap 和 createMap 方法將 ThreadLocal 的維護職責
由 threadLocals 轉移給了 inheritableThreadLocals
threadLocals 和 inheritableThreadLocals 型別一樣
是 ThreadLocal 中的靜態內部類 ThreadLocalMap
為了維護執行緒本地變數定製化的雜湊map, 兩者由 Thread 持有

回到上文 TheadLocal set方法
首先獲取當前執行緒，入參呼叫 getMap 方法獲取當前執行緒的 inheritableThreadLocals

• map不為null
  將 ttl 做為 key，value 作為值，放入當前執行緒的 inheritableThreadLocals
  
  
• map為null
  將 ttl 和 value 構造一個新的 ThreadLocalMap，初始化當前執行緒的 inheritableThreadLocals
  
  

1.2 接下來看 CompletableFuture 的 supplyAsync 方法
這個方法呼叫棧很深，如果多執行緒功力不深，基本看不懂
但這不妨礙排查這個問題
supplyAsync 預設用的 ForkJoinPool 跑任務
那麼必然會啟一個執行緒
即必然會呼叫 Thread 的 init 方法初始化執行緒

首先將斷點加到 CompletableFuture.supplyAsync(supplier); 這行
debug跑起來
然後將斷點加到 Thread init 方法的第一行
（防止jvm啟動初始化的執行緒產生干擾，比如 c2 complier thread）

點選檢視程式碼

Thread parent = currentThread();

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

重點是這裡，判斷父執行緒的 inheritableThreadLocals 如果不為 null
就把父執行緒的 inheritableThreadLocals 複製到子執行緒

1.3 接著看 ttl 的 get 方法

點選檢視程式碼

public final T get() {
    T value = super.get();
    if (disableIgnoreNullValueSemantics || null != value) addThisToHolder();
    return value;
}

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

同樣走的 ThreadLocal 的 get 方法
首先獲取當前執行緒，getMap 獲取其對應的 inheritableThreadLocals
順利拿到之前父執行緒設定的變數
到這裡，第一個問題算是解了

二. 同樣的程式碼跑在單元測試，沒有修飾的任務為什麼不能跨越執行緒池傳遞 ttl

這裡我有理由懷疑是 spring 容器在拉起的時候，提前用到了 ForkJoinPool 的 commonPool
但是專案依賴眾多，如何定位
既然用到了，那麼將斷點加在 ForkJoinPool 啟動執行緒
然後沿著呼叫棧幀一直向上找不就行了
將斷點加在 ForkJoinPool 的 createWorker方法的第一行
開始找

點選檢視程式碼

/* 檢查邏輯刪除欄位只能有最多一個 */
    Assert.isTrue(fieldList.parallelStream().filter(TableFieldInfo::isLogicDelete).count() < 2L,
        String.format("annotation of @TableLogic can't more than one in class : %s.", clazz.getName()));

果然，熟悉的身影，mybatis plus
spring容器拉起時在初始化 SqlSessionFactory 時
會呼叫 TableInfoHelper 的 initTableFields 方法初始化表主鍵和欄位
注意這裡用的 stream 的並行流 parallel stream，很熟悉了
底層預設用的 ForkJoinPool 的 commonPool
那麼在主執行緒設定的 TTL，執行緒池中的執行緒之前已經初始化，當然就拿不到了
好，這是第二個問題

三. 為什麼專案中自定義執行緒池獲取不到前面主執行緒建立的 ttl
和二是相同的問題，執行操作前，執行緒池已經被排程執行任務了
執行緒如果池化，那麼後續在跑非同步任務時就沒有父子執行緒之說了
那麼現在只剩最後一個問題

四. 為什麼專案中任務加了包裝器後又拿到了

點選檢視程式碼

TtlWrappers.wrap(() -> create(def, variable))

沒有什麼辦法，跟進去吧

4.1 看看 TtlWrappers 的靜態方法 wrap 做了什麼

點選檢視程式碼

public static <T> Supplier<T> wrap(@Nullable Supplier<T> supplier) {
    if (supplier == null) return null;
    else if (supplier instanceof TtlEnhanced) return supplier;
    else return new TtlSupplier<T>(supplier);
}

看樣子，大概是想用裝飾模式包裝 Supplier 為 TTL wrapper
new 了一個 TtlSupplier，這是 TtlWrappers 的一個靜態內部類
繼續進去

點選檢視程式碼

TtlSupplier(@NonNull Supplier<T> supplier) {
    this.supplier = supplier;
    this.capture = capture();
}

supplier完成賦值後，重點是後面的 capture

點選檢視程式碼

/**
 * Capture all {@link TransmittableThreadLocal} and registered {@link ThreadLocal} values in the current thread.
 *
 * @return the captured {@link TransmittableThreadLocal} values
 * @since 2.3.0
 */
@NonNull
public static Object capture() {
    return new Snapshot(captureTtlValues(), captureThreadLocalValues());
}

capture 靜態方法位於 ttl 的靜態內部類 Transmitter 中
註釋很清晰，捕獲當前執行緒的所有 ttl 和 ThreadLocal 的值
new Snapshort 繼續跟進去

點選檢視程式碼

private static class Snapshot {
    final WeakHashMap<TransmittableThreadLocal<Object>, Object> ttl2Value;
    final WeakHashMap<ThreadLocal<Object>, Object> threadLocal2Value;

    private Snapshot(WeakHashMap<TransmittableThreadLocal<Object>, Object> ttl2Value, WeakHashMap<ThreadLocal<Object>, Object> threadLocal2Value) {
        this.ttl2Value = ttl2Value;
        this.threadLocal2Value = threadLocal2Value;
    }
}

Snaphost 同樣是 ttl 的靜態內部類
構造方法的第一個引數方法 captureTtlValues 跟進去

點選檢視程式碼

private static WeakHashMap<TransmittableThreadLocal<Object>, Object> captureTtlValues() {
    WeakHashMap<TransmittableThreadLocal<Object>, Object> ttl2Value = new WeakHashMap<TransmittableThreadLocal<Object>, Object>();
    for (TransmittableThreadLocal<Object> threadLocal : holder.get().keySet()) {
        ttl2Value.put(threadLocal, threadLocal.copyValue());
    }
    return ttl2Value;
}

同樣來自 Transmitter
好，程式碼並不複雜，重點是 holder

點選檢視程式碼

// Note about holder:
// 1. holder self is a InheritableThreadLocal(a *ThreadLocal*).
// 2. The type of value in holder is WeakHashMap<TransmittableThreadLocal<Object>, ?>.
//    2.1 but the WeakHashMap is used as a *Set*:
//        - the value of WeakHashMap is *always null,
//        - and never be used.
//    2.2 WeakHashMap support *null* value.
private static InheritableThreadLocal<WeakHashMap<TransmittableThreadLocal<Object>, ?>> holder =
        new InheritableThreadLocal<WeakHashMap<TransmittableThreadLocal<Object>, ?>>() {
            @Override
            protected WeakHashMap<TransmittableThreadLocal<Object>, ?> initialValue() {
                return new WeakHashMap<TransmittableThreadLocal<Object>, Object>();
            }

            @Override
            protected WeakHashMap<TransmittableThreadLocal<Object>, ?> childValue(WeakHashMap<TransmittableThreadLocal<Object>, ?> parentValue) {
                return new WeakHashMap<TransmittableThreadLocal<Object>, Object>(parentValue);
            }
        };

holder 為 ttl 的靜態成員變數，型別為 InheritableThreadLocal 的匿名內部類
重寫了 initialValue 和 childValue 方法
再看註釋
這裡 value 的 type 是 WeakHashMap 並且這個 map 被當作 set 用了
還記得上文分析 ttl 的 set 方法嗎，有一塊沒有講
對，就是 else 的 addThisToHolder

點選檢視程式碼

private void addThisToHolder() {
    if (!holder.get().containsKey(this)) {
        holder.get().put((TransmittableThreadLocal<Object>) this, null); // WeakHashMap supports null value.
    }
}

set方法
第一步將 ttl 做為 key，value 作為值，放入當前執行緒的 inheritableThreadLocals
第二步 addThisToHolder
holder.get()方法獲取當前執行緒的 inheritableThreadLocals 變數
key為 holder本身，這裡不要亂了
第一次肯定是拿不到的，那麼這裡為什麼沒有npe？
上文提到他重寫了 initialValue 方法
繼續
首先判斷當前執行緒的 inheritableThreadLocals 是否包含 holder
第一次肯定沒有
那麼把 ttl 本身作為key, 放入當前執行緒的 inheritableThreadLocals 維持的 map 中

至此 TtlSupplier 的 capture 屬性已經持有了主執行緒的所有 ttl 快照

4.2 接下來看 TtlSupplier 重寫的 get 方法

這是核心的行為，可以斷定，其必然做了增強

點選檢視程式碼

public T get() {
    final Object backup = replay(capture);
    try {
        return supplier.get();
    } finally {
        restore(backup);
    }
}

結構很清晰，先replay，再執行核心行為，最後restore
replay 跟進去

點選檢視程式碼

/**
 * Replay the captured {@link TransmittableThreadLocal} and registered {@link ThreadLocal} values from {@link #capture()},
 * and return the backup {@link TransmittableThreadLocal} values in the current thread before replay.
 *
 * @param captured captured {@link TransmittableThreadLocal} values from other thread from {@link #capture()}
 * @return the backup {@link TransmittableThreadLocal} values before replay
 * @see #capture()
 * @since 2.3.0
 */
@NonNull
public static Object replay(@NonNull Object captured) {
    final Snapshot capturedSnapshot = (Snapshot) captured;
    return new Snapshot(replayTtlValues(capturedSnapshot.ttl2Value), replayThreadLocalValues(capturedSnapshot.threadLocal2Value));
}

同樣是位於 ttl 的靜態內部類 Transmitter 的靜態方法
上文 Tramsmitter capture()方法捕獲的主執行緒快照這裡用到了
replayTtlValues 方法跟進去

點選檢視程式碼

private static WeakHashMap<TransmittableThreadLocal<Object>, Object> replayTtlValues(@NonNull WeakHashMap<TransmittableThreadLocal<Object>, Object> captured) {
    WeakHashMap<TransmittableThreadLocal<Object>, Object> backup = new WeakHashMap<TransmittableThreadLocal<Object>, Object>();

    for (final Iterator<TransmittableThreadLocal<Object>> iterator = holder.get().keySet().iterator(); iterator.hasNext(); ) {
        TransmittableThreadLocal<Object> threadLocal = iterator.next();

        // backup
        backup.put(threadLocal, threadLocal.get());

        // clear the TTL values that is not in captured
        // avoid the extra TTL values after replay when run task
        if (!captured.containsKey(threadLocal)) {
            iterator.remove();
            threadLocal.superRemove();
        }
    }

    // set TTL values to captured
    setTtlValuesTo(captured);

    // call beforeExecute callback
    doExecuteCallback(true);

    return backup;
}

private static void setTtlValuesTo(@NonNull WeakHashMap<TransmittableThreadLocal<Object>, Object> ttlValues) {
    for (Map.Entry<TransmittableThreadLocal<Object>, Object> entry : ttlValues.entrySet()) {
        TransmittableThreadLocal<Object> threadLocal = entry.getKey();
        threadLocal.set(entry.getValue());
    }
}

注意這裡已經到子執行緒了
首先透過 holder 獲取當前執行緒的 inheritableThreadLocals 變數
很可能沒有
但是如果執行緒之前已經池化用完沒有remove，這裡是有的
遍歷 map 的 key ttl
這裡為了不汙染子執行緒上下文，先做了備份
對於快照中不包含的 ttl 資訊依次 remove
然後遍歷快照資訊設定到當前執行緒的 inheritableThreadLocals
doExecuteCallback 方法是 ttl 為開發者留的一個勾子方法
時機在任務執行前
最後返回子執行緒 ttl 備份

再回到 TtlSupplier 的 get 方法
supplier 的 get 方法執行任務
最後還剩 restore，傳入上面子執行緒的 ttl 備份

點選檢視程式碼

/**
 * Restore the backup {@link TransmittableThreadLocal} and
 * registered {@link ThreadLocal} values from {@link #replay(Object)}/{@link #clear()}.
 *
 * @param backup the backup {@link TransmittableThreadLocal} values from {@link #replay(Object)}/{@link #clear()}
 * @see #replay(Object)
 * @see #clear()
 * @since 2.3.0
 */
public static void restore(@NonNull Object backup) {
    final Snapshot backupSnapshot = (Snapshot) backup;
    restoreTtlValues(backupSnapshot.ttl2Value);
    restoreThreadLocalValues(backupSnapshot.threadLocal2Value);
}

private static void restoreTtlValues(@NonNull WeakHashMap<TransmittableThreadLocal<Object>, Object> backup) {
    // call afterExecute callback
    doExecuteCallback(false);

    for (final Iterator<TransmittableThreadLocal<Object>> iterator = holder.get().keySet().iterator(); iterator.hasNext(); ) {
        TransmittableThreadLocal<Object> threadLocal = iterator.next();

        // clear the TTL values that is not in backup
        // avoid the extra TTL values after restore
        if (!backup.containsKey(threadLocal)) {
            iterator.remove();
            threadLocal.superRemove();
        }
    }

    // restore TTL values
    setTtlValuesTo(backup);
}

主要看 restoreTtlValues 方法
doExecuteCallback 和上面邏輯類似
區別在時機在任務執行後
holder 獲取子執行緒的 inheritableThreadLocals 變數
遍歷 map 的 key ttl
對於不在備份的 ttl 全部刪除
最後恢復子執行緒的 ttl

彷彿一切沒有發生過

至此最後一個問題解決

你對的不一定對，你錯了一定是錯了
原始碼面前，沒有什麼秘密可言了