Spring @Async propagate context information

Question

I've a Spring Boot 2.2 application. I created a service like this:

@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)));
}

this is my @Async configuration:

@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;
    }
}

Using SecurityContextHolder.MODE_INHERITABLETHREADLOCAL I see the Security context is passed to the @Async method. In my multi-tenant application I use a ThreadLocal to set the tenant's 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();
    }

}

Because ThreadLocal is related to the thread, it's not available in the @Async method. Furthemore my custom @PlanAuthorization aop needs it to perform verifications of the tenant's plan. Is there a clean way to set TenantContext in any @Async method in my application?

Answer 1

I ended up to use a 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();
            }
        };
    }
}

and I used it in this way:

@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;
    }
}

It works and it seems also a neat solution.

Answer 2

The solution for such case is to:

1. configure custom thread pool so that you override it's execute method to sets up your thread local (or executes any task from your main context), decorate the task and submit decorated task for execution instead of original one

2. instruct @Async annotation to use concreate thread pool

    @Bean("tenantExecutor) public Executor threadLocalAwareThreadPool() { final CustomizableThreadFactory threadNameAwareFactory = new CustomizableThreadFactory("threadAsync"); final ThreadPoolExecutor threadPoolExecutor = new ThreadPoolExecutor(2, 10, 0L, TimeUnit.MILLISECONDS, new ArrayBlockingQueue<>(500), threadNameAwareFactory) { // override original method of thread pool @Override public void execute(Runnable originalTask) { final String tenantId = tenantThreadLocal.get(); // read data from current before passing the task to async thread // decorate the actual task by creating new task (Runnable) where you first set up the thread local and then execute your actual task super.execute(() -> { tenantThreadLocal.set(tenantId); // set data in actual async thread originalTask.run(); }); } }; return threadPoolExecutor; }

Nowe we tell spring use our custom exectuor

  @Async("tenantExecutor") 
    public Future<AuditCdr> saveCDR(Cdr3CXDto cdrRecord) {
    // your code
    }

Answer 3

Instead of ThreadLocal you must use InheritableThreadLocal. Then you will see the values from the parent thread.

API Doc: https://docs.oracle.com/javase/8/docs/api/java/lang/InheritableThreadLocal.html

Here is an article about this in combination with Spring: https://medium.com/@hariohmprasath/async-process-using-spring-and-injecting-user-context-6f1af16e9759