3 ThreadLocal-手写TransmittableThreadLocal

jkf3333

已于 2022-09-21 23:35:48 修改

阅读量230

点赞数

分类专栏： ThreadLocal 文章标签： ThreadLocal

于 2022-09-19 09:41:14 首次发布

本文链接：https://blog.csdn.net/u010652576/article/details/126925292

版权

ThreadLocal 专栏收录该内容

4 篇文章 0 订阅

订阅专栏

TTL的原理

为了解决InheritableThreadLocal无法处池内线程ThreadLocal传递问题。

线程池ThreadLocal的特殊性

池线程中的ThreadLocal，和父子线程间的ThreadLocal传递有一定的特殊性：

池内线程有自己的ThreadLocal。
a. 线程池可以指定ThreadFactory，通过ThreadFactory创建的Thread，可以设置自己的ThreadLocal。
b. 池内线程创建时，同样是业务线程触发，因此业务线程可以看做池线程的父线程，因此业务线程的InheritableThreadLocal，同样会被复制到池线程中。
不同任务的ThreadLocal存在不一致性。例如：业务线程A的CurrentUserThreadLocal存储的是用户甲，业务线程B的CurrentUserThreadLocal存储的事用户乙。

ThreadLocal的传递时机

线程池中线程的执行过程分析，唯一可以变动ThreadLocal的时机，只有run方法，因为run方法是业务唯一可以控制的代码。当然这也是一句废话，因为池内线程只处理run方法。

因此大致逻辑如下：

提交业务线程时，备份该业务线程的ThreadLocal。
池线程执行时，备份池线程的ThreadLocal。
根据步骤一的备份，恢复备份数据到池线程中，完成业务ThreadLocal到池线程ThreadLocal的传递。
根据步骤二的备份，复原池线程的ThreadLocal。

为什么备份业务线程的ThreadLocal？
因为业务线程可能会消亡，消亡后，是无法获取ThreadLocal对应的值，因此要备份。

为什么要备份池线程的ThreadLocal，执行完成后，还有复原，而不是直接移除掉。
因为线程池的一个拒绝策略是：callerRun，当触发后，其实执行run方法的不再是池线程，而是业务线程本身，如果移除掉，就会导致业务线程后续无法再次获取。

自定义TTL

单一ThreadLocal传递

只需要传递CurrentUser，那么可以自定义如下。

public class CurrentUserRunnable implements Runnable {
    private String currentUser;
    private Runnable runnable;

    @Override
    public void run() {
        try {
            //设置池线程的currentUser
            CurrentUserUtil.set(currentUser);
            //真正的执行run方法
            runnable.run();
        } finally {
            //移除池线程的CurrentUser
            CurrentUserUtil.remove();
        }
    }

    public CurrentUserRunnable(Runnable runnable) {
        this.runnable = runnable;
        //备份 业务线程的CurrentUser
        currentUser = CurrentUserUtil.get();
    }

}

CurrentUserUtil就是对    
private static ThreadLocal<String> currentUserTL = new ThreadLocal<>();
做了包装，提供了静态get、set、remove方法，就不做展示。



运行Demo
public class CurrentUserDemo {
    public static void main(String[] args) {
        CurrentUserUtil.set("jkf");
        ExecutorService executorService = Executors.newFixedThreadPool(1);
        executorService.execute(new CurrentUserRunnable(new Runnable() {
            @Override
            public void run() {
                System.out.println(Thread.currentThread().getName() + ":" + CurrentUserUtil.get());
            }
        }));
    }
}
输出结果：
pool-1-thread-1:jkf

当然上述也存在问题：run方法的finally移除了currentUser，如果触发calllerRun拒绝策略，就会导致业务线程再也无法获取CurrentUser。

稍微的优化

添加ExecutorThreadUtil，包装：

private static ThreadLocal<Boolean> executorThreadTL= new ThreadLocal() {
        protected Boolean initialValue() {
            return false;
        }
    };



修改CurrentUserDemo，创建executorService后，添加ThreadFactory
public class CurrentUserDemo {
    public static void main(String[] args) {
       
        //设置创建ThreadFactory，表明当前线程为池线程。
        ((ThreadPoolExecutor) executorService).setThreadFactory(new ThreadFactory() {
            @Override
            public Thread newThread(Runnable r) {
                Thread t = new Thread(r);
                //表明当前线程是池线程
                ExecutorThreadUtil.set();
                return t;
            }
        });
    }
}

修改CurrentUserRunnable的finall
//当前线程是池线程
            if(ExecutorThreadUtil.isExecutorThread()){
                //移除池线程的CurrentUser
                CurrentUserUtil.remove();
            }

通过ExecutorThreadUtil标明当前线程是否是池线程，当触发callerRun时，会给当前线程的ThreadLocalMap中，添加executorThreadTL，value=false，会稍微占用一些内存。

缺点：上述方案只能针对特定的ThreadLocal，进行处理，例如：需要传递CurrentOrg，那么就需要有CurrentOrgRunnable；如果要传递CurrentUser和CurrentOrg，那么还需要特殊处理，当然上述情况可以提取一个ThreadLocalRunnable，把ThreadLocal数组作为入参。

多ThreadLocal传递

刚才说了单一ThreadLocal传递的方案，但是针对多ThreadLoal传递的情况，需要抽象一个ThreadLocalRunnable，具体优化如下：

public class ThreadLocalRunnable implements Runnable {
	//存储业务的ThreadLocal对应的值。
    private Map<ThreadLocal, Object> tls;
    private Runnable runnable;

    /**
     * 创建任务时，把需要的ThreadLocal全部传递进来
     */
    public ThreadLocalRunnable(Runnable run, List<ThreadLocal> allThreadLocal) {
        this.runnable = run;
        //备份业务的ThreadLocal
        backupThreadLocal(allThreadLocal);
    }

    @Override
    public void run() {
        try {
        	//恢复业务的ThreadLocal值到池线程中
            recoveryThreadLocal();
            //真正的执行run方法
            runnable.run();
        } finally {
        	//移除恢复的ThreadLocal
            removeThreadLocal();
        }
    }

    /**
     * 备份传递进来的ThreadLocal
     */
    private void backupThreadLocal(List<ThreadLocal> allThreadLocal) {
        if (!CollectionUtils.isEmpty(allThreadLocal)) {
            for (ThreadLocal tl : allThreadLocal) {
                //备份 业务线程的Tl的value
                tls.put(tl, tl.get());
            }
        }
    }

    /**
     * 恢复备份的业务ThreadLocal到池线程内
     */
    private void recoveryThreadLocal() {
        if (ExecutorThreadUtil.isExecutorThread()) {
            //池线程，恢复业务线程的ThreadLocal
            for (Map.Entry<ThreadLocal, Object> entry : tls.entrySet()) {
                //此时的Thread，是池线程
                entry.getKey().set(entry.getValue());
            }
        }
    }

    /**
     * 移除池线程中的业务ThreadLocal，防止ThreadLocal数据泄漏
     */
    private void removeThreadLocal() {
        if (ExecutorThreadUtil.isExecutorThread()) {
            //池线程，移除恢复的业务ThreadLocal
            for (Map.Entry<ThreadLocal, Object> entry : tls.entrySet()) {
                entry.getKey().remove();
            }
        }
    } 
}

通过上述优化，解决了多个ThreadLocal传递的问题，同时优化了代码结构，但是最严重的一个问题没有解决，即：run方法执行期间，你不知道需要业务线程的哪些ThreadLocal，尤其是当run方法内部，调用了其他服务，不确定用到哪些，就不确定构造函数的入参，run方法存在获取ThreadLocal对应值为null的风险。

最好是可以直接把业务线程的所有ThreadLocal全部备份，这样就不需要考虑run方法中需要哪些业务线程的ThreadLocal。
即修改ThreadLocalRunnable的backupThreadLocal方法。

private void backupThreadLocal() {
        Map<ThreadLocal, Object> allThreadLocals = Thread.currentThread().threadLocalMap;
        if (!CollectionUtils.isEmpty(allThreadLocals)) {
            for (Map.Entry<ThreadLocal, Object> entry : allThreadLocals.entrySet()) {
                //备份 业务线程的Tl的value
                tls.put(entry.getKey(), entry.getValue());
            }
        }
    }

即创建任务时，直接获取当前线程的所有ThreadLocal，这样就不需要创建任务时，把业务的ThreadLocal作为入参了。

上述backupThreadLocal方法是理想中的好方法，但是Thread的ThreadLocalMap不允许跳过ThreadLocal获取值，因此可以把优化点放在：获取当前线程的所有ThreadLocal

ThreadLocalMap暴露访问

暴露ThreadLoalMap，

继承Thread（直接访问）
通过ThreadLocalMap的源码查看，ThreadLocalMap的访问权限是：default，只允许包内类访问，不允许子类访问，因此通过继承Thread访问ThreadLocalMap行不通。
继承ThreadLocal（间接访问）
通过继承ThreadLocal，重写set、get、remove方法，保存ThreadLocal变量，通过变量访问值，间接的获取当前线程的所有ThreadLocal。

自定义MyTTL

通过上述分析，只能通过继承ThreadLocal，持有所有的ThreadLocal才能实现间接的访问ThreadLocalMap。

public class MyTTL<V> extends ThreadLocal<V> {

    //1. 继承ThreadLocal，当有set/get/remove操作时，更新MyTTL记录（增删）。
    //2. 为了持有所有MyTTL变量，因此需要是静态变量，保证记录是唯一一份。
    //3. Set集合：保证MyTTL唯一，MyTTL变量可能多次set/get/remove操作
    //4. 通过ThreadLocalsHolder持有线程的所有TTL变量，这样可以通过访问threadLocalsHolder
    // 获取线程所有的MyTTL变量。
    //5. 因为MyTTL继承ThreadLocal支持多线程，因此通过ThreadLocal区分不同线程的所有ThreadLocal
    //6. threadLocalsHolder持有了所有线程的所有MyTTL变量
    private static ThreadLocal<Set<MyTTL>> threadLocalsHolder = new ThreadLocal() {
        /**初始值*/
        protected Set<MyTTL> initialValue() {
            return new HashSet<>();
        }
    };

    /**
     * 设置数据，因为是继承ThreadLocal，因此需要先设置父类
     *
     * @param v
     */
    public void set(V v) {
        super.set(v);
        if (v == null) {
            //MyTTL值设置为null，相当于移除MyTTL。
            removeHolder();
        }
        //添加
        add();
    }

    /**
     * 记录ThreadLocal
     */
    private void add() {
        threadLocalsHolder.get().add(this);
    }

    /**
     * 获取数据
     */
    public V get() {
        V v = super.get();
        if (v == null) {
            removeHolder();
        }
        return v;
    }

    private void removeHolder() {
        threadLocalsHolder.get().remove(this);
    }

    /**
     * 生成当前Thread的所有ThreadLocal的备份（快照）
     */
    public static Map<MyTTL, Object> backup() {
        Set<MyTTL> currentTls = threadLocalsHolder.get();
        if (currentTls.isEmpty()) {
            return new HashMap<>();
        }
        Map<MyTTL, Object> backup = new HashMap<>();
        for (MyTTL ttl : currentTls) {
            backup.put(ttl, ttl.get());
        }
        return backup;
    }

    /**
     * 根据ThreadLocals，恢复当前线程
     */
    public static void resetThreadLocal(Map<MyTTL, Object> allThreadLocals) {
        //重建MyTTL
        for (Map.Entry<MyTTL, Object> entry : allThreadLocals.entrySet()) {
            MyTTL tl = entry.getKey();
            //当前线程添加TTL
            tl.set(entry.getValue());
        }
        //移除不存在的MyTTL
        for (MyTTL myTTL : threadLocalsHolder.get()) {
            if (!allThreadLocals.containsKey(myTTL)) {
                threadLocalsHolder.get().remove(myTTL);
            }
        }
    }
}

上述代码还存在一定的缺陷：threadLocalsHolder强应用MyTTL，在ThreadLocal章节中，就已经提到ThreadLocalMap的key是弱引用，防止强引用导致ThreadLocal对象无法回收，而threadLocalsHolder现在就是强引用MyTTL对象，也存在上述问题，因此threadLocalsHolder的Set中key要改为弱引用，但是没有WeakHashSet，只有WeakHashMap，又因为WeakHashMap也能实现set效果，因此直接使用WeakHashMap代替。

优化MyTTL的强引用

public class MyTTL<V> extends ThreadLocal<V> {

    //1. 继承ThreadLocal，当有set/get/remove操作时，更新MyTTL记录（增删）。
    //2. 为了持有所有MyTTL变量，因此需要是静态变量，保证记录是唯一一份。
    //3. Set集合：保证MyTTL唯一，MyTTL变量可能多次set/get/remove操作
    //4. 通过ThreadLocalsHolder持有线程的所有TTL变量，这样可以通过访问threadLocalsHolder
    // 获取线程所有的MyTTL变量。
    //5. 因为MyTTL继承ThreadLocal支持多线程，因此通过ThreadLocal区分不同线程的所有ThreadLocal
    //6. threadLocalsHolder持有了所有线程的所有MyTTL变量
    //7. 为了解决threadLocalsHolder强引用MyTTL对象，导致MyTTL对象无法回收的问题，
    //   修改为弱引用，又因为不存在WeakHashSet，只好用WeakHashMap代替，但是相关的value设置为null
    private static ThreadLocal<Map<MyTTL<?>, ?>> threadLocalsHolder = new ThreadLocal() {
        /**初始值*/
        protected WeakHashMap<MyTTL<?>, ?> initialValue() {
            return new WeakHashMap();
        }
    };

    /**
     * 设置数据，因为是继承ThreadLocal，因此需要先设置父类
     *
     * @param v
     */
    public void set(V v) {
        super.set(v);
        if (v == null) {
            //MyTTL值设置为null，相当于移除MyTTL。
            removeHolder();
        }
        //添加
        add();
    }

    /**
     * 记录ThreadLocal
     */
    private void add() {
        //实现Set效果
        if (!threadLocalsHolder.get().containsKey(this)) {
            threadLocalsHolder.get().put(this, null);
        }
    }

    /**
     * 获取数据
     */
    public V get() {
        V v = super.get();
        if (v == null) {
            removeHolder();
        }
        return v;
    }
    
	 /**
     * 移除
     */
    public void remove() {
        super.remove();
        removeHolder();
    }

    private void removeHolder() {
        threadLocalsHolder.get().remove(this);
    }

    /**
     * 生成当前Thread的所有ThreadLocal的备份（快照）
     */
    public static Map<MyTTL, Object> backup() {
        Map<MyTTL<?>, ?> currentTls = threadLocalsHolder.get();
        if (currentTls.isEmpty()) {
            return new HashMap<>();
        }
        Map<MyTTL, Object> backup = new HashMap<>();
        for (Map.Entry<MyTTL<?>, ?> entry : currentTls.entrySet()) {
            backup.put(entry.getKey(), entry.getKey().get());
        }
        return backup;
    }

    /**
     * 根据ThreadLocals，恢复当前线程
     */
    public static void resetThreadLocal(Map<MyTTL, Object> allThreadLocals) {
        //重建MyTTL
        for (Map.Entry<MyTTL, Object> entry : allThreadLocals.entrySet()) {
            MyTTL tl = entry.getKey();
            //当前线程添加TTL
            tl.set(entry.getValue());
        }
        //移除不存在的MyTTL
        for (MyTTL myTTL : threadLocalsHolder.get().keySet()) {
            if (!allThreadLocals.containsKey(myTTL)) {
                threadLocalsHolder.get().remove(myTTL);
            }
        }
    }
}

优化MyTTL的继承

MyTTL继承自ThreadLocal，ThreadLocal本身是不能处理父子线程的传递问题，而InheritableThreadLocal可以实现父子线程间的传递，因为MyTTL继承修改为InheritableThreadLocal。

public class MyTTL<V> extends InheritableThreadLocal<V> {}

这样TTL既支持父子线程的传递，也支持线程池的线程复用。

MyTTLRunnbale

在ThreadLocal的传递时机中，描述了如何传递的问题，因此还需要特殊的Runnable接口配合MyTTL才能实现传递功能。

public class MyTTLRunnable implements Runnable {
    //实际的run方法
    private Runnable runnable;
    //存储业务线程的所有MyTTL，此时可以看做是快照
    //业务线程后续再添加的MyTTL，run方法方位时，仍然为null
    private Map<MyTTL, Object> bizThreadLocals= new HashMap<>();

    public MyTTLRunnable(Runnable runnable) {
        this.runnable = runnable;
        //备份业务线程的MyTTL
        bizThreadLocals= MyTTL.backup();
    }

    @Override
    public void run() {
        //备份池线程的MyTTL
        Map<MyTTL, Object> currentThreadBackup = MyTTL.backup();
        try {
            //恢复业务MyTTL到池线程中
            MyTTL.resetThreadLocal(bizThreadLocals);
            runnable.run();
        } finally {
            //恢复池线程的MyTTL
            MyTTL.resetThreadLocal(currentThreadBackup);
        }
    }

    public static MyTTLRunnable of(Runnable runnable) {
        return new MyTTLRunnable(runnable);
    }
}

测试

public class MyTTLDemo {
    private static MyTTL<String> currentUser = new MyTTL<>();
    private static ExecutorService executorService = Executors.newFixedThreadPool(1);

    public static void main(String[] args) {
        for (int i = 0; i < 10; i++) {
            //模拟10个请求
            String user = "user_" + i;
            currentUser.set(user);
            executorService.execute(MyTTLRunnable.of(new Runnable() {
                @Override
                public void run() {
                    System.out.println(Thread.currentThread().getName() + ":" + currentUser.get());
                }
            }));
        }
        executorService.shutdown();
    }
}

输出：
pool-1-thread-1:user_0
pool-1-thread-1:user_1
pool-1-thread-1:user_2
pool-1-thread-1:user_3
pool-1-thread-1:user_4
pool-1-thread-1:user_5
pool-1-thread-1:user_6
pool-1-thread-1:user_7
pool-1-thread-1:user_8
pool-1-thread-1:user_9

启动只有一个线程的线程池，输出了不同的请求的用户信息，这就表明MyTTL支持线程池的复用。

优化

上述测试是普通线程池运行MyTTLRunnable，如果需要使用MyTTL，就需要变动普通的Runnable为MyTTLRunnable，改动可能很大，因此可以换个角度：MyTTL线程池，运行普通Runnable。

public class TTLExecutorService extends ThreadPoolExecutor {

    public TTLExecutorService(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue);
    }

    /**
     * 重写execute，内部把Runnable替换为MyTTLRunnable
     */
    public void execute(Runnable command) {
        super.execute(MyTTLRunnable.of(command));
    }
}

继承ThreadPoolExecutor，重新execute方法，这样使用就比较方便。

个人思考

池线程的MyTTL恢复，有两种方案：1 备份恢复。2 用完删除。

备份恢复
先备份池线程的MyTTL，恢复业务MyTTL到池线程，然后再根据备份的MyTTL，恢复池线程的TTL。这样可以完整的保证池线程的MyTTL正确性。
用完删除：恢复业务MyTTL到池线程，完成后，移除池线程的业务MyTTL。
方案一：存在数据copy和替换的问题，不过都是浅拷贝，速度还可控。
方案二存在严重的后果：线程池存在callerRun拒绝策略，如果触发，此时的池线程却是业务线程本身，如果移除了业务MyTTL，就导致业务线程后续无法继续使用。

但是如果能区分出池线程和业务线程，那么方案二是否可行？

业务TTL用完删除

其他与MyTTLRunnable一致
public class MyTTLRunnableTwo implements Runnable {
    @Override
    public void run() {
        try {
            //恢复业务MyTTL到池线程中
            MyTTL.resetThreadLocal(bizThreadLocals);
            runnable.run();
        } finally {
            //删除业务TTL
            removeBizMyTTL(bizThreadLocals);
        }
    }

    private void removeBizMyTTL(Map<MyTTL, Object> bizThreadLocals) {
        if (ExecutorThreadUtil.isExecutorThread()) {
            //只有池线程才会移除
            for (MyTTL myTTL : bizThreadLocals.keySet()) {
                myTTL.remove();
            }
        }
    }   
}

核心是如何ExecutorThreadUtil.isExecutorThread()。线程池是支持ThreadFactory，那么是否可以通过ThreadFactory创建线程时标明。

public class MyTTLRunnableTwoDemo {
    private static MyTTL<String> currentUser = new MyTTL<>();
    private static ExecutorService executorService;

    static {
        executorService = Executors.newFixedThreadPool(1);
        //线程池指定ThreadFactory
        ((ThreadPoolExecutor) executorService).setThreadFactory(new ThreadFactory() {
            @Override
            public Thread newThread(Runnable r) {
                Thread t = new Thread(r);
                //当前线程设置为池线程
                ExecutorThreadUtil.set();
                t.setName("pool_" + t.getName());
                return t;
            }
        });
    }

    public static void main(String[] args) {
        for (int i = 0; i < 10; i++) {
            //模拟10个请求
            String user = "user_" + i;
            currentUser.set(user);
            executorService.execute(MyTTLRunnableTwo.of(new Runnable() {
                @Override
                public void run() {
                    System.out.println(Thread.currentThread().getName() + ":" + currentUser.get());
                }
            }));
        }
        executorService.shutdown();
    }
}

输出：
pool_Thread-0:user_0
pool_Thread-0:user_1
pool_Thread-0:user_2
pool_Thread-0:user_3
pool_Thread-0:user_4
pool_Thread-0:user_5
pool_Thread-0:user_6
pool_Thread-0:user_7
pool_Thread-0:user_8
pool_Thread-0:user_9