쿼리 카운터 만들기
프로젝트를 하면서 하나의 API에서 발생하는 쿼리들을 한 번에 모아 볼 수 있는 기능이 있다면 좋겠다고 생각했다! 그래서 그 동안 했던 프로젝트 행록과 펫쿠아에는 AOP와 프록시 커넥션을 사용한 쿼리 카운터 로직이 있는데, 이것을 포스팅해보려고 한다.
우선 API 쿼리 정보를 저장할 클래스를 생성한다.
data class QueryInfo(
var count: Int = 0, // 쿼리 개수와
var time: Long = 0L, // 쿼리 소요 시간을 측정할 계획이다!
) {
fun increaseCount() {
count++
}
fun addTime(time: Long) {
this.time += time
}
}여기서 생성한 QueryInfo객체는 ThreadLocal을 사용하여 스레드별로 쿼리 개수와 소요시간을 측정할 예정이다. API는 각자 별도의 스레드에서 실행되므로 ThreadLocal을 활용해 API에서 발생한 쿼리 정보를 측정할 수 있다. (물론 API에서 비동기 메서드를 실행한다면, 하나의 API가 두 개 이상의 Thread를 사용할 것이다. 이 부분은 아래에서 다시 설명!)
PreparedStatement는 SQL 쿼리 내용을 가지고 있으며, 쿼리를 실행할 때 사용된다. PreparedStatement에서 execute() executeQuery(), executeUpdate()와 같이 execute로 시작하는 메서드를 실행하면 쿼리를 실행하고 ResultSet을 받아온다.
때문에 PrepareStatement에서 execute로 시작하는 메서드를 실행할 때 인터셉트하여 execute를 수행하는데 소요된 시간을 전부 측정한다면, 하나의 API에서 발생한 쿼리 소요 시간을 측정할 수 있다.
private const val EXECUTE = "execute"
class PreparedStatementProxyHandler(
private val queryInfo: QueryInfo,
) : MethodInterceptor {
@Nullable
override operator fun invoke(@Nonnull invocation: MethodInvocation): Any? {
var result: Any?
val executionTime = measureTimeMillis {
result = invocation.proceed()
}
val method = invocation.method
if (method.name.contains(EXECUTE)) {
queryInfo.addTime(executionTime)
queryInfo.increaseCount()
}
return result;
}
}다음은 DB 커넥션을 추상화한 객체 Connection을 프록시 객체로 생성한다. 현재 Hikari를 사용하고 있기에 커넥션 중 HikariProxyConnection에만 프록시가 적용되게 한다.
(HikariProxyConnection에만 강제하는 이유는, p6spy등 다른 라이브러리를 사용할 때 P6Connection이 생성되었는데 P6Connection은 실제 DB 커넥션을 감싼 프록시로, 로깅 기능을 추가하기 위해 생성된 커넥션이다. 이렇게 다른 라이브러리등을 사용해 추가적인 프록시 커넥션을 사용하고 있다면, 로그가 프록시 커넥션의 개수만큼 N배 증가할 수 있다. 때문에 실제로 DB와 통신하는 커넥션을 특정하여 프록시로 감싸는 게 확실하다)
커넥션이 HikariProxyConnection일 경우, 앞에서 생성한 프록시 PreparedStatement를 반환하게 설정해준다.
private const val HIKARY_PROXY_CONNECTION = "HikariProxyConnection"
private const val PREPARED_STATEMENT = "prepareStatement"
class ConnectionProxyHandler(
private val connection: Any,
private val queryInfo: QueryInfo,
) : MethodInterceptor {
override fun invoke(invocation: MethodInvocation): Any? {
val result = invocation.proceed()
if (result != null && isPreparedStatement(invocation)) {
val proxyFactory = ProxyFactory(result)
proxyFactory.addAdvice(PreparedStatementProxyHandler(queryInfo))
return proxyFactory.proxy
}
return result
}
private fun isPreparedStatement(invocation: MethodInvocation): Boolean {
val targetObject = invocation.`this` ?: return false
val targetClass: Class<*> = targetObject.javaClass
val targetMethod = invocation.method
return targetClass.name.contains(HIKARY_PROXY_CONNECTION) && targetMethod.name == PREPARED_STATEMENT
}
fun getProxy(): Any {
val proxyFactory = ProxyFactory(connection)
proxyFactory.addAdvice(this)
return proxyFactory.proxy
}
}이제 AOP를 사용하여 DB에서 커넥션을 가져올 때, 위에서 생성한 프록시 커넥션을 가져오게 하면 된다. 프록시 커넥션을 생성할 때, 해당 스레드에 QueryInfo가 존재하지 않다면 새로 생성해서 ThreadLocal에 저장한다.
이제 커넥션에서 PreparedStatement를 호출해 쿼리를 execute할 때마다, 우리가 바라는 대로 프록시 PreparedStatement가 쿼리 개수와 소요시간을 측정헤 QueryInfo에 저장해 줄 것이다! API가 컨트롤러를 빠져나갈 때, QueryInfo의 내용을 로그로 출력해주고 비워주자!
@Aspect
@Component
class QueryCounterAop {
private val log = LoggerFactory.getLogger(QueryCounterAop::class.java)
private val queryInfoStorage = ThreadLocal<QueryInfo>()
@Around("execution( * javax.sql.DataSource.getConnection())")
fun getProxyConnection(joinPoint: ProceedingJoinPoint): Any {
val connection = joinPoint.proceed()
val queryInfo = queryInfoStorage.getOrSet { QueryInfo() }
return ConnectionProxyHandler(connection, queryInfo).getProxy()
}
@After("within(@org.springframework.web.bind.annotation.RestController *)")
fun logQueryInfo() {
val attributes = RequestContextHolder.getRequestAttributes() as? ServletRequestAttributes
val request = attributes?.request
request?.let {
val queryInfo = queryInfoStorage.getOrSet { QueryInfo() }
log.info("METHOD: ${request.method}, URI: ${request.requestURI}, QUERY_COUNT: ${queryInfo.count}, QUERY_TIME: ${queryInfo.time}ms")
}
queryInfoStorage.remove()
}
}API에서 비동기 메서드를 호출한다면?
결론 먼저 말하면, 비동기 메서드에서 발생한 쿼리 로그까지 하나의 API 로그에 포함시키는 것은 포기했다...!
위 방식은 ThreadLocal을 사용해 쿼리 정보를 저장하기에, API가 비동기 메서드를 호출한다면 스레드가 달라져 추적할 수 없다.
비동기 메서드의 쿼리도 전부 합해서 로그로 출력할 순 없을까? 이 조건을 만족하려면, 한 API에서 실행된 모든 스레드가 종료되는 시점에 QueryInfo들을 모아 로그로 출력해야 한다.
자동적으로 @After("within(@org.springframework.web.bind.annotation.RestController *)") 시점에 로그를 출력할 수 없게 된다. 스레드별로 어떤 API에서 발생한 것인지를 저장해두고, 한 API에서 파생된 스레드 정보들을 데이터로 가지고 있어야 한다. 이후, 특정 스레드가 종료될 때, 해당 스레드의 API를 찾고, 해당 API의 스레드들이 전부 종료되었다면 로그를 출력하고, 종료되지 않은 스레드라면 무시하는 로직을 추가해야할 것이다.(이외에 다른 방법은 생각나지 않았다...)
위 방법을 구현하면 쿼리 카운터 로직이 매우 복잡해질 거라고 생각된다. 때문에 비동기 메서드에서 실행된 쿼리를 하나의 로그로 모으지 않고, 각 API에서 수행된 스레드 별로 로그를 각자 출력하기로 했다.
그리고 로그에 API별 고유한 아이디를 포함시켜, 로그 대시보드에서 아이디를 기준으로 검색해서 API에서 발생한 쿼리 정보들을 확인하는 방식을 선택했다.
한 API에서 발생한 모든 쿼리 정보를 하나의 로그로 모으는 것엔 실패했지만, 한 API에서 현재 최대 2개의 스레드를 가지기에 별도의 로그로 출력해도 확인하기에 어려운 부분은 없을거라 생각한다.
