Spring @Async propagate context information

user_2012635 0
Viewed 6908

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?

spring-async spring-boot spring java
3 Answers

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.

user_20126350

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....
    }
user_18741130
user_10451420
Related