ThreadLocal을 사용할 때 주의할 점

이번 글에서는 Java ThreadLocal 클래스에 대해 소개합니다.
 
ThreadLocal 클래스는 어떻게 구현되어 있는지 코드를 분석해보고, 간단한 사용법과 주의해야 할 점에 대해 정리해보겠습니다.
 

1. ThreadLocal?

Java ThreadLocal 클래스를 확인해보면 아래와 같이 ThreadLocal에 대한 전반적인 소개로 시작합니다.

ThreadLocal Class

요약하면, ThreadLocal 클래스는 쓰레드 단위로 로컬 변수를 관리할 수 있는 방법을 제공하고 각 변수는 쓰레드별로 get/set 메서드로 접근할 수 있다는 것인데요.
 
user ID, Transaction ID와 같이 쓰레드와 관련된 상태를 연관해서 저장해야 하는 상황에서 ThreadLocal 인스턴스를 사용할 수 있다고 예시를 들어주고 있습니다.
 
즉, ThreadLocal은 다중 쓰레드 환경에서 각 쓰레드에 고유한 데이터를 저장하고 관리하는 데 사용되는 클래스라고 정리할 수 있습니다.
 
 

2. ThreadLocal 코드 살펴보기

2-1. ThreadLocal 객체 생성, withInitial 메서드

ThreadLocal은 아래와 같이 제네릭 타입의 초기값을 가지며 기본적으로는 null을 반환합니다.
 
다만, jdk1.8부터 도입된 withInitial 메서드를 통해 ThreadLocal 변수에 초기값을 지정하여 리턴하는 SuppliedThreadLocal 객체를 생성할 수도 있습니다.

public class ThreadLocal<T> {
    
    // .. 생략
    
    protected T initialValue() {
        return null;
    }
    
    // .. 생략
    
    public static <S> ThreadLocal<S> withInitial(Supplier<? extends S> supplier) {
        return new SuppliedThreadLocal<>(supplier);
    }
    
    // .. 생략
    
    static final class SuppliedThreadLocal<T> extends ThreadLocal<T> {

        private final Supplier<? extends T> supplier;

        SuppliedThreadLocal(Supplier<? extends T> supplier) {
            this.supplier = Objects.requireNonNull(supplier);
        }

        @Override
        protected T initialValue() {
            return supplier.get();
        }
    }

withInitial() 메서드는 Supplier 함수형 인터페이스를 파라미터로 받는 것을 볼 수 있는데요. Supplier 함수형 인터페이스는 람다식을 통해 매개변수를 받지 않고 리턴값을 반환하는 방법을 제공합니다. 
 
아래 테스트 코드를 통해 간단한 ThreadLocal 변수 초기화 방법을 확인해보겠습니다.

@Test
@DisplayName("withInitial 메서드를 통해 쓰레드로컬 변수 초기값을 설정할 수 있다.")
void threadlocal_withInitial_returns_initialized_threadlocal_var() {
    ThreadLocal<Integer> threadLocal = ThreadLocal.withInitial(() -> 0);

    Integer value = threadLocal.get();

    Assertions.assertEquals(0, value);
}

테스트 성공

 

2-2. set / get 메서드

ThreadLocal의 set(), get() 메서드 코드 구현은 아래와 같습니다.

public class ThreadLocal<T> {
    // .. 생략

    public void set(T value) {
        Thread t = Thread.currentThread();
        ThreadLocalMap map = getMap(t);
        if (map != null) {
            map.set(this, value);
        } else {
            createMap(t, value);
        }
    }

    public T get() {
        Thread t = Thread.currentThread();
        ThreadLocalMap map = getMap(t);
        if (map != null) {
            ThreadLocalMap.Entry e = map.getEntry(this);
            if (e != null) {
                @SuppressWarnings("unchecked")
                T result = (T)e.value;
                return result;
            }
        }
        return setInitialValue();
    }
    
    // .. 생략

    static class ThreadLocalMap {

        
        static class Entry extends WeakReference<ThreadLocal<?>> {
            /** The value associated with this ThreadLocal. */
            Object value;

            Entry(ThreadLocal<?> k, Object v) {
                super(k);
                value = v;
            }
        }
         
    // .. 생략  
    
    ThreadLocalMap getMap(Thread t) {
        return t.threadLocals;
    }

    void createMap(Thread t, T firstValue) {
        t.threadLocals = new ThreadLocalMap(this, firstValue);
    }

ThreadLocalMap이라는 내부 static 클래스를 활용해서 값을 저장하고 조회하는데요.
 
set 메서드의 경우 현재 쓰레드의 ThreadLocalMap을 확인하고 존재한다면 맵에 값을 넣어주고 없다면 정적인 ThreadLocalMap을 생성하여 값을 지정하고 반환합니다.
 
get 메서드는 현재 쓰레드의 ThreadLocalMap을 확인하여 값을 리턴합니다.
 
참고로 Thread 클래스는 아래와 같이 ThreadLocalMap을 필드로 가지고 있습니다.
 

[Thread.java]

public class Thread implements Runnable {
    // .. 생략

    /* ThreadLocal values pertaining to this thread. This map is maintained by the ThreadLocal class. */
    ThreadLocal.ThreadLocalMap threadLocals = null;

이를 통해 해당 코드의 주석(This map is maintained by the ThreadLocal class.)의 내용처럼 ThreadLocal에서 개별 쓰레드의 변수 값을 조회할 수 있게 합니다. 
 

2-3. remove 메서드

remove 메서드 역시 아래와 같이 현재 쓰레드의 ThreadLocalMap을 가져오고 ThreadLocalMap에 할당된 Entry 값을 제거합니다.

     public void remove() {
         ThreadLocalMap m = getMap(Thread.currentThread());
         if (m != null) {
             m.remove(this);
         }
     }
     
     static class ThreadLocalMap {
        static class Entry extends WeakReference<ThreadLocal<?>> {
            /** The value associated with this ThreadLocal. */
            Object value;

            Entry(ThreadLocal<?> k, Object v) {
                super(k);
                value = v;
            }
        }

		// .. 생략

        /**
         * Remove the entry for key.
         */
        private void remove(ThreadLocal<?> key) {
            Entry[] tab = table;
            int len = tab.length;
            int i = key.threadLocalHashCode & (len-1);
            for (Entry e = tab[i];
                 e != null;
                 e = tab[i = nextIndex(i, len)]) {
                if (e.get() == key) {
                    e.clear();
                    expungeStaleEntry(i);
                    return;
                }
            }
        }

 
아래와 같은 테스트 코드로 remove() 메서드의 동작을 간단히 확인해보겠습니다.

@Test
@DisplayName("remove 메서드를 호출하면 threadLocal에 할당된 값이 제거된다.")
void threadlocal_remove_test() {
    ThreadLocal<Integer> threadLocal = ThreadLocal.withInitial(() -> 1);

    threadLocal.remove();

    Integer threadLocalValue = threadLocal.get();
    Assertions.assertNull(threadLocalValue);
}

이 테스트는 처음에 ThreadLocal의 변수를 1로 초기화한 뒤, remove() 메서드를 호출하고 다시 로컬 변수 값을 가져오고 있습니다.
 
remove() 메서드를 호출했기 때문에 할당된 변수가 없으므로 null이 결과가 되어야 할 것 같은데요.
하지만 아래와 같이 테스트는 실패합니다.
 

테스트가 실패한다.

이 이유에 대한 답은 remove() 메서드의 설명을 보면 알 수 있습니다.

Removes the current thread's value for this thread-local variable. If this thread-local variable is subsequently read by the current thread, its value will be reinitialized by invoking its initialValue method, unless its value is set by the current thread in the interim. This may result in multiple invocations of the initialValue method in the current thread.

 
'remove() 메서드는 ThreadLocal의 변수를 제거하지만, 만약 remove() 이후에 이어서 ThreadLocal의 변수를 읽으려고 하면 initialValue 메서드가 호출되어 변수가 다시 할당된다'는 이야기를 하고 있습니다.
 
따라서, 어떻게 보면 당연하겠지만 ThreadLocal의 자원을 반납하는 remove() 메서드를 호출한 이후에는 재차 해당 ThreadLocal의 변수값을 읽지 말아야 합니다.
 

3. ThreadPool과 ThreadLocal

앞서 살펴본 것처럼 ThreadLocal은 특정 쓰레드에 대해 ThreadLocalMap을 가지고 로컬 변수를 관리합니다.

 

지금부터는 이러한 ThreadLocal의 성격으로 인해 발생할 수 있는 문제점을 짚어보려고 합니다.

 

먼저, 아무 문제가 없는 기본적인 ThreadLocal 사용에 대한 코드를 보겠습니다.

3-1. 기본 코드의 동작

[JwtContext.java]

public class JwtContext {
    private String jwt;

    public String getJwt() {
        return this.jwt;
    }
}

유저의 토큰 정보를 저장하는 간단한 JwtContext 클래스입니다.
 

[JwtContextHolder.java]

public class JwtContextHolder {

    private static final ThreadLocal<JwtContext> jwtContextHolder = new ThreadLocal<>();

    public static void setJwtContext(JwtContext jwtContext) {
        jwtContextHolder.set(jwtContext);
    }

    public static JwtContext getJwtContext() {
        JwtContext jwtContext = jwtContextHolder.get();
        if (jwtContext == null) {
            return new JwtContext("No token!");
        }
        return jwtContext;
    }

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

JwtContextHolder는 ThreadLocal을 활용해서 유저의 JwtContext를 저장하고 관리할 수 있는 클래스입니다.
 
그리고 아래와 같이 Spring MVC에서 서블릿 필터를 통해 유저의 Token 정보를 저장할 수 있는 Interceptor 클래스를 하나 생성하겠습니다.
 

[JwtContextInterceptor.java]

import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import org.springframework.web.servlet.HandlerInterceptor;

public class JwtContextInterceptor implements HandlerInterceptor {

    private static final String AUTHORIZATION_NAME = "X-AUTH";

    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) {
        String token = request.getHeader(AUTHORIZATION_NAME);
        JwtContext jwtContext = new JwtContext(token);
        JwtContextHolder.setJwtContext(jwtContext);

        return true;
    }
}

 
해당 Interceptor 클래스와 JwtContextHolder가 정상적으로 동작하는지 확인하기 위해 아래 테스트 코드를 수행해보겠습니다.

public class ThreadLocalTest {

    private JwtContextInterceptor interceptor;

    @Mock
    private Object handler;

    @BeforeEach
    public void setUp() {
        MockitoAnnotations.initMocks(this);
        interceptor = new JwtContextInterceptor();
    }

    @Test
    void jwt_token_interceptor_working_test() {

        MockHttpServletRequest request = new MockHttpServletRequest();
        MockHttpServletResponse response = new MockHttpServletResponse();

        // Set the "X-AUTH" header in the mock request
        request.addHeader("X-AUTH", "your-jwt-token-here");

        boolean result = interceptor.preHandle(request, response, handler);

        assertTrue(result); // Ensure that preHandle returns true

        JwtContext jwtContext = JwtContextHolder.getJwtContext();
        assertNotNull(jwtContext);

        assertEquals("your-jwt-token-here", jwtContext.getJwt());
    }
    
    // .. 생략
}

테스트가 성공한다.

해당 테스트에서 Interceptor는 true를 반환하고, JwtContextHolder에는 우리가 요청한 토큰 정보가 들어있음을 검증했습니다.

 

여기까지는 아무런 문제가 없습니다.
 

3-2. ThreadPool에서 재사용하는 Thread 문제

하지만 쓰레드를 재사용할 수 있는 ThreadPool을 사용할 때는 ThreadLocal의 로컬 변수 역시 재사용될 수 있다는 것을 기억해야 하는데요.
 
아래 코드를 보면서 이 때 발생할 수 있는 문제를 살펴보겠습니다.

 

아래와 같이 3개의 쓰레드를 재사용하는 ThreadPool을 만들고 테스트를 실행해봅니다.

@Test
    void jwt_token_interceptor_with_threadpool_working_test() {
        ExecutorService executorService = Executors.newFixedThreadPool(3);

        for (int threadCount = 0; threadCount < 5; threadCount++) {
            int token_postfix = threadCount;
            executorService.execute(() -> {
                System.out.println("jwt_before_request = " + JwtContextHolder.getJwtContext().getJwt());

                String token = "my-token" + "-" + token_postfix;
                MockHttpServletRequest request = new MockHttpServletRequest();
                MockHttpServletResponse response = new MockHttpServletResponse();
                request.addHeader("X-AUTH", token);
                interceptor.preHandle(request, response, handler);

                System.out.println("jwt_after_request = " + JwtContextHolder.getJwtContext().getJwt());
            });
        }

        executorService.shutdown();

        // 스레드 풀 종료 대기
        while (true) {
            try {
                if (executorService.awaitTermination(10, TimeUnit.SECONDS)) {
                    break;
                }
            } catch (InterruptedException e) {
                System.err.println("Error: " + e);
                executorService.shutdownNow();
            }
        }
        System.out.println("All threads are finished");
    }

해당 테스트 코드에서는 각 Thread의 실행 시점에 JwtContextHolder에 담긴 토큰 정보를 출력해보는데요.
 
일반적인 비즈니스 로직을 생각해보면 각 쓰레드는 고유한 유저의 토큰 정보를 가지고 있다는 것을 기대하지만, 해당 테스트 코드를 실행해보면 특정 쓰레드가 호출될 때 그 이전의 쓰레드에서 할당한 토큰 값(my-token-0, my-token-2)이 이미 ThreadLocal에 할당된 것을 볼 수 있습니다.

Request 이전에 쓰레드에 토큰값이 할당되어 있다.

 

만약, 어떤 쓰레드에서 유저의 토큰 정보를 확인한 뒤 API를 호출하는 등의 작업이 이어진다면 제대로 된 토큰이 요청으로 들어오지 않았음에도 불구하고 그 다음 작업이 이어질 수 있는 장애가 발생할 수 있다는 것을 예측해볼 수 있습니다.
 
뿐만 아니라, ThreadLocal에 할당된 토큰 정보를 제대로 삭제하지 않으면 Memory Leak이 발생할 수 있는 지점이 됩니다.
 
따라서, 쓰레드풀을 사용하면서 ThreadLocal을 사용할 때는 remove() 메서드를 호출하지 않으면 예외를 발생시키거나 remove() 메서드의 호출을 강제하는 코드를 작성하는 것을 권장합니다.
 
예를 들어 위에서 작성했던 Interceptor 코드에 아래 내용을 추가하고 테스트 코드를 조금 수정해보겠습니다.

public class JwtContextInterceptor implements HandlerInterceptor {

    //.. 기존과 동일

    public void afterCompletion(HttpServletRequest request, HttpServletResponse response, Object handler,
                                 @Nullable Exception ex) {

        JwtContextHolder.clear();
    }
}

@Test
void jwt_token_interceptor_with_threadpool_working_test() {
    ExecutorService executorService = Executors.newFixedThreadPool(3);

    for (int threadCount = 0; threadCount < 5; threadCount++) {
        int token_postfix = threadCount;
        executorService.execute(() -> {
            System.out.println("jwt_before_request = " + JwtContextHolder.getJwtContext().getJwt());

            String token = "my-token" + "-" + token_postfix;
            MockHttpServletRequest request = new MockHttpServletRequest();
            MockHttpServletResponse response = new MockHttpServletResponse();
            request.addHeader("X-AUTH", token);
            interceptor.preHandle(request, response, handler);
            System.out.println("jwt_after_request = " + JwtContextHolder.getJwtContext().getJwt());
            interceptor.afterCompletion(request, response, handler, null);
        });
    }
    
    // .. 생략

요청이 끝나면 토큰이 초기화된다.

쓰레드는 자신의 ThreadLocal에 할당된 값을 remove()하고 종료합니다. 
 
따라서, 테스트 결과를 보면 재사용되는 쓰레드에서도 현재 ThreadLocal에 있는 JwtToken이 초기화되어있어 No token! 이라는 메시지가 출력됩니다.
 
약간은 인위적인 코드지만 재사용되는 쓰레드에서 ThreadLocal을 사용할 때 어떤 점을 주의해야 하는지 이해하면 좋을 것 같습니다.
 

 

3-3. 부모 Thread의 ThreadLocal VS 자식 Thread의 ThreadLocal 문제

두번째 문제는 부모 Thread와 자식 Thread 사이에서 발생할 수 있는 ThreadLocal 문제입니다.

 

기본적으로 부모 쓰레드와 자식 쓰레드는 서로 다른 쓰레드이므로 서로 다른 ThreadLocal 값을 가지고 있습니다.

 

즉, 자식 쓰레드는 부모 쓰레드의 ThreadLocal에 접근할 수 없습니다.

@Test
void threadlocal_with_parent_and_child_thread() {
    JwtContext jwtContext = new JwtContext("myToken");
    JwtContextHolder.setJwtContext(jwtContext);

    // Main 쓰레드에서의 ThreadLocal
    Thread currentThread = Thread.currentThread();
    System.out.println("currentThread = " + currentThread);
    System.out.println("currentThread jwt = " + JwtContextHolder.getJwtContext().getJwt());

    // 중간 비즈니스 로직..

    // ThreadPool 에서 ThreadLocal
    ExecutorService executorService = Executors.newFixedThreadPool(1);
    executorService.execute(() -> {
        Thread executorThread = Thread.currentThread();
        System.out.println("executorThread = " + executorThread);
        System.out.println("currentThread jwt = " + JwtContextHolder.getJwtContext().getJwt());
    });
}

그래서 위 코드를 보면 [1] Main 쓰레드에서 JwtContextHolder에 'myToken' 이라는 JwtContext를 저장하고 [2] ThreadPool 안에서는 해당 유저에 대한 병렬적인 작업(API 호출 등)을 처리하고자 하는 의도로 작성되었으나, [3] ThreadPool 안에서 생성되는 Thread는 부모 쓰레드에서 갖고 있던 'myToken' 이라는 값에 접근할 수 없고 코드를 실행해보면 아래와 같은 결과가 출력되게 됩니다.

자식 쓰레드에서 부모 쓰레드의 ThreadLocal에 접근하지 못한다.

 

 

이러한 문제는 parent Thread의 정보를 그대로 상속받는 InheritableThreadLocal을 사용해서 해결할 수 있는데요.

public class JwtContextHolder {

    private static final ThreadLocal<JwtContext> jwtContextHolder = new InheritableThreadLocal<>();

InheritableThreadLocal은 ThreadLocal을 확장한 클래스로 상위 쓰레드에서 하위 쓰레드로의 값의 상속을 제공하여 하위 쓰레드가 생성되면 상위 쓰레드에 저장된 값을 상속받아서 사용할 수 있습니다.

 

그래서 위와 같이 JwtContextHolder의 코드를 수정하고 테스트 코드를 실행하면 아래와 같이 쓰레드는 다르지만 동일한 토큰값에 접근하는 것을 확인할 수 있습니다.

InheritableThreadLocal 사용 시 자식 쓰레드가 부모 쓰레드의 ThreadLocal에 접근할 수 있다.

 

 

4. 정리 / Reference

이번 글에서는 Java ThreadLocal의 기본적인 개념과 사용방법, 그리고 쓰레드풀에서 ThreadLocal을 사용할 때 발생할 수 있는 문제점에 대해 살펴봤습니다.
 

ThreadLocal은 쓰레드 별로 변수를 관리할 수 있다는 편의를 제공하지만 이 말을 반대로 해석하면 "Thread가 달라지는 순간마다 ThreadLocal에 어떤 값이 할당되어있을까?"에 대한 고민과 주의를 필요로 합니다.

 

적절한 상황에서 신중하게 ThreadLocal을 사용하면 좋을 것 같습니다.

 
- https://madplay.github.io/post/java-threadlocal
- https://www.baeldung.com/java-threadlocal