通过ThreadLocal实现一个上下文管理组件

背景

  • https://blog.csdn.net/qq_39687472/article/details/129213948
    上文中已经解析了ThreadLocal原理以及应用场景
  • 本文基于ThreadLocal原理,实现了一个上下文状态管理组件Scope,通过开启一个自定义的Scope,在Scope范围内,可以通过Scope各个方法读写数据;
  • 通过自定义线程池实现上下文状态数据的线程间传递;
  • 提出了一种基于FilterScopeRequest粒度的上下文管理方案。

文章源码地址:
github:https://github.com/pengchengSU/demo-request-scope

自定义上下文Scope

在工作中,我们经常需要维护一些上下文,这样可以避免在方法调用过程中传入过多的参数,需要查询/修改一些数据的时候,直接在当前上下文中操作就行了。举个具体点的例子:当web服务器收到一个请求时,需要解析当前登录态的用户,在后续的业务执行流程中都需要这个用户名。

如果只需要维护一个上下文状态数据还比较好处理,可以通过方法传参的形式,执行每个业务方法的时候都通过添加一个表示用户名方法参数传递进去,但是如果需要维护上下文状态数据比较多的话,这个方式就不太优雅了。

一个可行的方案是通过Threadlocal实现请求线程的上下文,只要是同一线程的执行过程,不同方法间不传递上下文状态变量,直接操作ThreadLocal对象实现状态数据的读写。当存在多个上下文状态的话,则需要维护多个ThreadLocal,似乎也可以勉强接受。但是当遇到业务流程中使用线程池的情况下,从Tomcat传递这些ThreadLocal到线程池中的线程里就变的比较麻烦了。

基于以上考虑,下面介绍一种基于Threadlocal实现的上下文管理组件Scope

Scope.java

public class Scope {

    // 静态变量,维护不同线程的上下文Scope
    private static final ThreadLocal<Scope> SCOPE_THREAD_LOCAL = new ThreadLocal<>();

    // 实例变量,维护每个上下文中所有的状态数据,为了区分不同的状态数据,使用ScopeKey类型的实例作为key
    private final ConcurrentMap<ScopeKey<?>, Object> values = new ConcurrentHashMap<>();

    // 获取当前上下文
    public static Scope getCurrentScope() {
        return SCOPE_THREAD_LOCAL.get();
    }

    // 在当前上下文设置一个状态数据
    public <T> void set(ScopeKey<T> key, T value) {
        if (value != null) {
            values.put(key, value);
        } else {
            values.remove(key);
        }
    }

    // 在当前上下文读取一个状态数据
    public <T> T get(ScopeKey<T> key) {
        T value = (T) values.get(key);
        if (value == null && key.initializer() != null) {
            value = key.initializer().get();
        }
        return value;
    }

    // 开启一个上下文
    public static Scope beginScope() {
        Scope scope = SCOPE_THREAD_LOCAL.get();
        if (scope != null) {
            throw new IllegalStateException("start a scope in an exist scope.");
        }
        scope = new Scope();
        SCOPE_THREAD_LOCAL.set(scope);
        return scope;
    }

    // 关闭当前上下文
    public static void endScope() {
        SCOPE_THREAD_LOCAL.remove();
    }
}

ScopeKey.java

public final class ScopeKey<T> {

    // 初始化器,参考 ThreadLocal 的 withInitial()
    private final Supplier<T> initializer;

    public ScopeKey() {
        this(null);
    }

    public ScopeKey(Supplier<T> initializer) {
        this.initializer = initializer;
    }

    // 统一初始化所有线程的 ScopeKey 对应的值,参考 ThreadLocal 的 withInitial()
    public static <T> ScopeKey<T> withInitial(Supplier<T> initializer) {
        return new ScopeKey<>(initializer);
    }

    public Supplier<T> initializer() {
        return this.initializer;
    }

    // 获取当前上下文中 ScopeKey 对应的变量
    public T get() {
        Scope currentScope = getCurrentScope();
        return currentScope.get(this);
    }

    // 设置当前上下文中 ScopeKey 对应的变量
    public boolean set(T value) {
        Scope currentScope = getCurrentScope();
        if (currentScope != null) {
            currentScope.set(this, value);
            return true;
        } else {
            return false;
        }
    }
}

使用方式

@Test
public void testScopeKey() {
    ScopeKey<String> localThreadName = new ScopeKey<>();

    // 不同线程中执行时,开启独占的 Scope
    Runnable r = () -> {
        // 开启 Scope
        Scope.beginScope();
        try {
            String currentThreadName = Thread.currentThread().getName();
            localThreadName.set(currentThreadName);
            log.info("currentThread: {}", localThreadName.get());
        } finally {
            // 关闭 Scope
            Scope.endScope();
        }
    };

    new Thread(r, "thread-1").start();
    new Thread(r, "thread-2").start();

    /** 执行结果
     * [thread-1] INFO com.example.demo.testscope.TestScope - currentThread: thread-1
     * [thread-2] INFO com.example.demo.testscope.TestScope - currentThread: thread-2
     */
}

@Test
public void testWithInitial() {
    ScopeKey<String> initValue = ScopeKey.withInitial(() -> "initVal");

    Runnable r = () -> {
        Scope.beginScope();
        try {
            log.info("initValue: {}", initValue.get());
        } finally {
            Scope.endScope();
        }
    };

    new Thread(r, "thread-1").start();
    new Thread(r, "thread-2").start();

    /** 执行结果
     * [thread-1] INFO com.example.demo.testscope.TestScope - initValue: initVal
     * [thread-2] INFO com.example.demo.testscope.TestScope - initValue: initVal
     */
}

上面的测试用例中在代码中手动开启和关闭Scope不太优雅,可以在Scope中添加两个静态方法包装下RunnableSupplier接口:

//开启一个新的 Scope 执行 Runnable
public static <X extends Throwable> void runWithNewScope(@Nonnull ThrowableRunnable<X> runnable)
        throws X {
    supplyWithNewScope(() -> {
        runnable.run();
        return null;
    });
}

//开启一个新的 Scope 执行 Supplier
public static <T, X extends Throwable> T
        supplyWithNewScope(@Nonnull ThrowableSupplier<T, X> supplier) throws X {
    beginScope();
    try {
        return supplier.get();
    } finally {
        endScope();
    }
}

两个接口

@FunctionalInterface
public interface ThrowableRunnable<X extends Throwable> {
    void run() throws X;
}

@FunctionalInterface
public interface ThrowableSupplier<T, X extends Throwable> {
    T get() throws X;
}

以新的Scope执行,可以这样写:

@Test
public void testRunWithNewScope() {
    ScopeKey<String> localThreadName = new ScopeKey<>();

    ThrowableRunnable r = () -> {
        String currentThreadName = Thread.currentThread().getName();
        localThreadName.set(currentThreadName);
        log.info("currentThread: {}", localThreadName.get());
    };

    // 不同线程中执行时,开启独占的 Scope
    new Thread(() -> Scope.runWithNewScope(r), "thread-1").start();
    new Thread(() -> Scope.runWithNewScope(r), "thread-2").start();

    /** 执行结果
     * [thread-2] INFO com.example.demo.TestScope.testscope - currentThread: thread-2
     * [thread-1] INFO com.example.demo.TestScope.testscope - currentThread: thread-1
     */
}

在线程池中传递Scope

在上一节中实现的Scope,通过ThreadLocal实现了了一个自定义的上下文组件,在同一个线程中通过ScopeKey.set() / ScopeKey.get()读写同一个上下文中的状态数据。

现在需要实现这样一个功能,在一个线程执行过程中开启了一个Scope,随后使用线程池执行任务,要求在线程池中也能获取当前Scope中的状态数据。典型的使用场景是:服务收到一个用户请求,通过Scope将登陆态数据存到当前线程的上下文中,随后使用线程池执行一些耗时的操作,希望线程池中的线程也能拿到Scope中的登陆态数据。
由于线程池中的线程和请求线程不是一个线程,按照目前的实现,线程池中的线程是无法拿到请求线程上下文中的数据的。
解决方法是,在提交runnable时,将当前的Scope引用存到runnable对象中,当获得线程执行时,将Scope替换到执行线程中,执行完成后,再恢复现场。在Scope中新增如下静态方法:

// 以给定的上下文执行 Runnable
public static <X extends Throwable> void runWithExistScope(Scope scope, ThrowableRunnable<X> runnable) throws X {
    supplyWithExistScope(scope, () -> {
        runnable.run();
        return null;
    });
}

// 以给定的上下文执行 Supplier
public static <T, X extends Throwable> T supplyWithExistScope(Scope scope, ThrowableSupplier<T, X> supplier) throws X {
    // 保留现场
    Scope oldScope = SCOPE_THREAD_LOCAL.get();
    // 替换成外部传入的 Scope
    SCOPE_THREAD_LOCAL.set(scope);
    try {
        return supplier.get();
    } finally {
        if (oldScope != null) {
            // 恢复线程
            SCOPE_THREAD_LOCAL.set(oldScope);
        } else {
            SCOPE_THREAD_LOCAL.remove();
        }
    }
}

实现支持Scope切换的自定义线程池ScopeThreadPoolExecutor

public class ScopeThreadPoolExecutor extends ThreadPoolExecutor {

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

    public static ScopeThreadPoolExecutor newFixedThreadPool(int nThreads) {
        return new ScopeThreadPoolExecutor(nThreads, nThreads, 0L, TimeUnit.MILLISECONDS,
                new LinkedBlockingQueue<Runnable>());
    }

    /**
     * 只要override这一个方法就可以
     * 所有submit, invokeAll等方法都会代理到这里来
     */
    @Override
    public void execute(Runnable command) {
        Scope scope = getCurrentScope();
        // 提交任务时,把执行 execute 方法的线程中的 Scope 传进去
        super.execute(() -> runWithExistScope(scope, command::run));
    }
}

测试下ScopeThreadPoolExecutor是否生效:

@Test
public void testScopeThreadPoolExecutor() {
    ScopeKey<String> localVariable = new ScopeKey<>();
    Scope.beginScope();

    try {
        localVariable.set("value out of thread pool");
        Runnable r = () -> log.info("localVariable in thread pool: {}", localVariable.get());

        // 使用线程池执行,能获取到外部Scope中的数据
        ExecutorService executor = ScopeThreadPoolExecutor.newFixedThreadPool(10);
        executor.execute(r);
        executor.submit(r);

    } finally {
        Scope.endScope();
    }

    /** 执行结果
     * [pool-1-thread-1] INFO com.example.demo.testscope.TestScope - localVariable in thread pool: value out of thread pool
     * [pool-1-thread-2] INFO com.example.demo.testscope.TestScope - localVariable in thread pool: value out of thread pool
     */
}

@Test
public void testScopeThreadPoolExecutor2() {
    ScopeKey<String> localVariable = new ScopeKey<>();
    Scope.runWithNewScope(() -> {
        localVariable.set("value out of thread pool");
        Runnable r = () -> log.info("localVariable in thread pool: {}", localVariable.get());

        // 使用线程池执行,能获取到外部Scope中的数据
        ExecutorService executor = ScopeThreadPoolExecutor.newFixedThreadPool(10);
        executor.execute(r);
        executor.submit(r);
    });

    /** 执行结果
     * [pool-1-thread-2] INFO com.example.demo.testscope.TestScope - localVariable in thread pool: value out of thread pool
     * [pool-1-thread-1] INFO com.example.demo.testscope.TestScope - localVariable in thread pool: value out of thread pool
     */
}

以上两个测试用例,分别通过手动开启Scope、借助runWithNewScope工具方法自动开启Scope两种方式验证了自定义线程池ScopeThreadPoolExecutorScope可传递性。

通过Filter、Scope实现Request上下文

接下来介绍如何通过FilterScope实现Request粒度的Scope上下文。思路是:通过注入一个拦截器,在进入拦截器后开启Scope作为一个请求的上下文,解析Request对象获取获取相关状态信息(如登陆用户),并在Scope中设置,在离开拦截器时关闭Scope

AuthScope.java

// 获取登录态的工具类
public class AuthScope {
    private static final ScopeKey<String> LOGIN_USER = new ScopeKey<>();

    public static String getLoginUser() {
        return LOGIN_USER.get();
    }

    public static void setLoginUser(String loginUser) {
        if (loginUser == null) {
            loginUser = "unknownUser";
        }
        LOGIN_USER.set(loginUser);
    }
}

ScopeFilter.java

@Lazy
@Order(0)
@Service("scopeFilter")
public class ScopeFilter extends OncePerRequestFilter {

    @Override
    protected String getAlreadyFilteredAttributeName() {
        return this.getClass().getName();
    }

    @Override
    protected void doFilterInternal(HttpServletRequest request, HttpServletResponse response,
                                    FilterChain filterChain) throws ServletException, IOException {
        // 开启Scope
        beginScope();
        try {
            Cookie[] cookies = request.getCookies();
            String loginUser = "unknownUser";
            if (cookies != null) {
                for (Cookie cookie : cookies) {
                    if (cookie.getName().equals("login_user")) {
                        loginUser = cookie.getValue();
                        break;
                    }
                }
            }

            // 设置该 Request 上下文对用的登陆用户
            AuthScope.setLoginUser(loginUser);

            filterChain.doFilter(request, response);
        } finally {
            // 关闭Scope
            endScope();
        }
    }
}

注入Filter

@Slf4j
@Configuration
public class FilterConfig {

    @Bean
    public FilterRegistrationBean<ScopeFilter> scopeFilterRegistration() {
        FilterRegistrationBean<ScopeFilter> registration = new FilterRegistrationBean<>();
        registration.setFilter(new ScopeFilter());
        registration.addUrlPatterns("/rest/*");
        registration.setOrder(0);
        log.info("scope filter registered");
        return registration;
    }
}

UserController.java

@Slf4j
@RestController
@RequestMapping("/rest")
public class UserController {

    // curl --location --request GET 'localhost:8080/rest/getLoginUser' --header 'Cookie: login_user=zhangsan'
    @GetMapping("/getLoginUser")
    public String getLoginUser() {
        return AuthScope.getLoginUser();
    }

    // curl --location --request GET 'localhost:8080/rest/getLoginUserInThreadPool' --header 'Cookie: login_user=zhangsan'
    @GetMapping("/getLoginUserInThreadPool")
    public String getLoginUserInThreadPool() {
        ScopeThreadPoolExecutor executor = ScopeThreadPoolExecutor.newFixedThreadPool(4);
        executor.execute(() -> {
            log.info("get login user in thread pool: {}", AuthScope.getLoginUser());
        });

        return AuthScope.getLoginUser();
    }
}

通过以下请求验证,请求线程和线程池线程是否能获取登录态,其中登录态通过Cookie模拟:

curl --location --request GET 'localhost:8080/rest/getLoginUser' --header 'Cookie: login_user=zhangsan'
curl --location --request GET 'localhost:8080/rest/getLoginUserInThreadPool' --header 'Cookie: login_user=zhangsan'

感谢阅读,希望对你有所帮助 😃

敬请期待我的下一篇文章,谢谢。
在这里插入图片描述

来源:juejin.cn/post/7153287656624324638

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值