An alternative approach to ThreadLocal using Spring

1. OVERVIEW

In a blog post published sometime ago, Multitenant applications using Spring Boot JPA Hibernate and Postgres I included some code to set to and to retrieve from, the tenant identifier (a discriminator for selecting its associated datasource) using a ThreadLocal reference:

A context holder class for holding the tenant data:

...
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();
  }
}

A Spring MVC interceptor (which could have been also done using a servlet filter) to set and clear such tenant identifier:

public class DvdRentalMultiTenantInterceptor extends HandlerInterceptorAdapter {

  private static final String TENANT_HEADER_NAME = "X-TENANT-ID";

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

And somewhere in the application:

String currentTenantId = DvdRentalTenantContext.getTenantId();

I normally try to avoid using ThreadLocal and would also advise to limit their usage. They can be handy, solve difficult problems but can also introduce memory leaks.

In this post I discuss how to use Spring’s ThreadLocalTargetSource to prevent dealing directly with the dangerous ThreadLocal while practicing dependency injection and proper mocking in unit tests.

Note: This approach would still use a ThreadLocal reference, but delegates setting it, injecting it and clearing it, to the Spring framework.

2. REQUIREMENTS

Java 7+.
Maven 3.2+.
Familiarity with Spring Framework.

3. CREATE THE DEMO APP

curl "https://start.spring.io/starter.tgz"
 -d dependencies=actuator,web
 -d language=java
 -d type=maven-project
 -d baseDir=threadlocaltargetsource-demo
 -d groupId=com.asimio
 -d artifactId=threadlocaltargetsource-demo
 -d version=0-SNAPSHOT
 | tar -xzvf -

This command will create a Maven project in a folder named threadlocaltargetsource-demo with the actuator and web-related Spring Boot dependencies.

Tip: Read on if you are interested in adding Spring Boot dependencies using the BOM approach.

Let’s take a look at the relevant classes included in this demo application:

TenantStore.java:

package com.asimio.demo.tenant;
...
public class TenantStore {

  private String tenantId;

  public void clear() {
    this.tenantId = null;
  }

  // Getters, Setters
}

This class serves as the tenant data holder.

AppConfig.java:

package com.asimio.demo.config;
...
@Configuration
public class AppConfig {

  @Bean
  public Filter tenantFilter() {
    return new TenantFilter();
  }

  @Bean
  public FilterRegistrationBean tenantFilterRegistration() {
    FilterRegistrationBean result = new FilterRegistrationBean();
    result.setFilter(this.tenantFilter());
    result.setUrlPatterns(Lists.newArrayList("/*"));
    result.setName("Tenant Store Filter");
    result.setOrder(1);
    return result;
  }

  @Bean(destroyMethod = "destroy")
  public ThreadLocalTargetSource threadLocalTenantStore() {
    ThreadLocalTargetSource result = new ThreadLocalTargetSource();
    result.setTargetBeanName("tenantStore");
    return result;
  }

  @Primary
  @Bean(name = "proxiedThreadLocalTargetSource")
  public ProxyFactoryBean proxiedThreadLocalTargetSource(ThreadLocalTargetSource threadLocalTargetSource) {
    ProxyFactoryBean result = new ProxyFactoryBean();
    result.setTargetSource(threadLocalTargetSource);
    return result;
  }

  @Bean(name = "tenantStore")
  @Scope(scopeName = "prototype")
  public TenantStore tenantStore() {
    return new TenantStore();
  }
}

This is where Spring’s beans are instantiated. tenantFilter() and tenantFilterRegistration() are straightforward, tenantFilter() instantiates a servlet filter and tenantFilterRegistration() implements a Spring’s mechanism that allows dependencies to be injected in TenantFilter.java, a regular servlet filter.

Tip: Read on Implementing and configuring Servlets, Filters and Listeners in Spring Boot applications for more information.

The other beans look interesting:

threadLocalTenantStore bean. ThreadLocalTargetSource is useful when you need an object, a TenantStore.java instance in this case, to be created for each incoming request. The target (a TenantStore.java object) will be instantiated only once in each thread and will get removed from each thread’s threadLocals map when ThreadLocalTargetSource’s destroy() is called, for instance when the application is shutdown.
tenantStore bean is the target object stored in each thread. It is required to be prototype-scoped in AbstractPrototypeBasedTargetSource, parent class of ThreadLocalTargetSource.
proxiedThreadLocalTargetSource bean. According to the documentation and source code comments TargetSources must run in a BeanFactory since they need to call the getBean() method to create a new prototype instance.

DemoResource.java:

package com.asimio.demo.web;
...
@RestController
@RequestMapping(value = "/demo")
public class DemoResource {

  @Autowired
  private TenantStore tenantStore;

  @RequestMapping(method = RequestMethod.GET)
  public String getDemo() {
    return String.format("Tenant: %s", this.tenantStore.getTenantId());
  }
}

This class implements a simple API in which a TenantStore.java instance is injected.

We will see later how it would appear there is only one instance of TenantStore.java class for different requests, each holding different tenant data.

TenantFilter.java:

package com.asimio.demo.web;
...
public class TenantFilter implements Filter {
...
  @Autowired
  private TenantStore tenantStore;
...
  @Override
  public void doFilter(ServletRequest servletRequest, ServletResponse servletResponse, FilterChain chain)
      throws  IOException, ServletException {

    HttpServletRequest request = (HttpServletRequest) servletRequest;
    String tenantId = request.getHeader(TENANT_HEADER_NAME);
    try {
      this.tenantStore.setTenantId(tenantId);
      chain.doFilter(servletRequest, servletResponse);
    } finally {
      // Otherwise when a previously used container thread is used, it will have the old tenant id set and
      // if for some reason this filter is skipped, tenantStore will hold an unreliable value
      this.tenantStore.clear();
    }
  }
...

This servlet filter takes care of setting the tenant identifier to the TenantStore.java holder and clearing it up during the servlet filter chain’s way out.

We will see later how it would appear there is only one instance of TenantStore.java class for different requests, each holding different tenant data.

4. PACKAGING AND RUNNING THE APP

This application can be run from your preferred IDE as a regular Java application or from command line:

cd <path to demo application>
mvn clean package
java -jar target/threadlocaltargetsource-demo.jar

mvn spring-boot:run

Let’s place a breakpoint in the TenantFilter.java class:

TenantFilter breakpoint

and send a couple of simultaneous requests to the /demo endpoint:

curl -v -H "X-TENANT-ID:tenant_1" "http://localhost:8080/demo"
*   Trying ::1...
* TCP_NODELAY set
* Connected to localhost (::1) port 8080 (#0)
> GET /demo HTTP/1.1
> Host: localhost:8080
> User-Agent: curl/7.54.0
> Accept: */*
> X-TENANT-ID:tenant_1
>

and

curl -v -H "X-TENANT-ID:tenant_2" "http://localhost:8080/demo"
*   Trying ::1...
* TCP_NODELAY set
* Connected to localhost (::1) port 8080 (#0)
> GET /demo HTTP/1.1
> Host: localhost:8080
> User-Agent: curl/7.54.0
> Accept: */*
> X-TENANT-ID:tenant_2
>

Both requests should have been suspended in the Tenant filter breakpoint, lets take a look at partial stack traces and attributes values:

TenantFilter breakpoint request #1 First request to /demo endpoint

Notice that tenantStore's id is 119 and that tenantStore's tenantId is set to tenant_1.

TenantFilter breakpoint request #2 Second request to /demo endpoint

As for the second request, tenantStore's id is also 119 but the tenantStore's tenantId is set to tenant_2. Interesting eh. All these while both requests are still being processed. Also notice http-nio-8080-exec-1’s stack trace which corresponds to processing request #1.

As the execution of the requests is resumed, the response to the first request looks like:

< HTTP/1.1 200
< Content-Type: text/plain;charset=UTF-8
< Content-Length: 16
< Date: Mon, 27 Nov 2017 05:41:23 GMT
<
* Connection #0 to host localhost left intact
Tenant: tenant_1

and the response to the second request:

< HTTP/1.1 200
< Content-Type: text/plain;charset=UTF-8
< Content-Length: 16
< Date: Mon, 27 Nov 2017 05:41:23 GMT
<
* Connection #0 to host localhost left intact
Tenant: tenant_2

Let’s now see how a unit test for the DemoResource.java class would look like:

Tip: Read on if you are interested in learning how to Split Unit and Integration Tests using Maven and Surefire plugin or Reporting Code Coverage using Maven and JaCoCo plugin.

DemoResourceTest.java:

package com.asimio.demo.web;
...
@RunWith(MockitoJUnitRunner.class)
public class DemoResourceTest {

  @Mock
  private TenantStore tenantStore;

  @InjectMocks
  private DemoResource controller;

  @Test
  public void shouldReturnBlahTenantId() {
    // Given
    BDDMockito.when(this.tenantStore.getTenantId()).thenReturn("blahTenantId");
    // When
    String result = this.controller.getDemo();
    // Then
    Assert.assertThat(result, Matchers.equalTo("Tenant: blahTenantId"));
  }
}

TenantStore dependency is being mocked using Mockito and injected in the API implementation. A case could be made to prevent using @InjectMocks but that goes out of the scope of this post.

Compare this:

DDMockito.when(this.tenantStore.getTenantId()).thenReturn("blahTenantId");
...

to how I used it earlier this year:

try {
  DvdRentalTenantContext.setTenantId("tenant_1");
  ...
} finally {
  DvdRentalTenantContext.clear();
}

which is error-prone since the ThreadLocal would have to be cleared in every test that follows this pattern to prevent a possible misuse if another unit test that uses the same ThreadLocal runs in the same thread.

And that’s all.

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.