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?

You could create your own authenticationToken and set the tenant there, so you would have it available via the SecurityContext. (Just an idea) — Markus Schreiber, May 08 '20 at 10:52
What if you split your service into two layers - first has authorization stuff in it and calls second layer, while second is defined as async, and accepts `tenant code` as one of its method parameters? — M. Prokhorov, May 08 '20 at 11:11
And a tiny improvement: instead of code inside `getCurrentTenantId`, you could initialize `ThreadLocal ~~code = ThreadLocal.withInitial(() -> TENANT_DEFAULT);`~~ — M. Prokhorov, May 08 '20 at 11:14
@M.Prokhorov thanks for your hints. I wanted to avoid to split the service to be honest. I was looking for a more transparent solution as using Async annotation — drenda, May 08 '20 at 12:11
Anything involved context sharing in async methods in Spring requires extra leg work, usually on the part of Spring like with security context. I can't say off the top of my head that you can or cannot replicate behavior with security context. If your `TenantCode` has a meaning of `GrantedAuthority`, you could try putting in in Spring's `Authorization`, or look at the Spring code that shares contexts to see if you can plug into it. — M. Prokhorov, May 08 '20 at 12:19

drenda · Answer 1 · 2020-05-12T10:45:27.763

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.

Really helpful! In my case, I want to set RequestContext, so I added this in the TaskDecorator:`RequestAttributes attributes = RequestContextHolder.getRequestAttributes();` — Eden Li, Dec 19 '22 at 14:16

score 13 · Accepted Answer · edited Sep 03 '21 at 07:10

The solution for such case is to :

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

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....
    }

score 0 · Answer 3 · answered May 08 '20 at 10:54

0

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

answered May 08 '20 at 10:54

Simon Martinelli

34,053
5
48
82

5

Spring doesn't necessarily create a child thread for each async call, so `Inheritable` TL isn't guaranteed to work either. – M. Prokhorov May 08 '20 at 11:10
This does not guaranteed to work. Usually it works untill first corepoolsize calls as thread pool will create new thread until then. Once your corepoolsize threads created it will reuse existing threads in pool so it will not inherit thread local variables even though they are Inheritable. Above solution given by @drenda works fine and neat to implement. – Prakash Boda Dec 18 '20 at 09:03

score 0 · Answer 4 · answered Jun 21 '23 at 19:02

For Log4j2 ThreadContext

import java.util.Map;

import org.apache.logging.log4j.ThreadContext;
import org.springframework.core.task.TaskDecorator;

public class RequestContextTaskDecorator implements TaskDecorator {

    @Override
    public Runnable decorate(Runnable runnable) {
        // Right now:Web Thread context !
        // (Grab the current thread context map data)
        Map<String, String> contextMap = ThreadContext.getImmutableContext();
        return () -> {
            try {
                // Right now: @Async thread context !
                // (Restore the web thread context's data)
                ThreadContext.putAll(contextMap);
                runnable.run();
            } finally {
                ThreadContext.clearAll();
            }
        };

    }

}

Spring @Async propagate context information

4 Answers4