[속닥속닥] 🚀 우당탕탕 쿼리 카운터 개발기

헌치·2022년 11월 1일

속닥속닥 우테코

우아한테크코스

목록 보기

24/30

1. 들어가며✨

속닥속닥은 Spring Data JPA를 사용했습니다.

그말인 즉슨, N+1 문제를 개선하고, 쿼리 개수를 모니터링할 필요가 있다는 뜻이었습니다.
이에 저와 이스트는 특정 커넥션에서 사용된 쿼리 개수를 세는 "쿼리 카운터"를 개발했습니다!

~~sql의 count() 아닙니다~~

쿼리 카운터를 구현하는 방법으로 고려했던 것은 총 2가지였습니다.

하이버네이트의 StatementInspector 사용하기

직접 AOP 기반 쿼리카운터 만들기

저희는 2가지 방법 중 2번을 선택했습니다.

1번이 구현하기엔 더 편리합니다.
하지만 이후 하이버네이트 이외의 구현체를 사용하거나 사정상 직접 JdbcTemlate을 사용했을 경우에도 범용성있는 쿼리 카운터를 구현하고 싶어 AOP로 구현했습니다.

다음은 쿼리 카운터 세팅 과정입니다!

2. 세팅 과정

1. Spring AOP를 통한 PerformanceAspect 구현

우선 저희 쿼리카운터는 AOP 기반으로 구현했습니다.
AOP는 기존 기능을 담은 객체(타겟)의 프록시에 부가기능(관점, Aspect)을 추가해주는 방식입니다.

스프링에서는 어노테이션을 통한 쉬운 AOP 구현을 지원합니다.
하지만 문제가 있었는데요, 원 객체가 Bean으로 등록된 싱글턴 객체여야 한다는 점입니다.
저희가 타겟으로 선정하고 싶은 건 Connection이고, 당연히 싱글턴 객체가 아닙니다.
그렇다면 저희의 AOP를 통한 쿼리카운팅 관점 적용 계획은 물거품이 된 걸까요?😢

다행히 아닙니다!😄
DataSource에서 getConnection() 메소드를 통해 Connection을 얻게 되는데요.
이 DataSource가 Bean 등록 되어있기 때문입니다!

만약 getConnection()이 호출될 때마다,
진짜 커넥션이 아닌
이를 낚아채 쿼리 카운터 기능을 추가한 프록시 커넥션 객체를 리턴한다면?!

다음은 해당 아이디어를 통해 프록시 커넥션을 리턴하는 클래스, PerformanceAspect 입니다.

import java.lang.reflect.Proxy;
import java.sql.Connection;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.springframework.stereotype.Component;

@Aspect
@Component
@Slf4j
public class PerformanceAspect {

    private final ThreadLocal<QueryCounter> queryCounter = new ThreadLocal<>();

    @Pointcut("execution(* javax.sql.DataSource.getConnection(..))")
    public void performancePointcut() {
    }

    @Around("performancePointcut()")
    public Object start (ProceedingJoinPoint point) throws Throwable {
        final Connection connection = (Connection) point.proceed();
        queryCounter.set(new QueryCounter());
        final QueryCounter counter = this.queryCounter.get();

        final Connection proxyConnection = getProxyConnection(connection, counter);
        queryCounter.remove();
        return proxyConnection;
    }

    private Connection getProxyConnection(Connection connection, QueryCounter counter) {
        return (Connection) Proxy.newProxyInstance(
                getClass().getClassLoader(),
                new Class[]{Connection.class},
                new ConnectionHandler(connection, counter)
        );
    }
}

2. ConnectionHandler를 통한 다이나믹 프록시 구현😄

PerformanceAspect 클래스 속 getProxyConnection의 구현이

return (Connection) Proxy.newProxyInstance(
                getClass().getClassLoader(),
                new Class[]{Connection.class},
                new ConnectionHandler(connection, counter)
        );

다음과 같이 되어있는데요. 다이나믹 프록시를 통해 커넥션의 프록시 객체를 생성했기 때문입니다.

다이나믹 프록시란, 런타임 중에 인터페이스가 구현된 인스턴스를 동적으로 구현한 프록시를 의미합니다.

프록시 객체를 직접 생성하기 위해서는 타겟의 모든 메소드에 대해 직접 구현해줘야 해 불편합니다.
특히 저희가 프록시로 만들어야 할 Connection의 경우, 수십 개의 메소드가 있는데요.
이를 모두 구현한 프록시 객체라니... 상상만 해도 구현하기 귀찮네요.

그래서 다이나믹 프록시를 통해 프록시 커넥션 객체를 만들어 줬습니다.

세팅방법

다이나믹 프록시는 Java java.lang.reflect.Proxy package에서 제공해주는 API를 이용하여 손쉽게 만들 수 있습니다.

준비물은 총 세가지입니다.
클래스로더, 원 타겟 클래스(인터페이스), 동적으로 프록시를 설정할 ConnectionHandler입니다.
이중 직접 구현이 필요한 ConnectionHandler로 들어가 보겠습니다.

@Slf4j
public class ConnectionHandler implements InvocationHandler {

    private final Object target;
    private final QueryCounter counter;

    public ConnectionHandler(Object target, QueryCounter counter) {
        this.target = target;
        this.counter = counter;
    }

    @Override
    public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
        countPrepareStatement(method);
        logQueryCount(method);
        return method.invoke(target, args);
    }

    private void logQueryCount(Method method) {
        if (method.getName().equals("close")) {
            warnTooManyQuery();
            log.info("\n====== count : {} =======\n", counter.getCount());
        }
    }

    private void countPrepareStatement(Method method) {
        if (method.getName().equals("prepareStatement")) {
            counter.increase();
        }
    }

    private void warnTooManyQuery() {
        if (counter.isWarn()) {
            log.warn("\n======= Too Many Query !!!! =======");
        }
    }
}

해당 핸들러는 InvocationHandler 인터페이스의 구현이 꼭 필요합니다.

이중 invoke() 메소드가 핵심입니다.

해당 메소드는 구현할 타겟 클래스의 모든 메소드에 직접 접근하고, 동적으로 리턴값을 바꾸거나 부가관점을 추가할 수 있습니다.

저희는 connection 객체가 prepareStatement() 메소드를 호출할 때마다 쿼리 카운트를 증가시켰습니다.
커넥션을 통해 쿼리가 실행될 때마다 해당 메소드가 한번 씩 호출된다는 점에 착안했습니다.

3. QueryCounter 구현😄

이제 대망의 QueryCounter 입니다.

별 구현은 필요 없고, count 변수를 감싼 객체입니다.
increase() 를 통해 쿼리 실행 수를 1 증가시키고
isWarn() 을 통해 11개 이상의 쿼리 카운트 여부를 반환합니다.

public class QueryCounter {

    private int count;

    public void increase() {
        count++;
    }

    public int getCount() {
        return count;
    }

    public boolean isWarn() {
        return count > 10;
    }
}

4. ThreadLocal을 통한 스레드 별 카운터 객체 자원 할당😄

자, 이제 해당 카운터 객체가 스레드별로 할당되게 해야 합니다.

ThreadLocal을 이용하면 되는데요.
스레드 영역에 변수를 설정할 수 있기 때문에, 특정 스레드가 실행하는 모든 코드에서 그 쓰레드에 설정된 변수 값을 사용할 수 있게 됩니다.

사용법

ThreadLocal.set() : 현재 쓰레드의 로컬 변수에 값을 저장한다.
ThreadLocal.get() : 현재 쓰레드의 로컬 변수 값을 읽어온다.
ThreadLocal.remove() : 현재 쓰레드의 로컬 변수 값을 삭제한다.

@Slf4j
public class ConnectionHandler implements InvocationHandler {

    ...

    @Override
    public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
        countPrepareStatement(method);
        logQueryCount(method);
        return method.invoke(target, args);
    }

    private void logQueryCount(Method method) {
        if (method.getName().equals("close")) {
            warnTooManyQuery();
            log.info("\n====== count : {} =======\n", counter.getCount());
        }
    }

    private void warnTooManyQuery() {
        if (counter.isWarn()) {
            log.warn("\n======= Too Many Query !!!! =======");
        }
    }
}

저희는 connection이 끝날 때에 쿼리 카운트 개수를 로그로 보여주고 싶었습니다.

이에 다이나믹 프록시에서,
커넥션 속 close() 메소드가 실행되었을 때 로그가 출력되도록 구현했습니다!

3. 앞으로의 방향

1. 쿼리 정보도 필요하다👀

현재 너무 많은 쿼리가 카운트되었을 때 경고해주는 warnTooManyQuery() 메소드에서는, 해당 쿼리나 메소드 관련 정보가 제공되지 않는데요.
한번 더 다이나믹 프록시를 적용하던가 해서, 마지막 쿼리 상태라도 확인할 수 있게 로깅하는 방법까지 고려하고 있습니다.