[kerdy] AOP, ThreadLocal을 사용하여 N+1 detector 만들어보기

728x90

글을 쓰게 된 이유

이번 글은 커디 팀뿐만 아니라 우아한테크코스(이하 우테코) 크루들에게 모두 도움이 되지 않는 글일까 생각합니다.

레벨 3 프로젝트를 진행하면서 기존에 JPA를 사용했던 분들, 처음 사용해 보신 분들이 있으실 겁니다.

기존에 사용했던 분들은 N+1을 알고 있지만 기능 만들기에 급해서 일단 구현 먼저 하셨을 것이고,

처음 사용해 보신 분들은 N+1 들어는 봤는데 이걸 어떻게 찾는 거지?라는 분들이 계실 거예요.

레벨 3가 끝나고 N+1을 해결해야 하는데 만든 기능들이 너무 많아서 API 호출하고 나서 로그에 찍힌 쿼리들을 하나하나 읽기 힘들더라고요,,

그래서 이 부분을 어떻게 해결할 수 있을까 생각해 보니 AOP(프록시)를 통해서 해당 API에서 쿼리의 개수를 측정할 수 있는 N+1 detector를 만들어보면 어떨까라는 생각을 했고, 어떻게 만들었는지 소개해보려고 합니다.

이 글을 읽으신 분들이 모두 N+1을 잘 발견하셔서 해결하셨으면 좋겠네요.

사전 지식

먼저 AOP와 프록시를 알아야 합니다.

Spring에서는 AOP를 적용할 때, 프록시 개념을 사용하고 있습니다.

인터페이스에 대한 프록시를 만들 때는 JDK 동적 프록시,

구체 클래스에 대한 프록시를 만들 때는 CGLIB를 기본적으로 사용합니다.

하지만 스프링에서는 인터페이스, 구체 클래스의 프록시를 나눠서 관리하기보다는 한번 더 추상화해서 개발자들에게 더 나은 경험을 제공하기 위해 ProxyFactory 를 제공하고 있습니다.

ProxyFactory를 통해 프록시 클래스를 만들어주면 내부에서 인터페이스면 JDK 동적 프록시를 구체 클래스면 CGLIB를 통해 만들어줍니다.

그래서 이번에 ProxyFactory를 사용하여 프록시 클래스를 만들어서 사용합니다.

저장할 데이터 클래스 정의

query의 개수와 시간, 해당 api를 호출한 http method, url를 기록하기 위해 아래와 같은 클래스를 정의했습니다.

package com.support;

import lombok.Getter;
import lombok.ToString;

@Getter
@ToString
public class LoggingForm {

  private String apiUrl;
  private String apiMethod;
  private Long queryCounts = 0L;
  private Long queryTime = 0L;

  public void queryCountUp() {
    queryCounts++;
  }

  public void addQueryTime(final Long queryTime) {
    this.queryTime += queryTime;
  }

  public void setApiUrl(final String apiUrl) {
    this.apiUrl = apiUrl;
  }

  public void setApiMethod(final String apiMethod) {
    this.apiMethod = apiMethod;
  }
}

이 데이터 클래스를 어떻게 가지고 있을까요?

ThreadLocal를 사용하여 데이터 클래스 관리

많은 API 요청 사이에서 Thread-safe 하게 해당 데이터를 어떻게 관리할 수 있을까 생각해 보았습니다.

Thread-safe 하지 않다면 객관적인 쿼리 시간이나 개수를 얻을 수 없기 때문이죠.

자바에서 제공하고 있는 ThreadLocal을 사용해보려 합니다.

쉽게 말하자면 ThreadLocal은 각 스레드에 고유한 개인 사물함이 있다고 생각하시면 됩니다.

그래서 다른 thread가 액세스 할 수 없어서 변경이나 충돌에 자유롭습니다. 우린 객관적인 데이터를 얻을 수 있는 것이지요.

중요한 점은 메모리 누수가 나지 않게 ThreadLocal을 사용하고 나서 remove를 해줘야 합니다.

아이디어

먼저 이런 ‘부가 기능’들을 추가할 때는 AOP 사용하는 것을 알고 있었습니다.

그래서 AOP를 사용하여 원래의 요청들을 가져와서 쿼리가 나갈 때마다 개수를 세고,

메서드가 시작할 때 시간을 측정하고, 끝날 때 시간을 측정하여 그 차이를 통해서 쿼리 시간을 알 수 있었습니다.

그러면 어떤 요청들을 중간에 가로채야 하는지 먼저 생각해 봤습니다.

우리가 미션 1 때 DB에 접근하기 위해서 아래와 같은 코드를 작성했습니다.

public void addUser(final User user) {
    final  Stringquery = "INSERT INTO user VALUES(?, ?)";
    try (final Connection connection = getConnection();
         final PreparedStatement preparedStatement = connection.prepareStatement(query)) {
        preparedStatement.setString(1, user.userId());
        preparedStatement.setString(2, user.name());
        preparedStatement.executeUpdate();
    } catch (final SQLException e) {
        throw new RuntimeException(e);
    }
}

DriverManager를 통해 Connection을 가져오고, PreparedStatement를 가져와서 executeUpdate 메서드를 통해 쿼리를 보냅니다.

공식문서에서 보면 DriverManager, DataSource 모두 Connection을 가져오는 방법이지만 DataSource가 DriverManager보다 더 빠르게 connection을 가져온다고 합니다.

그러면 DataSource, Connection, PreparedStatement 중에 요청을 가로채서 AOP를 적용해야 합니다. 하지만 스프링 AOP는 스프링 빈에만 적용할 수 있습니다.

DataSource는 스프링 빈으로 등록되는 것을 알고 있었지만 혹시나 Connection이나 PreparedStatement도 등록되어 있는지 확인해 봤습니다.

@Autowired
private ApplicationContext ac;

final DataSource dataSource = ac.getBean(DataSource.class);
  System.out.println("bean.getClass() = " + dataSource.getClass());
//    final Connection connection = ac.getBean(Connection.class);
//    System.out.println("connection. = " + connection.getClass());
  final PreparedStatement preparedStatement = ac.getBean(PreparedStatement.class);
  System.out.println("preparedStatement.getClass() = " + preparedStatement.getClass());

아쉽게도 없더라고요,,

Connection과 PreparedStatement는 우리가 빈으로 등록할 수 없기 때문에 DataSource가 getConnection 하는 요청을 가로채서 쿼리 개수를 세는 기능과, 시간을 측정하는 기능을 부가적으로 넣어야 합니다.

그림을 그리면 위와 같습니다.

Connection의 프록시, PreparedStatement 프록시를 만들어서 쿼리를 보내는 메서드에 해당 프록시를 호출하도록 하면 됩니다.

글로 이해하기 어려우실 테니까 이제 코드를 보도록 하겠습니다.

PreparedStatement Proxy Handler

먼저 PreparedStatement 프록시 클래스를 handle 하는 클래스입니다.

이때는 DB에 직접적인 쿼리를 날리기 때문에 해당 쿼리를 날리는 메서드가 호출되면 시간과 쿼리 개수를 증가시켜 줍니다.

@RequiredArgsConstructor
public class PreparedStatementProxyHandler implements MethodInterceptor { // 1

  private static final List<String> JDBC_QUERY_METHOD =
      List.of("executeQuery", "execute", "executeUpdate");

  private final LoggingForm loggingForm;

  @Nullable
  @Override
  public Object invoke(@Nonnull final MethodInvocation invocation) throws Throwable {

    final Method method = invocation.getMethod(); // 2

    if (JDBC_QUERY_METHOD.contains(method.getName())) { // 3
      final long startTime = System.currentTimeMillis();
      final Object result = invocation.proceed();
      final long endTime = System.currentTimeMillis();

      loggingForm.addQueryTime(endTime - startTime);
      loggingForm.queryCountUp();

      return result;
    }

    return invocation.proceed();
  }
}

(1)

여기서 ProxyFactory를 사용하기 위해서는 MethodInterceptor를 사용해야 합니다.

MethodInterceptor는 CGLIB Proxy를 만들 때도 사용하는 인터페이스(다른 것)여서 아래와 같은 패키지에 있는 MethodInterceptor를 사용해야 합니다.

import org.aopalliance.intercept.MethodInterceptor;

(2)

invocation는 호출된 클래스, 즉, PreparedStatement를 의미합니다.

앞서 말했듯이 PreparedStatement의 클래스 요청을 우리가 가로챈 것이지요.

PreparedStatement의 어떤 메서드를 호출했는지 알기 위해 invocation.getMethod를 한 것입니다.

(3)

PreparedStatement에서 DB에 쿼리를 날릴 때 executeQuery, execute, executeUpdate를 사용하기 때문에 해당 메서드를 호출하게 되면 해당 기능이 수행됩니다.

즉, PreparedStatement에서 JDBC_QUERY_METHOD에 있는 메서드를 호출하게 되면 트리거가 되어 if 절이 실행됩니다.

Connection Proxy Handler

package com.support;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import lombok.RequiredArgsConstructor;
import org.aopalliance.intercept.MethodInterceptor;
import org.aopalliance.intercept.MethodInvocation;
import org.springframework.aop.framework.ProxyFactory;

@RequiredArgsConstructor
public class ConnectionProxyHandler implements MethodInterceptor {

  private static final String JDBC_PREPARE_STATEMENT_METHOD_NAME = "prepareStatement";

  private final Object connection;
  private final LoggingForm loggingForm;

  @Nullable
  @Override
  public Object invoke(@Nonnull final MethodInvocation invocation) throws Throwable {
    final Object result = invocation.proceed(); // 1

    if (hasConnection(result) && hasPreparedStatementInvoked(invocation)) { // 2
      final ProxyFactory proxyFactory = new ProxyFactory(result);
      proxyFactory.addAdvice(new PreparedStatementProxyHandler(loggingForm));
      return proxyFactory.getProxy();
    }

    return result;
  }

  private boolean hasPreparedStatementInvoked(final MethodInvocation invocation) { // 3
    return invocation.getMethod().getName().equals(JDBC_PREPARE_STATEMENT_METHOD_NAME);
  }

  private boolean hasConnection(final Object result) {
    return result != null;
  }

  public Object getProxy() {
    final ProxyFactory proxyFactory = new ProxyFactory(connection);
    proxyFactory.addAdvice(this);
    return proxyFactory.getProxy();
  }
}

PreparedStatementProxyHandler와 거의 유사합니다.

(1)

Connection 클래스의 요청을 가로챕니다

(2)

Connection 객체가 있고, prepareStatement 메서드가 호출되면 if 문을 수행합니다.

여기서 중요한 점은 if 조건을 만족하게 되면 그냥 클래스가 아닌 이전에 만들었던 ProxyHandler를 주입해 줍니다.

즉, preparedStatement메서드를 호출하게 되면 PreparedStatementProxyHandler를 주입하여 해당 요청이 PreparedStatementProxyHandler로 가도록 하고,

메서드를 호출하지 않으면 일반 PreparedStatement를 주입해 주는 것입니다.

AOP 적용하기

package com.support;

import javax.servlet.http.HttpServletRequest;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.After;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

@Aspect
@Slf4j
public class NPlus1DetectorAop {

  private final ThreadLocal<LoggingForm> currentLoggingForm;

  public NPlus1DetectorAop() {
    this.currentLoggingForm = new ThreadLocal<>();
  }

  @Around("execution( * javax.sql.DataSource.getConnection())") // 1
  public Object captureConnection(final ProceedingJoinPoint joinPoint) throws Throwable {
    final Object connection = joinPoint.proceed();

    return new ConnectionProxyHandler(connection, getCurrentLoggingForm()).getProxy();
  }

  private LoggingForm getCurrentLoggingForm() {
    if (currentLoggingForm.get() == null) { // 2
      currentLoggingForm.set(new LoggingForm());
    }

    return currentLoggingForm.get();
  }

  @After("within(@org.springframework.web.bind.annotation.RestController *)") // 3
  public void loggingAfterApiFinish() {
    final LoggingForm loggingForm = getCurrentLoggingForm();

    ServletRequestAttributes attributes =
        (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();

    if (isInRequestScope(attributes)) {
      HttpServletRequest request = attributes.getRequest();

      loggingForm.setApiMethod(request.getMethod());
      loggingForm.setApiUrl(request.getRequestURI());
    }

    log.info("{}", getCurrentLoggingForm());
    currentLoggingForm.remove(); // 4
  }

  private boolean isInRequestScope(final ServletRequestAttributes attributes) {
    return attributes != null;
  }
}

이제 AOP 기능을 한번 만들어보겠습니다.

이전까지는 리플렉션을 통해서 프록시 클래스를 만든 과정이었다면

이제는 제일 처음 과정인 DataSource에서 Connection의 요청을 가로채고, API 가 끝나면 로깅하는 과정입니다.

(1)

@Around를 사용하여 getConnection이 호출되면 Connection을 호출하는 것이 아닌 Proxy 클래스의 Connection을 호출하도록 합니다.

(2)

ThreadLocal이 비어있으면 사물함을 하나 만들어주고, 있다면 기존에 있던 사물함을 반환해 줍니다.

(3)

RestController를 포인트컷에 사용한 이유는 API 요청에서의 쿼리를 세기 때문에 @RestController를 사용하였습니다.

즉, @RestController 이 있는 메서드를 사용하고, 메서드가 끝이 나면 uri, http method 및 해당 요청에서의 쿼리 시간, 쿼리 개수를 ThreadLocal의 loggingForm에 저장하고 로깅합니다.

(4)

ThreadLocal에서의 주의할 점은 사용하고 나서 무조건 remove를 해줘야 합니다. 메모리 누수가 발생하지 않기 위해,,

API를 호출하고 나서의 로깅된 값입니다. 정확히 9개의 쿼리가 나갔고 queryCounts가 9개로 정확하게 측정됐네요

그리고 AOP는 스프링 빈으로 등록돼있어야 합니다~~~

결론

이번 글을 통해서는 AOP, 프록시, ThreadLocal 등의 개념을 배울 수 있었습니다.

기존에 있던 Connection, PreparedStatement 객체가 아닌

프록시를 사용한 Connection, PreparedStatement를 통해서 쿼리 시간, 갯수, api url, http method를 알 수 있었습니다.

Spring에서는 프록시를 만들기 위한 방법도 여러 개가 있다는 것두요.

이 글을 통해 제 팀뿐만 아니라 우테코 전체 크루원들에게 도움 될 생각에 재밌게 기능을 만들고, 글을 작성한 것 같아요.

마지막 예시에서 행사를 저장하는 기능인데 쿼리가 9개 나가니까 바로 수정하러 가보겠습니다,,

여러분들도 적용해보시고 나서 치명적인 N+1을 발견해서 해결하셨으면 좋겠습니다

후기 글들,,,

감사합니다 디노 :)