异步任务中传递用户信息的一种优雅写法

前言

在异步任务中，我们通常会遇到子任务获取当前用户的场景。我们可能会采取ThreadLocal来存储主线程传递的用户信息。然后在业务开始时set，业务结束时remove，来保证不会出现OOM的场景。

基础写法

@SpringBootTest
public class UserAsyncTransmitTest {

    @Resource
    private Executor myTaskExecutor;

    @Resource
    private TestService testService;

    @Test
    void test() throws InterruptedException {
        LoginUser loginUser = LoginUser.builder().username("zhangsan").build();
        System.out.println("主线程start==================");

        myTaskExecutor.execute(() -> {
            UserContext.setLoginUser(loginUser);
            try {
                testService.doSomething();
            } catch (Exception e) {
                e.printStackTrace();
            } finally {
                UserContext.remove();
            }
        });
        Thread.sleep(5000L);
        System.out.println("主线程end===================");
    }
}

测试示例

运行结果：

这种写法，每次写异步任务需要用户时，都要把用户信息set，remove。那么有没有一种写法可以让写业务代码时更加纯粹一点，让程序员更关注业务，代码更加美观呢？
其实，我们可以自定义Runnable加上自定义线程池来实现。

升级写法

• UserContextRunnable.java

public class UserContextRunnable implements Runnable {

    private final LoginUser currentUser;

    private final Runnable runnable;

    public UserContextRunnable(LoginUser currentUser, Runnable runnable) {
        this.currentUser = currentUser;
        this.runnable = runnable;
    }

    @Override
    public void run() {
        try {
            UserContext.setLoginUser(currentUser);
            runnable.run();
        } finally {
            UserContext.remove();
        }
    }
}

• UserContextTaskExecutor.java

public class UserContextTaskExecutor extends ThreadPoolTaskExecutor {

    private static final long serialVersionUID = 7334165009514720624L;

    private LoginUser currentUser;

    public void user(LoginUser currentUser) {
        this.currentUser = currentUser;
    }

    public void execute(LoginUser currentUser, @NotNull Runnable task) {
        super.execute(new UserContextRunnable(currentUser, task));
    }

    public void execute(@NotNull Runnable task) {
        if (currentUser != null) {
            super.execute(new UserContextRunnable(currentUser, task));
        }else{
            super.execute(task);
        }
    }
}

测试示例

@SpringBootTest
public class UserAsyncTransmitTest {

    @Resource
    private Executor userContextTaskExecutor;

    @Resource
    private TestService testService;

    @Test
    void test() throws InterruptedException {
        LoginUser loginUser = LoginUser.builder().username("zhangsan").build();
        System.out.println("主线程start==================");

        userContextTaskExecutor.execute(new UserContextRunnable(loginUser, ()-> 
                testService.doSomething()
        ));
        Thread.sleep(5000L);
        System.out.println("主线程end===================");
    }
}

运行结果：