들어가기에 앞서..

---

회사 프로젝트에서 API Request Header 에 포함되어 들어온 사용자정보를 Interceptor 에서 검증하고, 이를 UserContext 에 담아 전역적으로 사용하기 위한 구성을 했고, 이를 구현하는데 ThreadLocal<UserContext> 를 사용했다.

Request Per Thread 방식으로 동작하는 Tomcat 기반의 서버였기 때문에 예상한대로 동작했다.

WebFlux 에서 기본적으로 사용되는 Event-Loop 방식의 Netty 에서는 어떤 문제가 발생하고, 이를 어떤 방식으로 해결해 나갈수 있는지 고민한 과정을 포스팅에 담아보려고 한다.

ThreadLocal 이란?

---

ThreadLocal 은 쉽게 말해 Thread 별로 가지는 변수다. 내부적으로 Thread ID 를 Key 로하는 Map<K, V> 에 데이터를 저장하기 때문에 동일한 Thread 라면 어디서든 해당 값에 접근할 수 있다.

이러한 ThreadLocal 은 실제로 Spring Security 의 SecurityContext, Spring Web MVC 의 RequestContext, Logback 의 MDC 등에서 사용되고 있는 방식이다.

WebFlux 에서 사용해보자

---

Request Per Thread 에서는 요청이 들어오고 응답이 나갈때까지 동일한 Thread 가 요청을 처리하기 때문에 ThreadLocal 을 사용해도 문제가 없다.

하지만, Non-Blocking Event-Loop 방식으로 동작하는 WebFlux(Netty) 에서는 하나의 요청을 처리되는 Thread 가 계속 변경될 수 있기 때문에 ThreadLocal 에 저장된 값이 없거나 뒤죽박죽 변경될 수 있다. 따라서 이러한 처리되는 Thread 가 변경되는 환경에서는 ThreadLocal 사용에 반드시 주의를 기울여야 한다.

예를 들어, ThreadLocal 에 데이터를 넣는 시점에는 Thread-A 였지만, 요청이 처리되는 중간에 Thread-B 로 변경되고 ThreadLocal 에서 데이터를 꺼낸다면 기존에 저장된 데이터가 없을 것이다. 심지어 다른 요청이 저장한 값이 잘못 확인될 수도 있다.

실제로 이커머스를 대표하는 모 대기업에서 Spring Web MVC 에서 Spring WebFlux 로 전환하던 중 ThreadLocal 을 고려하지 못해 마이페이지에 다른 사람의 정보가 노출되는 이슈가 있었다고 한다.

이제 테스트를 통해 실제로 이런 문제가 발생하는지 확인해보고 개선해보자.

테스트 환경은 아래와 같다.

@GetMapping("/test/{id}") 
fun startTest(@PathVariable id: String): Mono<String> {
	printLog(id, "start")
    return webClient.get().uri("/delay/2").exchangeToMono {
    	printLog(id, "end")
        it.bodyToMono(String::class.java)
    }
}

이런 구조에서 Spring MVC Service 의 /delay/{seconds} API 에 2초의 딜레이를 준다면 2초 조금 넘는 정도에서 100개의 요청이 모두 처리된다. 실제로 처리된 로그를 보자.

동일한 요청임에도 start, end 를 처리하는 Thread 가 다르다는 것을 확인할 수 있다. 이런 상황에서 ThreadLocal 에 데이터를 넣는다면 값이 없거나 다른 요청이 넣은 값을 잘못 가져오는 등의 문제가 발생할 수 있다.

이제 ThreadLocal 변수를 만들어서 테스트해보자.

object TestContextHolder {
	private val context: ThreadLocal<Long> = ThreadLocal()
    
    fun setContext(id: Long) = context.set(id)
    fun getContext(): Long = context.get()
    fun clearContext() = context.remove()
}

fun printLog(id: String, status: String) {
	logger.info("$status -> ${Thread.currentThread().name} $id ${TestContextHolder.getContext()}")
}

요청에 포함된 Path Variable 값이 ID 를 ThreadLocal 변수에 저장하고 꺼내서 로그를 찍는 방식으로 확인해보자.

start 시점에는 파라미터로 받은 id 와 Context 에 저장된 값이 일치하지만, end 시점에는 Context 에 저장된 값이 id 와 다르다는 것을 확인할 수 있다. 처리하는 Thread 가 변경되면서 기존에 저장한 값이 아니라 다른 요청에 의해서 저장된 다른 값이 나오게된다.

InheritableThreadLocal 을 통해서 테스트도 해봤는데 마찬가지로 동일한 결과가 나왔다.

당연한 결과인데 Event-Loop 에서 사용되는 각 Thread 가 서로 부모-자식 관계가 아니기 때문에 이 상황에서의 해결책이 될 수 없다.

그럼 이 문제를 어떻게 해결해야할까?

Reactor 의 contextWrite & deferContextual

---

Reactor 에서는 contextWrite 와 deferContextual 을 이용해 ThreadLocal 문제를 해결한다. 옛날 버전에서는 contextWrite 대신 subscriberContext 을 사용했는데 현재는 Deprecated 상태이다.

val key: String = "KEY"
val result: Mono<String> = Mono.deferContexctual { ctx -> Mono.jus("Hello " + ctx.get(KEY) }
				.contextWrite { ctx -> ctx.put(KEY, "World") }

결론부터 말하면 최종적으론 이런 형태의 구조를 만들어야한다.
저장하고 싶은 Context 정보를 실행중인 Thread 에 의존적인 ThreadLocal 에 저장하는 것이 아니라 Flow 자체에 포함시킬 수 있는 별도 Context 로 등록해줘야하는 것이다.

간략한 구조를 설명하면 아래와 같다.

1. deferContextual { } 을 통해서 Context 에 값이 들어있다고 가정하고 로직을 작성한다.
2. 뒤에서 contextWrite 를 통해서 실제 값을 넣어준다.

순서를 주의하자. Context 에서 값을 꺼내는 로직이 먼저오고 넣어주는 코드가 뒤에와서 헷갈리기 쉽다. 이러한 구조를 사용하는 이유는 Reactor 의 Flow 는 Subscribe 가 발생하는 순간 흘러가는 구조를 가지기 때문이다.

위 구조를 사용해 Spring MVC 에서 사용하던 RequestContextHolder & RequestContext 구조를 만들어보자!

먼저 Request 정보를 담고, 전역적으로 활용할 수 있도록 RequestContextHolder 객체를 만들어보자.

object RequestContextHolder {
	const val CONTEXT_KEY = "REQUEST_CONTEXT"
    fun getRequest(): Mono<ServerHttpRequest> =
    	Mono.deferContextual { ctx -> Mono.just(ctx[CONTEXT_KEY]) }
}

그리고 해당 Context 에 값을 채워주기 위한 Interceptor 를 만들었다.

@Configuration
@ConditionalOnWebApplication(type = ConditionalOnWebApplication.Type.REACTIVE)
class RequestContextInterceptor: WebFilter {
	override fun filter(exchange: ServerWebExchange, chain: WebFilterChain): Mono<Void> {
    	val request = exchange.request
        return chain.filter(exchange)
        	.contextWrite { ctx -> ctx.put(RequestContextHolder.CONTEXT_KEY, request) }
    }
}

Interceptor 에서 요청을 가로채 해당 요청을 contextWrite 를 통해 넣어준다. chain.filter(exchange) 뒤에 넣어줘야하는 것이 중요하다. 앞서 설명했듯이 우선 Context 에 값이 있다는 가정하에 로직이 작성되고 실제로 값을 넣어주는 부분이 뒤에있어야 하기 때문이다.

그럼 이제 실제 API 에서 RequestContextHolder 에서 값을 꺼내 사용해보도록 하자.

@GetMapping("/request")
fun request(@RequestParam id: String): Mono<String> {
	return RequestContextHolder.getRequest()
    	.map { req -> req.uri.query }
}

항상 RequestContextHolder.getRequest() 이 반환해주는 Mono 객체를 통해 Flow 를 만들어나가야한다.

그럼 가장 처음 소개했던 이슈를 해결하고 마무리해보자!

마지막 테스트

---

object UserContextHolder {
	private const val CONTEXT_KEY = "USER_CONTEXT"
    
    fun setContext(userContext: UserContext, context: Context) = context.put(CONTEXT_KEY, userContext)
    fun getContext(): Mono<UserContext) = 
    	Mono.deferContextual { context -> Mono.just(context[CONTEXT_KEY]) }
}

기존에 ThreadLocal 을 활용해 만들었던 ContextHolder 는 이런식으로 변경된다.

@Configuration
@ConditionalOnWebApplication(type = ConditionalOnWebApplication.Type.REACTIVE)
class UserContextInterceptor(
	val objectMapper: ObjectMapper
): WebFilter {
	
    override fun filter(exchange: ServerWebExchange, chain: WebFilterChain): Mono<Void> {
    	val xPayAccount = exchange.request.headers.getFirst("x-kakao-pay-account")
        val payAccount = try {
        	objectMapper.readValue(xPayAccount, PayAccount::class.java)
        } catch (e: Exception) {
        	throw RuntimeException()
        }
       	return chain.filter(exchange)
        	.contextWrite { ctx -> UserContextHolder.setContext(UserContext(payAccount), context) }
    }

}

UserContextInterceptor 에서 Request Header 에 포함된 사용자 정보를 contextWrite 를 통해 UserContext 에 값을 저장한다.

그리고 어플리케이션 코드에서는 UserContextHolder.getContext() 호출시 Mono.deferContextual 을 통해 UserContext 에 저장된 사용자 정보를 꺼내 사용할 수 있다.

아래는 테스트할 때 사용할 코드이다.

@GetMapping("/test/{id}")
fun test(@PathVariable id: String): Mono<String> = UserContextHolder.getContext().map { userContext ->
	printLog(id, "start", userContext)
    userContext
}.flatMap { userContext ->
	webClient.get().uri("/delay/2").exchangeToMono { resp ->
    	resp.bodyToMono(String::class.java).doOnNext { printLog(id, "end", userContext) }
    }
}

포스팅 처음에서 테스트했던 코드와 동일한 기능을 수행한다.

그 결과는...

start 에서 넣었던 ID 가 동일하게 end 시점에도 확인되는 것을 확인할 수 있다.

또 다른 예시로 Logback 의 MDC 를 WebFlux 환경에서 사용하기 위한 코드를 소개하면 이번 포스팅을 마친다.

@Component
class CustomHttpHandlerDecorator: HttpHandlerDecoratorFactory {

	override fun apply(httpHandler: HttpHandler): HttpHandler {
    	return (request, response) -> {
        	httpHandler.handle(request, response).contextWrite(ctx -> {
            	val traceId: String = getTraceId()
                MDC.put(MDC_KEY_TRACE_ID, traceId)
                return Context.of(MDC_KEY_TRACE_ID, traceId)
            })
        }
    }
    
    private fun getTraceId(): String {
    	return UUID.randomUUID().toString()
    }
    
	companion object {
    	private const val MDC_KEY_TRACE_ID = "traceId"
    }
}