TransmittableThreadLocal

scj1022

已于 2024-02-21 23:12:40 修改

阅读量573

点赞数 7

分类专栏：并发与多线程源码分析文章标签： java

于 2024-02-18 19:22:50 首次发布

本文链接：https://blog.csdn.net/scj1022/article/details/136157685

版权

并发与多线程同时被 2 个专栏收录

10 篇文章 0 订阅

订阅专栏

源码分析

10 篇文章 0 订阅

订阅专栏

TransmittableThreadLocal

一、概述

1、背景

大多数系统会使用线程池技术，这已经不仅仅是InheritableThreadLocal能够解决的了。

TransmittableThreadLocal 是Alibaba开源，用于实现 在使用线程池等会缓存线程的组件时传递ThreadLocal。

2、使用场景

下面是几个典型场景例子。

分布式跟踪系统
应用容器或上层框架跨应用代码给下层SDK传递信息
日志收集记录系统上下文

3、maven依赖

<!-- TransmittableThreadLocal -->
<dependency>
    <groupId>com.alibaba</groupId>
    <artifactId>transmittable-thread-local</artifactId>
    <version>2.14.2</version>
</dependency>

4、生效前提

当我们使用线程池时，要使得线程池的变量传递起效果，需要：

使用TtlRunnable.get(runnable)对Runnable进行包装

使用TtlCallable.get(callable)对Callable进行包装
使用TtlExecutors.getTtlExecutor对Executor进行包装

使用TtlExecutors.getTtlExecutorService对 ExecutorService进行包装

使用TtlExecutors.getTtlScheduledExecutorService对ScheduledExecutorService进行包装

第二种方式其实就是在 execute 和 submit 时通过 TtlRunnable 或 TtlCallable 对原任务进行包装

至于为什么要这样，可以在下面的源码分析中找到答案。

二、源码分析 - TransmittableThreadLocal

TransmittableThreadLocal 继承自 InheritableThreadLocal

public class TransmittableThreadLocal<T> extends InheritableThreadLocal<T> implements TtlCopier<T> {
    
}

使得当用户利用 new Thread() 创建线程时仍然可以达到传递InheritableThreadLocal 的目的。

1、构造方法

// com.alibaba.ttl.TransmittableThreadLocal

public TransmittableThreadLocal() {
    this(false);
}

public TransmittableThreadLocal(boolean disableIgnoreNullValueSemantics) {
    // 是否忽略null值set，默认false
    this.disableIgnoreNullValueSemantics = disableIgnoreNullValueSemantics;
}

2、holder

TransmittableThreadLocal 相比较 InheritableThreadLocal 很关键的一点改进是引入了holder变量

这样就不必对外暴露 Thread 中的 inheritableThreadLocals，保持 ThreadLocal.ThreadLocalMap 的封装性。

// com.alibaba.ttl.TransmittableThreadLocal

// holder其实就是InheritableThreadLocal的一个匿名内部类
private static final InheritableThreadLocal<WeakHashMap<TransmittableThreadLocal<Object>, ?>> holder =
    new InheritableThreadLocal<WeakHashMap<TransmittableThreadLocal<Object>, ?>>() {

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

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

holder其实就是InheritableThreadLocal的一个匿名内部类。
holder 由 final static 修饰。
holder 存储的是 WeakHashMap，key 是 TransmittableThreadLocal<Object>对象，value 是对应的值
- WeakHashMap的key是WeakReference弱引用，方便回收

3、set 方法

// com.alibaba.ttl.TransmittableThreadLocal

// 调用 set() 方法时，同时处理 holder 中 this 指针
@Override
public final void set(T value) {
    if (!disableIgnoreNullValueSemantics && value == null) {
        // 不忽略null写入，则移除本地线程变量
        remove();
    } else {
        // 调用父类ThreadLocal的set方法
        super.set(value);
        // 添加到 holder 中
        addThisToHolder();
    }
}

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

4、remove 方法

// com.alibaba.ttl.TransmittableThreadLocal

// 调用 remove() 方法时，同时remove holder 中 this 指针
@Override
public final void remove() {
    // 从 holder 中移除
    removeThisFromHolder();
    // 调用父类ThreadLocal的remove方法
    super.remove();
}

private void removeThisFromHolder() {
    holder.get().remove(this);
}

5、get 方法

// com.alibaba.ttl.TransmittableThreadLocal


// 调用 get() 方法时，同时将 this 指针放入 holder
@Override
public final T get() {
    // 调用父类ThreadLocal的get方法
    T value = super.get();
    // 如果允许忽略null，或者value不为null，再次添加到holder
    if (disableIgnoreNullValueSemantics || value != null) 
        addThisToHolder();
    return value;
}

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

三、源码分析 - TtlExecutors

public final class TtlExecutors {

    public static Executor getTtlExecutor(@Nullable Executor executor) {
        if (TtlAgent.isTtlAgentLoaded() || executor == null || executor instanceof TtlEnhanced) {
            return executor;
        }
        // 包装Executor
        return new ExecutorTtlWrapper(executor, true);
    }

    public static ExecutorService getTtlExecutorService(@Nullable ExecutorService executorService) {
        if (TtlAgent.isTtlAgentLoaded() || executorService == null || executorService instanceof TtlEnhanced) {
            return executorService;
        }
        // 包装ExecutorService
        return new ExecutorServiceTtlWrapper(executorService, true);
    }
    
    public static ScheduledExecutorService getTtlScheduledExecutorService(@Nullable ScheduledExecutorService scheduledExecutorService) {
        if (TtlAgent.isTtlAgentLoaded() || scheduledExecutorService == null || scheduledExecutorService instanceof TtlEnhanced) {
            return scheduledExecutorService;
        }
        // 包装ScheduledExecutorService
        return new ScheduledExecutorServiceTtlWrapper(scheduledExecutorService, true);
    }

}

下面看一下几个包装的 Wrapper 的 execute 和 submit 方法

class ExecutorTtlWrapper implements Executor, TtlWrapper<Executor>, TtlEnhanced {
    @Override
    public void execute(@NonNull Runnable command) {
        // 通过TtlRunnable对原Runnable进行包装
        executor.execute(TtlRunnable.get(command, false, idempotent));
    }
    
    // ...
}

class ExecutorServiceTtlWrapper extends ExecutorTtlWrapper implements ExecutorService, TtlEnhanced {
    
    @NonNull
    @Override
    public <T> Future<T> submit(@NonNull Callable<T> task) {
        // 通过TtlCallable对原Callable进行包装
        return executorService.submit(TtlCallable.get(task, false, idempotent));
    }

    @NonNull
    @Override
    public Future<?> submit(@NonNull Runnable task) {
         // 通过TtlRunnable对原Runnable进行包装
        return executorService.submit(TtlRunnable.get(task, false, idempotent));
    }
    
    // ...
    
}

class ScheduledExecutorServiceTtlWrapper extends ExecutorServiceTtlWrapper implements ScheduledExecutorService, TtlEnhanced {

    @NonNull
    @Override
    public ScheduledFuture<?> schedule(@NonNull Runnable command, long delay, @NonNull TimeUnit unit) {
        // 通过TtlRunnable对原Runnable进行包装
        return scheduledExecutorService.schedule(TtlRunnable.get(command, false, idempotent), delay, unit);
    }

    @NonNull
    @Override
    public <V> ScheduledFuture<V> schedule(@NonNull Callable<V> callable, long delay, @NonNull TimeUnit unit) {
        // 通过TtlCallable对原Callable进行包装
        return scheduledExecutorService.schedule(TtlCallable.get(callable, false, idempotent), delay, unit);
    }
    
    // ...

}

可以看到，几个包装的 Wrapper 的 execute 和 submit 方法最后都通过 TtlRunnable 或 TtlCallable 包装了原任务。

四、源码分析 - TtlRunnable

TtlCallable 和 TtlRunnable 差不多，这里以 TtlRunnable 为例分析。

public final class TtlRunnable implements Runnable, TtlWrapper<Runnable>, TtlEnhanced, TtlAttachments {
    
}

1、get 方法

// com.alibaba.ttl.TtlRunnable

public static TtlRunnable get(@Nullable Runnable runnable, boolean releaseTtlValueReferenceAfterRun, boolean idempotent) {
    if (runnable == null) return null;

    if (runnable instanceof TtlEnhanced) {
        // avoid redundant decoration, and ensure idempotency
        if (idempotent) return (TtlRunnable) runnable;
        else throw new IllegalStateException("Already TtlRunnable!");
    }

    // 返回TtlRunnable
    return new TtlRunnable(runnable, releaseTtlValueReferenceAfterRun);
}

private TtlRunnable(@NonNull Runnable runnable, boolean releaseTtlValueReferenceAfterRun) {
    this.capturedRef = new AtomicReference<Object>(capture());
    // 原Runnable
    this.runnable = runnable;
    this.releaseTtlValueReferenceAfterRun = releaseTtlValueReferenceAfterRun;
}

这里的capture()方法调用的是Transmitter中的capture()方法

2、run 方法

// com.alibaba.ttl.TtlRunnable

public void run() {
    // 获取captured（就是Snapshot对象），里面存储了父线程的值
    final Object captured = capturedRef.get();
    if (captured == null || releaseTtlValueReferenceAfterRun && !capturedRef.compareAndSet(captured, null)) {
        throw new IllegalStateException("TTL value reference is released after run!");
    }
    // 传入capture方法捕获的ttl，然后在子线程重放，也就是调用ttl的set方法，
    // 这样就会把值设置到当前的线程中去，最后会把子线程之前存在的ttl返回
    final Object backup = replay(captured);
    try {
        // 调用被包装的 runnable 对象的 run() 方法执行实际的任务。
        runnable.run();
    } finally {
        // 恢复备份
        restore(backup);
    }
}

这里的replay 和 restore 方法调用的也是Transmitter中的replay 和 restore 方法

五、源码分析 - Transmitter

上面说到，TtlRunnable的 get 和 run 方法调用了 Transmitter 中的capture()、replay 和 restore 方法

// com.alibaba.ttl.TransmittableThreadLocal

private static final Set<Transmittee<Object, Object>> transmitteeSet = new CopyOnWriteArraySet<>();

public static class Transmitter {

    @NonNull
    public static Object capture() {
        final HashMap<Transmittee<Object, Object>, Object> transmittee2Value = newHashMap(transmitteeSet.size());

        // 遍历 transmitteeSet
        for (Transmittee<Object, Object> transmittee : transmitteeSet) {
            try {
                // 调用 transmittee 的 capture 方法
                transmittee2Value.put(transmittee, transmittee.capture());
            } catch (Throwable t) {
                // ...
            }
        }

        // 构建 Snapshot
        return new Snapshot(transmittee2Value);
    }

    @NonNull
    public static Object replay(@NonNull Object captured) {
        // captured -> Snapshot
        final Snapshot capturedSnapshot = (Snapshot) captured;

        final HashMap<Transmittee<Object, Object>, Object> transmittee2Value = newHashMap(capturedSnapshot.transmittee2Value.size());

        // 遍历 Snapshot 中的 transmittee2Value
        for (Map.Entry<Transmittee<Object, Object>, Object> entry : capturedSnapshot.transmittee2Value.entrySet()) {
            Transmittee<Object, Object> transmittee = entry.getKey();
            try {
                Object transmitteeCaptured = entry.getValue();
                // 调用 transmittee 的 replay 方法
                transmittee2Value.put(transmittee, transmittee.replay(transmitteeCaptured));
            } catch (Throwable t) {
                // ...
            }
        }

        // 构建新的Snapshot
        return new Snapshot(transmittee2Value);
    }

    public static void restore(@NonNull Object backup) {
        // backup -> Snapshot，遍历 Snapshot 中的 transmittee2Value
        for (Map.Entry<Transmittee<Object, Object>, Object> entry : ((Snapshot) backup).transmittee2Value.entrySet()) {
            Transmittee<Object, Object> transmittee = entry.getKey();
            try {
                Object transmitteeBackup = entry.getValue();
                // 调用 transmittee 的 restore 方法
                transmittee.restore(transmitteeBackup);
            } catch (Throwable t) {
                // ...
            }
        }
    }

    private static class Snapshot {
        // key是Transmittee的实现，value是HashMap（存储了TransmittableThreadLocal及其值）
        final HashMap<Transmittee<Object, Object>, Object> transmittee2Value;

        public Snapshot(HashMap<Transmittee<Object, Object>, Object> transmittee2Value) {
            this.transmittee2Value = transmittee2Value;
        }
    }

}

可以看到 Transmitter 中的 capture()、replay 和 restore 方法，都是在操作 Snapshot

Snapshot 保存着 HashMap<Transmittee<Object, Object>, Object> transmittee2Value
调用的是 Transmittee 的 capture()、replay 和 restore 方法

六、源码分析 - Transmittee

1、Transmittee 接口

// com.alibaba.ttl.TransmittableThreadLocal

public interface Transmittee<C, B> {
    @NonNull
    C capture();

    @NonNull
    B replay(@NonNull C captured);

    @NonNull
    B clear();

    void restore(@NonNull B backup);
}

2、transmitteeSet 怎么来的

// com.alibaba.ttl.TransmittableThreadLocal

private static final Set<Transmittee<Object, Object>> transmitteeSet = new CopyOnWriteArraySet<>();

// 注册 transmitteeSet
static {
    registerTransmittee(ttlTransmittee);
    registerTransmittee(threadLocalTransmittee);
}

public static <C, B> boolean registerTransmittee(@NonNull Transmittee<C, B> transmittee) {
    return transmitteeSet.add((Transmittee<Object, Object>) transmittee);
}

public static <C, B> boolean registerTransmittee(@NonNull Transmittee<C, B> transmittee) {
    return transmitteeSet.add((Transmittee<Object, Object>) transmittee);
}


// 匿名内部类
private static final Transmittee<HashMap<TransmittableThreadLocal<Object>, Object>, HashMap<TransmittableThreadLocal<Object>, Object>> ttlTransmittee =
    new Transmittee<HashMap<TransmittableThreadLocal<Object>, Object>, HashMap<TransmittableThreadLocal<Object>, Object>>() {
    // 实现了Transmittee的几个接口...
}

// 匿名内部类
private static final Transmittee<HashMap<ThreadLocal<Object>, Object>, HashMap<ThreadLocal<Object>, Object>> threadLocalTransmittee =
    new Transmittee<HashMap<ThreadLocal<Object>, Object>, HashMap<ThreadLocal<Object>, Object>>() {
    // 实现了Transmittee的几个接口...
}

3、ttlTransmittee

1）capture

// com.alibaba.ttl.TransmittableThreadLocal.ttlTransmittee

@NonNull
@Override
public HashMap<TransmittableThreadLocal<Object>, Object> capture() {
    final HashMap<TransmittableThreadLocal<Object>, Object> ttl2Value = newHashMap(holder.get().size());
    // 遍历 holder 中的 TransmittableThreadLocal
    for (TransmittableThreadLocal<Object> threadLocal : holder.get().keySet()) {
        // copyValue 实际是调用了 TransmittableThreadLocal 的 get方法
        ttl2Value.put(threadLocal, threadLocal.copyValue());
    }
    return ttl2Value;
}

// com.alibaba.ttl.TransmittableThreadLocal

private T copyValue() {
    return copy(get());
}

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

public T copy(T parentValue) {
    return parentValue;
}

2）replay

// com.alibaba.ttl.TransmittableThreadLocal.ttlTransmittee

@NonNull
@Override
public HashMap<TransmittableThreadLocal<Object>, Object> replay(@NonNull HashMap<TransmittableThreadLocal<Object>, Object> captured) {
    // 备份
    final HashMap<TransmittableThreadLocal<Object>, Object> backup = newHashMap(holder.get().size());

    // 遍历 holder 中的 TransmittableThreadLocal
    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;
}

// com.alibaba.ttl.TransmittableThreadLocal

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

private static void doExecuteCallback(boolean isBefore) {
    // copy TTL Instances to avoid `ConcurrentModificationException`
    // even adjust TTL instances in biz lifecycle callbacks(beforeExecute/afterExecute)
    WeakHashMap<TransmittableThreadLocal<Object>, ?> ttlInstances = new WeakHashMap<TransmittableThreadLocal<Object>, Object>(holder.get());

    for (TransmittableThreadLocal<Object> threadLocal : ttlInstances.keySet()) {
        try {
            if (isBefore) threadLocal.beforeExecute();
            else threadLocal.afterExecute();
        } catch (Throwable t) {
            // ...
        }
    }
}

3）restore

// com.alibaba.ttl.TransmittableThreadLocal.ttlTransmittee

@Override
public void restore(@NonNull HashMap<TransmittableThreadLocal<Object>, Object> backup) {
    // call afterExecute callback
    doExecuteCallback(false);

    // 遍历 holder 中的 TransmittableThreadLocal
    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);
}

// com.alibaba.ttl.TransmittableThreadLocal

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