手撕一个解决ThreadLocal数据跨线程传递的实用工具

大家好，老吕在前面的文章中提到了如何解决不同场景下ThreadLocal数据传递问题，并且给出了阿里开源TTL框架和手写组件的方案。后来老吕还是觉得TTL的方案有点复杂，手写那个组件不够通用，所以就做了改进，使读写数据通用化。

设计比较简单实用，一个类就搞定了，思路还是通过代理Runnable和Callable接口来增加数据中转的能力，我一直认为这个思路是最划算的。本次加强了数据中转的通用能力。

一、主代码解析

package com.dhy.common.lib.taskproxy;
import lombok.extern.slf4j.Slf4j;


import java.lang.reflect.Method;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.Callable;


/**
 * https://github.com/lvaolin/dhy-common-lib/tree/main/dhy-common-lib-taskproxy
 *
 * 代理类：增加了ThreadLocal上绑定的数据中转功能
 */
@Slf4j
public class TaskProxy implements Runnable, Callable {


    /**
     * 被代理没有返回值的任务
     */
    private Runnable runnable;
    /**
     * 被代理有返回值的任务
     */
    private Callable callable;
    /**
     * 数据临时中转站
     */
    private Map localData = new HashMap<>();
    /**
     * ThreadLocal 工具类
     */
    private Class threadLocalHolder;




    public TaskProxy(Runnable runnable, Class threadLocalHolder){
        this.runnable = runnable;
        this.threadLocalHolder = threadLocalHolder;
        storeThreadLocal();
    }
    public TaskProxy(Callable callable, Class threadLocalHolder){
        this.callable = callable;
        this.threadLocalHolder = threadLocalHolder;
        storeThreadLocal();
    }


    @Override
    public void run() {
        restoreThreadLocal();
        this.runnable.run();
        clearThreadLocal();
    }


    @Override
    public Object call() throws Exception {
        restoreThreadLocal();
        V v = this.callable.call();
        clearThreadLocal();
        return v;
    }


    private void storeThreadLocal() {
        Method[] methods = threadLocalHolder.getDeclaredMethods();
        for (Method method : methods) {
            if (method.getName().startsWith("get")) {
                storeField(method);
            }
        }
    }


    private void storeField(Method method) {
        try {
            Object result = method.invoke(null, null);
            localData.put(method.getName(),result);
            log.info(method.getName()+" invoke");
        } catch (Throwable t){
            log.error(t.getMessage(),t);
        }
    }


    private void restoreThreadLocal() {
        Method[] methods = threadLocalHolder.getDeclaredMethods();
        for (Method method : methods) {
            if (method.getName().startsWith("set")) {
                restoreField(method);
            }
        }


    }


    private void restoreField(Method method) {
        try {
            Object filedValue = localData.get(method.getName().replaceFirst("s", "g"));
            method.invoke(null, filedValue);
            log.info(method.getName()+" invoke");
        } catch (Throwable t){
            log.error(t.getMessage(),t);
        }
    }


    private void clearThreadLocal() {
        Method[] methods = threadLocalHolder.getDeclaredMethods();
        for (Method method : methods) {
            if (method.getName().startsWith("remove")) {
                try {
                    method.invoke(null, null);
                    log.info(method.getName()+" invoke");
                } catch (Throwable t){
                    log.error(t.getMessage(),t);
                }
            }
        }
    }


}

二、使用说明

//使用 TaskProxy 类对 Runnable、Callable接口进行一个修饰代理即可
    
    Runnable runnable = (Runnable)new TaskProxy(new Runnable() {
                  @Override
                  public void run() {
                      System.out.println("使用代理Runnable："+ ThreadLocalHolder.getUserId());
                      System.out.println("使用代理Runnable："+ ThreadLocalHolder.getAppContext().getSessionId());
                  }
              },ThreadLocalHolder.class);
   
    executorService.submit(runnable);          
     
    //-----------------------------------------
    
    Callable callable = (Callable) new DhyTaskProxy(new Callable() {
                @Override
                public String call() throws Exception {
                    System.out.println("使用代理Callable："+ ThreadLocalHolder.getUserId());
                    System.out.println("使用代理Callable："+ ThreadLocalHolder.getAppContext());
                    return "ok";
                }
             }, ThreadLocalHolder.class);
             
     executorService.submit(callable);

ThreadLocal工具类如下

/**
 * ThreadLocal工具类
 */
public class ThreadLocalHolder {


    private static ThreadLocal userIdThreadLocal = new ThreadLocal<>();
    private static ThreadLocal traceIdThreadLocal = new ThreadLocal<>();
    private static ThreadLocal appContextThreadLocal = new ThreadLocal<>();


    public static void setUserId(String userId){
        userIdThreadLocal.set(userId);
    }
    public static void setTraceId(String traceId){
        traceIdThreadLocal.set(traceId);
    }
    public static void setAppContext(RequestContext appContext){
        appContextThreadLocal.set(appContext);
    }


    public static String getUserId(){
        return userIdThreadLocal.get();
    }
    public static String getTraceId(){
        return traceIdThreadLocal.get();
    }


    public static RequestContext getAppContext(){
        return appContextThreadLocal.get();
    }


    public static void removeUserId(){
        userIdThreadLocal.remove();
    }


    public static void removeTraceId(){
        traceIdThreadLocal.remove();
    }


    public static void removeAppContext(){
        appContextThreadLocal.remove();
    }
}

三、注意事项

1、代理的创建时机，每次使用之前要重新创建代理，不能重用，这个在阿里TTL框架中也提到这个了，原理本质上是从上下文中取数据，所以你的上下文变了后要重新生成代理。

2、本组件依靠ThreadLocal对象holder的static类型的get、set、remove方法来获取、设置、销毁数据的，所以你需要按照示例中的ThreadLocalHolder来写自己的ThreadLocalHolder，这是需要注意的。