계기

 N+1 문제가 발생했는지 알아보기 위해 쿼리의 개수를 세는 과정에서 불편함이 있었다. 특히 콘솔 창에서 일일이 수동으로 개수를 세다보니 실수할 여지가 많았다. 따라서 자동으로 쿼리의 개수를 측정해주는 쿼리 카운터를 개발하고 싶었다. 먼저 하이버네이트의 StatementInspector 인터페이스를 적용하면 쿼리의 개수를 간단히 셀 수 있지만 하이버네이트에 종속적이라 다른 구현체 혹은 JdbcTemplate을 사용하는 경우에는 카운트가 되지 않아 범용적이지 못하다고 판단하여 배제하였다. 나는 Spring aop와 다이나믹 프록시 기술을 활용하여 쿼리 카운터를 개발해보았다.

참고 : API의 쿼리 개수 세기 - (1) 하이버네이트를 이용한 카운팅

핵심 아이디어

(1) 카운트를 세는 시점

 먼저 클라이언트로부터 요청이 들어와 DB로 sql을 보내는 과정을 간략하게 살펴보자.

public class MemberRepositoryV0 {

public Member save(Member member) throws SQLException {
	String sql = "insert into member(member_id, money) values(?, ?)";
 	Connection con = null;
 	PreparedStatement pstmt = null;
 	try {
    	con = getConnection();
        pstmt = con.prepareStatement(sql);
        pstmt.setString(1, member.getMemberId());
        pstmt.setInt(2, member.getMoney());
        pstmt.executeUpdate();
 		return member;
    } 
    ...    
}

1. 클라이언트가 트랜잭션 AOP 프록시 호출
2. 트랜잭션 AOP 프록시가 스프링 컨테이너로부터 트랜잭션 매니저 획득
3. 트랜잭션 매니저가 데이터소스로부터 Connection 획득
4. 트랜잭션 매니저가 Connection을 트랜잭션 동기화 매니저(스레드 로컬)에 저장
5. 리포지토리에서 Connection 획득
6. 커넥션의 prepareStatement에 sql 저장

   ex) con.prepareStatement(sql);

7. executeUpdate, execute, executeQuery 메서드를 통해 데이터베이스에 sql 전송
8. 트랜잭션이 끝나면 해당 커넥션을 커넥션 풀에 반환

 처음에는 datasource로부터 Connection을 얻어올때 쿼리 카운트를 1 증가시키면 되지 않을까라고 생각했는데 트랜잭션을 적용하면 해당 트랜잭션 내에서는 Connection 동기화해서 같은 커넥션을 사용하므로 적합하지 않다. 따라서 Connection 객체로부터 prepareStatement 메서드를 호출할때 쿼리 카운트를 1 증가시키기로 했다.

(2) 프록시

datasource객체 -> Connection 객체 -> prepareStatement() 호출

 앞서서 Connection 객체로부터 prepareStatement()를 호출할때 쿼리 카운트를 1 증가시키로 결정했었다. 스프링에서는 특정 메서드가 호출되었을때 카운트 값을 증가시키는 것은 부가기능으로 AOP를 활용하면 좋다. 참고로 스프링 AOP는 스프링 컨테이너에 빈으로 등록되어 있어야 사용할 수 있다. datasource는 스프링 컨테이너에 빈으로 등록되어 있지만 Connection 객체는 빈으로 등록할 수 없다. 여기서 AOP 적용과 관련하여 정말 고민이 많았는데 결과론적으로 아래와 같은 플로우를 생각해냈다.

Connection을 프록시 적용할때 3가지 방법이 존재하는데 나는 아래와 같은 이유로 Cglib 다이나믹 프록시를 사용했다. 참고로 ProxyFactory를 사용하면 내부에서 자동으로 Cglib 다이나믹 프록시를 만들어준다. default가 Cglib지만 설정을 통해 Jdk 다이나믹 프록시를 생성할 수도 있다.

1. 직접 프록시 객체 생성 : Connection 인터페이스에 정의되어 있는 메서드들을 추가로 구현해줘야해서 불편하다.
2. JDK 다이나믹 프록시 : Connection이 인터페이스라 JDK 다이나믹 프록시를 생성할 수 있다. 하지만 내부에서 리플랙션 기술을 사용하므로 Cglib보다 비교적 느리다.
3. Cglib 다이나믹 프록시 : 원래는 final 메서드, 기본 생성자 호출 등 여러 문제가 있었는데 objenesis 라이브러리 등장 이후로 전부 해결되었다. 바이트 코드를 활용하므로 JDK 다이나믹 프록시보다 비교적 성능이 좋다.

🤫 JDK 다이나믹 프록시와 Cglib 다이나믹 프록시가 궁금하다면??
참고 : JDK Dynamic Proxy와 CGLIB의 차이점은 무엇일까?

(2) 동시성 이슈

 동시에 api 요청이 올 경우 각각의 api 마다 쿼리 카운트를 측정해야 하므로 스레드별로 카운트 값을 보존하고 있어야 한다. 이와 관련하여 스프링 웹 스코프를 사용하는 방식이랑 ThreadLocal을 사용하는 방식이 존재한다. 큰 차이는 없지만 나는ThreadLocal을 사용하면 스프링 웹 스코프 방식에 비해 순수 자바 환경(테스트)에서도 활용할 수 있어서 법용성 측면에서 더 좋다고 생각했다.

🤙 ThreadLocal을 다 사용하고 반드시 remove 해줘야한다.

결과

• DatasourceAspect : Datasource AOP

@Component
@Aspect
@RequiredArgsConstructor
public class DatasourceAspect {
    private final QueryCount queryCounter;

    @Around("execution(* javax.sql.DataSource.getConnection(..))")
    public Connection getConnection(ProceedingJoinPoint proceedingJoinPoint) throws Throwable {
        Connection connection = (Connection) proceedingJoinPoint.proceed();
        ProxyFactory proxyFactory = new ProxyFactory(connection);
        proxyFactory.addAdvice(new QueryAdvice(queryCounter));
        return (Connection) proxyFactory.getProxy();
    }
}

• QueryAdvice : dynamic proxy(datasource connection)

public class QueryAdvice implements MethodInterceptor{

    private final QueryCount queryCounter;
    public QueryAdvice(QueryCount queryCounter) {
        this.queryCounter = queryCounter;
    }

    @Nullable
    @Override
    public Object invoke(@Nonnull MethodInvocation invocation) throws Throwable {
        countPrepareStatement(invocation.getMethod());
        return invocation.proceed();
    }

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

• QueryCount : ThreadLocal을 활용한 count 값 관리

@Component
public class QueryCount {

    private final ThreadLocal<Integer> count = ThreadLocal.withInitial(() -> 0);

    public void increaseCount() {
        count.set(count.get() + 1);
    }

    public int getCount() {
        return count.get();
    }

    public void remove() {
        count.remove();
    }
}

• QueryCountInterceptor : 쿼리 로그 출력용

@Component
@Slf4j
public class QueryCountInterceptor implements HandlerInterceptor {

    private static final String QUERY_INFO_FORMAT = "QUERY_INFO : [{} {}] [STATUS CODE: {}] [QUERY_COUNT: {}]";
    private final QueryCount queryCounter;

    public QueryCountInterceptor(QueryCount queryCounter) {
        this.queryCounter = queryCounter;
    }

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

        int count = queryCounter.getCount();
        int status = response.getStatus();
        String requestURI = request.getRequestURI();
        String method = request.getMethod();
        log.info(QUERY_INFO_FORMAT, method, requestURI, status, count);
        queryCounter.remove();
    }
}

• 출력