Spring Boot Custom #Async Wrap Callable - spring

Question

I am working on an application that supports multi tenancy. The tenant's unqiue identifier is stored in a thread local and can be accessed via some service.
To allow parallel processing, I have created a Callable wrapper, sets the thread local variable:
class TenantAwareCallable<T> implements Callable<T> {
private final String tenantName;
private final Callable<T> delegate;
TenantAwareCallable(Callable<T> delegate, String tenantName) {
this.delegate = delegate;
this.tenantName = tenantName;
}
#Override
public T call() throws Exception {
// set threadlocal
TenantContext.setCurrentTenantName(tenantName);
try {
return delegate.call();
} catch (Exception e) {
// log and handle
} finally {
TenantContext.clear();
}
}
}
This can already be used in the application. But what I would like to have is some custom #Async annotation, like for example #TenantAwareAsync or #TenantPreservingAsync, that wraps the callable, created by Spring in this one and then executes it.
Is there some way to get started with this?
Thanks in advance!

Answer 1

I have a working solution for this, so I think sharing it here might help some people some day.
I solved it not by using the TenantAwareCallable but by customizing the Executor service. I chose to extend Spring's ThreadPoolTaskScheduler (since this is what we use in the project).
public class ContextAwareThreadPoolTaskScheduler extends ThreadPoolTaskScheduler {
#Override
protected ScheduledExecutorService createExecutor(int poolSize, ThreadFactory threadFactory, RejectedExecutionHandler rejectedExecutionHandler) {
return new ContextAwareThreadPoolTaskExecutor(poolSize, threadFactory, rejectedExecutionHandler);
}
}
The actual setting of the context data is done in a customized ScheduledThreadPoolExecutor:
public class ContextAwareThreadPoolTaskExecutor extends ScheduledThreadPoolExecutor {
public ContextAwareThreadPoolTaskExecutor(int poolSize, ThreadFactory threadFactory, RejectedExecutionHandler rejectedExecutionHandler) {
super(poolSize, threadFactory, rejectedExecutionHandler);
}
#Override
protected <V> RunnableScheduledFuture<V> decorateTask(Callable<V> callable, RunnableScheduledFuture<V> task) {
return new ContextAwareTask<V>(task);
}
#Override
protected <V> RunnableScheduledFuture<V> decorateTask(Runnable runnable, RunnableScheduledFuture<V> task) {
return new ContextAwareTask<V>(task);
}
static private class ContextAwareTask<T> implements RunnableScheduledFuture<T> {
private final RunnableScheduledFuture<T> delegate;
private final TenantContextHolder multitenantContextHolder;
private final LoggingContextHolder loggingContextHolder;
private final SecurityContext securityContext;
ContextAwareTask(RunnableScheduledFuture<T> delegate) {
this.delegate = delegate;
multitenantContextHolder = TenantContextHolder.newInstance();
loggingContextHolder = LoggingContextHolder.newInstance();
securityContext = SecurityContextHolder.getContext();
}
#Override
public void run() {
multitenantContextHolder.apply();
loggingContextHolder.apply();
SecurityContextHolder.setContext(securityContext);
delegate.run();
SecurityContextHolder.clearContext();
loggingContextHolder.clear();
multitenantContextHolder.clear();
}
// all other methods are just delegates
}
}
The Holders are basically just objects to store context state and apply it in the new thread.
public class TenantContextHolder {
private String tenantName;
public static TenantContextHolder newInstance() {
return new TenantContextHolder();
}
private TenantContextHolder() {
this.tenantName = TenantContext.getCurrentTenantName();
}
public void apply() {
TenantContext.setCurrentTenantName(tenantName);
}
public void clear() {
TenantContext.clear();
}
}
The custom implementation of the Scheduler can then be configured in the Spring environment.
#Configuration
public class AsyncConfiguration implements AsyncConfigurer {
private ThreadPoolTaskScheduler taskScheduler;
#Bean
public ThreadPoolTaskScheduler taskScheduler() {
if (taskScheduler == null) {
taskScheduler = new ContextAwareThreadPoolTaskScheduler();
}
return taskScheduler;
}
#Override
public Executor getAsyncExecutor() {
return taskScheduler();
}
}