java - Spring @Async 传播上下文信息

Question

我有一个 Spring Boot 2.2 应用程序。我创建了这样的服务：

@Async
@PreAuthorize("hasAnyRole('ROLE_PBX')")
@PlanAuthorization(allowedPlans = {PlanType.BUSINESS, PlanType.ENTERPRISE})
public Future<AuditCdr> saveCDR(Cdr3CXDto cdrRecord) {
    log.debug("Current tenant {}", TenantContext.getCurrentTenantId());  

    return new AsyncResult<AuditCdr>(auditCdrRepository.save(cdr3CXMapper.cdr3CXDtoToAuditCdr(cdrRecord)));
}

这是我的@Async 配置：

@Configuration
@EnableAsync
public class AsyncConfiguration implements AsyncConfigurer {

    @Override
    public Executor getAsyncExecutor() {
        SecurityContextHolder.setStrategyName(SecurityContextHolder.MODE_INHERITABLETHREADLOCAL);
        ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
        executor.setCorePoolSize(2);
        executor.setMaxPoolSize(10);
        executor.setQueueCapacity(500);
        executor.setThreadNamePrefix("threadAsync");
        executor.setWaitForTasksToCompleteOnShutdown(true);
        executor.initialize();
        return executor;
    }
}

使用SecurityContextHolder.MODE_INHERITABLETHREADLOCAL我看到安全上下文被传递给 @Async 方法。在我的多租户应用程序中，我使用 ThreadLocal 来设置租户的 ID：

public class TenantContext {
    public final static String TENANT_DEFAULT = "empty";
    private static final ThreadLocal<String> code = new ThreadLocal<>();

    public static void setCurrentTenantId(String code) {
        if (code != null)
            TenantContext.code.set(code);
    }

    public static String getCurrentTenantId() {
        String tenantId = code.get();
        if (StringUtils.isNotBlank(tenantId)) {
            return tenantId;
        }
        return TENANT_DEFAULT;
    }

    public static void clear() {
        code.remove();
    }

}

因为 ThreadLocal 与线程相关，所以在 @Async 方法中是不可用的。此外，我的自定义@PlanAuthorizationaop 需要它来执行租户计划的验证。是否有一种干净的方法可以在我的应用程序的任何 @Async 方法中设置 TenantContext？

标签： javaspringspring-bootspring-async

Answer 1

我最终使用了TaskDecorator：

@Log4j2
public class MdcTaskDecorator implements TaskDecorator {

    @Override
    public Runnable decorate(Runnable runnable) {
        // Right now: Web thread context !
        // (Grab the current thread MDC data)
        String tenantId = TenantContext.getCurrentTenantId();
        Long storeId = StoreContext.getCurrentStoreId();
        SecurityContext securityContext = SecurityContextHolder.getContext();
        Map<String, String> contextMap = MDC.getCopyOfContextMap();
        log.info("Saving tenant information for async thread...");
        return () -> {
            try {
                // Right now: @Async thread context !
                // (Restore the Web thread context's MDC data)
                TenantContext.setCurrentTenantId(tenantId);
                StoreContext.setCurrentStoreId(storeId);
                SecurityContextHolder.setContext(securityContext);
                MDC.setContextMap(contextMap);
                log.info("Restoring tenant information for async thread...");
                runnable.run();
            } catch (Throwable e) {
                log.error("Error in async task", e);
            } finally {
                MDC.clear();
            }
        };
    }
}

我以这种方式使用它：

@Configuration
@EnableAsync
public class AsyncConfiguration implements AsyncConfigurer {

    @Override
    public Executor getAsyncExecutor() {
        SecurityContextHolder.setStrategyName(SecurityContextHolder.MODE_INHERITABLETHREADLOCAL);
        ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
        executor.setCorePoolSize(1);
        executor.setMaxPoolSize(100);
        executor.setQueueCapacity(500);
        executor.setThreadNamePrefix("threadAsync");
        executor.setWaitForTasksToCompleteOnShutdown(true);
        executor.setTaskDecorator(new MdcTaskDecorator());
        executor.initialize();
        return executor;
    }
}

它有效，而且似乎也是一个巧妙的解决方案。