Propagating Tenant data to Spawned/Async Threads with ThreadLocalTargetSource, TaskDecorator in Multi-tenant Spring Boot applications

1. OVERVIEW

In late 2017, I published a blog post about using Spring’s ThreadLocalTargetSource as an alternative to dealing explicitly with ThreadLocal in Spring Boot Multi-tenant applications.

Where did I use ThreadLocal that I found Spring’s ThreadLocalTargetSource a useful alternative?

Let’s start with some context.

Earlier that same year I had published writing multi-tenant RESTful applications with Spring Boot, JPA, Hibernate and Postgres, where I took the shared application(s), different databases, multi-tenancy approach.

Multi-tenancy - Shared applications, different databases

The implementation used an HTTP Header with the Tenant identifier to choose the relational database from which to retrieve the data from.

You could send requests like:

curl -H "X-TENANT-ID: tenant_1" http://localhost:8080/...

curl -H "X-TENANT-ID: tenant_2" http://localhost:8080/...

The code also included a Spring MVC Interceptor to set and clear the Tenant identifier to and from the TenantContext using that HTTP custom header.

Note: You could have written a Servlet Filter instead.

DvdRentalMultiTenantInterceptor.java:

public class DvdRentalMultiTenantInterceptor extends HandlerInterceptorAdapter {

  @Override
  public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
    String tenantId = request.getHeader(TENANT_HEADER_NAME);
    DvdRentalTenantContext.setTenantId(tenantId);
    return true;
  }

  @Override
  public void postHandle(HttpServletRequest request, HttpServletResponse response, Object handler, ModelAndView modelAndView) throws Exception {
    DvdRentalTenantContext.clear();
  }

}

and

DvdRentalTenantContext.java:

public class DvdRentalTenantContext {

  private static final ThreadLocal<String> CONTEXT = new ThreadLocal<>();

  public static void setTenantId(String tenantId) {
    CONTEXT.set(tenantId);
  }

  public static String getTenantId() {
    return CONTEXT.get();
  }

  public static void clear() {
    CONTEXT.remove();
  }

}

You can now see how the TenantContext implementation uses Java’s ThreadLocal to set, retrieve and clear Tenant’s data, through the DvdRentalMultiTenantInterceptor.

A problem with ThreadLocal is that you have to be very careful how do you use it. You’ve got to be cautions where to set/clear data to/from the ThreadLocal to prevent the tenant’s data from leaking once the thread processing the request is done and back into the pool to process a future request.

And thus how I wrote about Spring’s ThreadLocalTargetSource as an alternative to using ThreadLocal directly in Spring Boot Multi-tenant applications.

But, there was a problem:

Multi-tenancy - ThreadLocalTargetSource and Async Requests

The Tenant identifier wasn’t being propagated when sending or processing async requests. And that’s what this blog post helps fixing.

Tip: You might want to take a look at the recent Code and Configuration Snippets / Recipes / Tips section, which helps you with common tasks without reading through a full blog post.

Note: Before moving into the next section:

This blog post doesn’t cover common code that instantiates and autowires ThreadLocalTargetSource, the TenantStore, and the TenantFilter.

You can read about it on:

Spring’s ThreadLocalTargetSource as an alternative to using ThreadLocal directly in Spring Boot Multi-tenant applications.

2. RESTful CONTROLLER

Let’s start with a simple Spring Boot RESTful Controller:

DemoController.java:

@RestController
public class DemoController {

  private final TenantStore tenantStore;
  private final AsyncTaskExecutor taskExecutor;

  @RequestMapping(path = "/sync", method = RequestMethod.GET)
  public String getDemo() {
    return this.formattedString("Sync Task");
  }

  @RequestMapping(path = "/async", method = RequestMethod.GET)
  public String getAsyncDemo() throws Exception {
    this.formattedString("Before Async");
    Future<String> futureTask = this.calculateAsync();
    String result = futureTask.get();
    this.formattedString("After Async");
    return result;
}

  private Future<String> calculateAsync() {
    return this.taskExecutor.submit(() -> this.formattedString("Async Task"));
  }

  private String formattedString(String prefix) {
    String result = String.format("%s. Tenant: %s", prefix, this.tenantStore.getTenantId());
    log.debug(result);
    return result;
  }

}

The REST Controller autowires a TenantStore Bean instead of using a TenantContext or TenantContextHolder with static methods that access the ThreadLocal directly.

It also includes an AsyncTaskExecutor Bean to process requests asynchronously.

Two GET methods are implemented:

getDemo() processes the request synchronously.
getAsyncDemo() processes the request asynchronously.

While both of them use the TenantStore class attribute to retrieve, log and return the tenantId.

3. TASK EXECUTOR CONFIGURATION

Let’s look at the AsyncTaskExecutor configuration:

TaskExecutorConfig.java:

@Configuration
public class TaskExecutorConfig {

  private final TaskExecutionProperties taskExecutionProperties;
  private final TenantStore tenantStore;

  @Bean
  public AsyncTaskExecutor taskExecutor() {
    return new ThreadPoolTaskExecutorBuilder()
      .corePoolSize(this.taskExecutionProperties.getPool().getCoreSize())
      .maxPoolSize(this.taskExecutionProperties.getPool().getMaxSize())
      .queueCapacity(this.taskExecutionProperties.getPool().getQueueCapacity())
      .threadNamePrefix(this.taskExecutionProperties.getThreadNamePrefix())
      .taskDecorator(this.tenantStoreTaskDecorator())
      .build();
  }

  @Bean
  public TaskDecorator tenantStoreTaskDecorator() {
    return new TenantStoreTaskDecorator(this.tenantStore);
  }

}

application.yml:

spring:
  task:
    execution:
      pool:
        core-size: 20
        max-size: 100
        queue-capacity: 3000
      thread-name-prefix: async-thread-

This is very simple and common ThreadPoolTaskExecutor instantiation.

The important pieces here are taskDecorator(this.tenantStoreTaskDecorator()) and the TenantStoreTaskDecorator class.

That’s what you need to propagate the Tenant identifier into spawned/async threads while being able to inject/autowire the TenantStore into other Spring Beans.

Tip: Read on to learn how to configure and switch between ThreadPoolTaskExecutor and VirtualThreadTaskExecutor in Spring Boot applications.

4. TENANT STORE TASK DECORATOR

TenantStoreTaskDecorator.java:

public class TenantStoreTaskDecorator implements TaskDecorator {

  private final TenantStore tenantStore;

  @Override
  public Runnable decorate(Runnable task) {
    String tenantId = this.tenantStore.getTenantId();
    return () -> {
      try {
        this.tenantStore.setTenantId(tenantId);
        task.run();
      } finally {
        this.tenantStore.setTenantId(null);
      }
    };
  }

}

This class implements Spring’s TaskDecorator interface, which as mentioned earlier, is key along with TaskExecutorConfig to get the Tenant identifier propagated into spawned/async threads.

It sets the Tenant identifier, runs the async process, then clears the tenantStore.

Cleaning the tenantStore class attribute doesn’t mean it sets the parent thread’s tenantStore.tenantId to null.

Remember tenantStore is prototype-scoped Spring Bean instantiated and proxied using ThreadLocalTargetSource, ProxyFactoryBean, and TenantStore classes.

5. SAMPLE REQUESTS

Now that we put all the pieces together, let’s send a couple of requests, watch the response and application logs.

curl -H "X-TENANT-ID: tenant_1" http://localhost:8080/demo/sync
Sync Task. Tenant: tenant_1

The Controller returns the Tenant identifier tenant_1 synchronously from the TenantStore.
That’s the same behavior the Controller had before including the Spring’s TaskDecorator changes.

The applications logs the same output as the response.

... [nio-8080-exec-1] com.asimiotech.demo.web.DemoController   : Sync Task. Tenant: tenant_1

Let’s now send a request that spawns a thread to process it.

curl -H "X-TENANT-ID: tenant_2" http://localhost:8080/demo/async
Async Task. Tenant: tenant_2

This time the application logs:

... [nio-8080-exec-1] com.asimiotech.demo.web.DemoController   : Before Async. Tenant: tenant_2
... [ async-thread-1] com.asimiotech.demo.web.DemoController   : Async Task. Tenant: tenant_2
... [nio-8080-exec-1] com.asimiotech.demo.web.DemoController   : After Async. Tenant: tenant_2

Tomcat’s nio-8080-exec-1 processing thread logs Before Async. Tenant: tenant_2.

TaskExecutor’s async-thread-1 spawned/async thread logs Async Task. Tenant: tenant_2

Tomcat’s nio-8080-exec-1 processing thread logs After Async. Tenant: tenant_2.

This confirms that the Tenant identifier is now propagated to spawned/async threads.
You can also see that cleaning/clearing the Tenant identifier in the TaskDecorator doesn’t clear it from the parent thread that spawned the async thread, in this case Tomcat’s nio-8080-exec-1.

But let’s take it a step further and automate these assertions in an integration test class.

6. INTEGRATION TEST CLASS

// ...
public class DemoControllerIntegrationTest {

// ...

  @Test
  public void shouldTenant2Async(CapturedOutput output) {
    String actualContent = RestAssured
      .given()
        .accept(ContentType.TEXT)
        .header("X-TENANT-ID", "tenant_2")
      .when()
        .get("/demo/async")
      .then()
        .statusCode(HttpStatus.OK.value())
        .extract().asString();

    MatcherAssert.assertThat(actualContent, Matchers.equalTo("Async Task. Tenant: tenant_2"));
    MatcherAssert.assertThat(output.getOut(), Matchers.containsString("auto-1-exec-1] com.asimiotech.demo.web.DemoController   : Before Async. Tenant: tenant_2"));
    MatcherAssert.assertThat(output.getOut(), Matchers.containsString("async-thread-1] com.asimiotech.demo.web.DemoController   : Async Task. Tenant: tenant_2"));
    MatcherAssert.assertThat(output.getOut(), Matchers.containsString("auto-1-exec-1] com.asimiotech.demo.web.DemoController   : After Async. Tenant: tenant_2"));
  }
// ...
}

Multi-tenancy - ThreadLocalTargetSource, Async Requests Integration Test

Similar HTTP response and console log assertions as those in the sample requests, but in an automated fashion.

Tip: Read on to learn how to capture and assert logs with JUnit 5 in Spring Boot integration tests.

7. CONCLUSION

This blog post covered using ThreadLocalTargetSource as an alternative to avoid using ThreadLocal directly in your Spring Boot applications.

It focussed on the configuration and code to propagate Tenant’s data to spawned/async threads and being able to inject/autowire it into other Spring Beans in Multitenancy architectures.

And it also included an integration test with assertions to verify the correctness of the Tenant’s data before, during, and after running an async task to process a RESTful request.

Tip: This blog post is a follow up and improvement of:

An alternative approach to ThreadLocal using Spring.

Thanks for reading and as always, feedback is very much appreciated. If you found this post helpful and would like to receive updates when content like this gets published, sign up to the newsletter.