Task Decorator (@Async 에서 MDC, SecurityContextHolder 정보 활용)

Task Decorator

• ThreadLocal은 스레드 별로 값이 저장되므로, @Async를 이용할 경우 다음의 데이터가 공유되지 않는 문제가 있다.
  • MDC를 이용하여 저장한 항목
  • SecurityContext 에 저장된 인증된 사용자 정보
• 이를 해결하기 위해 Task Decorator를 이용한다. 예는 다음과 같다.

 

• ThreadLocal이 정의 되어 있는 클래스

public class UserContextHolder {

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

    public static void set(String userId) {
        USER_ID.set(userId);
    }

    public static String get() {
        return USER_ID.get();
    }

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

 

• Main 스레드에 있는 ThreadLocal을 이용하여 runnable의 ThreadLocal에 값을 복사하는 부분

import org.springframework.core.task.TaskDecorator;

import java.util.Map;
import org.slf4j.MDC;

public class ContextCopyingTaskDecorator implements TaskDecorator {

    @Override
    public Runnable decorate(Runnable task) {

        // 1. 부모 스레드의 컨텍스트 캡처
        String userId = UserContextHolder.get();
        Map<String, String> mdcContext = MDC.getCopyOfContextMap();

        return () -> {
            try {
                // 2. 자식 스레드에 컨텍스트 설정
                if (userId != null) {
                    UserContextHolder.set(userId);
                }
                if (mdcContext != null) {
                    MDC.setContextMap(mdcContext);
                }

                // 3. 실제 작업 실행
                task.run();

            } finally {
                // 4. 반드시 clean
                UserContextHolder.clear();
                MDC.clear();
            }
        };
    }
}

• decorate 메서드 에서 runnable을 리턴하기전, main에 있는 ThreadLocal 항목을 불러와 값을 저장한다. 여기까지는 메인 스레드 영역
• return 에서 불러온 값을 다시 ThreadLocal에 추가 하고, runnable을 리턴한다. 여기부터 runnable의 실행인 새 스레드 영역
• 실행완료시 초기화 하는 코드를 넣는다.

 

매번 Async 작업을 할 때마다 위에 정의한 클래스를 이용하라면 매우 번거롭다. 따라서 @Async 설정에서 일괄 적용하도록 할 수도 있다. 

import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Bean;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

@Configuration
public class AsyncExecutorConfig {

    @Bean
    public ThreadPoolTaskExecutor asyncExecutor() {
        ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();

        executor.setCorePoolSize(5);
        executor.setMaxPoolSize(10);
        executor.setQueueCapacity(100);
        executor.setThreadNamePrefix("async-");

        executor.setTaskDecorator(new ContextCopyingTaskDecorator());

        executor.initialize();
        return executor;
    }
}

 

이렇게 할 경우, Async 동작시 일괄적으로 적용되어서 이용 할 수 있게 된다.

혹은 AsyncConfigurer 를 이용해 비슷하게 설정도 가능 하다. (https://dean83.tistory.com/393)