javaspringspring-bootspring-async

Spring @Async propagate context information


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?


Solution

  • 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 concrete 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;
      
        }
      

    Now we tell spring use our custom executor

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