AsyncLocalStorage를 이용해 Transaction 관심사 분리하기 (lines 40% 감소)

진행하는 프로젝트 특성상 DB 작업이 복잡하게 일어나지 않습니다. ORM에서 제공하는 트랜잭션, 쿼리문으로 충분한 수준입니다.

하지만 만약 DB 작업이 복잡해진다면 어떻게 해야할까요? 쿼리나 트랜잭션을 직접 관리해주는 것이 좋을 것이라 생각했습니다.

그래서 이번 프로젝트에서도 트랜잭션을 적용해보기로 했습니다.
프로젝트 Github

문제점

가장 기본적으로 트랜잭션은 이렇게 구현할 수 있습니다.

  const qr = this.getQueryRunner();
  try {
    await qr.connect();
    await qr.startTransaction();
    await qr.manager.save<UserEntity>(user); // 실제 쓰기가 일어나는 곳
    await qr.commitTransaction();
  } catch (e) {
    console.log(e);
    await qr.rollbackTransaction();
    throw new TransactionRollback();
  } finally {
    await qr.release();
  }

실제 쓰기가 일어나는 곳 하나를 작성하기위해 무수히 많은 코드가 필요합니다.

프로젝트에는 쿼리가 하나만 있는 것이 아니기 때문에 매번 데이터베이스 쓰기작업을 하는 함수마다 이 모든 걸 작성해주는 것은 귀찮고, 비효율적이고, 휴먼에러 발생확률이 매우 높습니다.

그래서 관심사 분리를 통해 공통 로직들을 분리합니다.

AOP

Aspect Oriented Programming, 관점 지향 프로그래밍

출처 https://tecoble.techcourse.co.kr/post/2021-06-25-aop-transaction/

위의 그림으로 가장 잘 AOP가 설명된다고 생각합니다.

Class A가 회원가입을 담당하고 있으면 A에서는 회원가입만하고 로깅, 보안, 트랜잭션은 다른 곳에서 하게 하겠다는 것이죠.

핵심 비즈니스 로직과 공통적으로 필요한 코드를 분리해서 단일책임원칙을 지킬 수 있게 됩니다.

Intercepter?

NestJs에는 인터셉터라는 개념이 존재합니다.

라우트핸들러가 있을때 인터셉터를 적용하면 라우트핸들러 실행 전, 후로 실행되어 생명주기를 관리할 수 있습니다.

트랜잭션 인터셉터는 다음과 같이 구현할 수 있습니다.

return 문 전까지는 핸들러 실행 전에 실행이 되고, next.handle()로 핸들러가 실행되고 pipe를 통해 그 이후의 로직들을 관리할 수 있습니다.

@Injectable()
export class TransactionInterceptor implements NestInterceptor {
  constructor(private readonly dataSource: DataSource) {}

  async intercept(
    context: ExecutionContext,
    next: CallHandler,
  ): Promise<Observable<any>> {
    const request = context.switchToHttp().getRequest();

    const qr = this.dataSource.createQueryRunner();

    await qr.connect();

    await qr.startTransaction();

    request.queryRunner = qr;

    return next.handle().pipe(
      catchError(async (e) => {
        await qr.rollbackTransaction();
        await qr.release();
        throw new TransactionRollback();
      }),
      tap(async () => {
        await qr.commitTransaction();
        await qr.release();
      }),
    );
  }
}

꽤나 괜찮아보이는지만 한 가지 문제점이 있습니다.

인터셉터는 컨트롤러 레벨에 적용된다는 것입니다.

실제로 트랜잭션이 적용되어야할 컨트롤러 함수에는 트랜잭션 하나로 관리되어야할 코드만 들어간다는 보장이 없습니다.

한 라우트핸들러에 복수 개의 트랜잭션이 필요할 수도 있고, 읽기 쓰기 작업이 혼용되기도 합니다.

저는 쓰기 작업에만 트랜잭션을 적용시키고 싶고요.

@Get('github-callback')
  async githubCallback(
    @Req() req: Request,
    @Res() res: Response,
    @Query('code') code: string,
  ) {
    const accessToken = await this.githubService.getGithubAccessToken(code);
    const user: UserDto = await this.githubService.getUserInfo(accessToken);
    let findUser = await this.userService.findUser(user);
    if (findUser === null) {
      await this.userService.addUser(user, 'github');
      findUser = await this.userService.findUser(user);
    }
    const returnTo: string = await this.authService.login(findUser, res, req);

    return res.redirect(returnTo);
  }

그리고 인터셉터로 구현하게 되면 함수의 인자로 쿼리러너를 전달해주고, 서비스 레벨에서 쿼리러너를 이용하려면 서비스레벨에도 쿼리러너를 전달해주어야하는 불편함이 생깁니다.

Decorator & AsyncLocalStorage

위의 불편함들 때문에 서비스 레벨에 적용시킬 수 있는 데코레이터를 만들어 트랜잭션을 관리하고 싶어졌습니다.

이것을 구현하기까지 두 가지 어려움이 있었습니다.

1. NestJs 라이프 사이클에 맞는 데코레이터 만들기
2. 데코레이터에서 생성한 쿼리러너를 서비스함수에서도 이용하기

NestJs 라이프 사이클에 맞는 데코레이터 만들기

NestJS는 싱글톤으로 인스턴스들을 관리해 앱이 초기화될 때 모든 클래스들이 초기화됩니다.

express에서 MethodDecorator를 구현할 때와 약간 동작에 차이점이 있었습니다.

데코레이터로 메타데이터 정보를 입력하고,

NestJS 라이프 사이클에 따라 클래스들이 초기화 되고 그 후 DiscoverService로 등록된 인스턴스들을 탐색하고 필터링하여 트랜잭션 적용이 필요한 곳을 찾아 원래함수를 wrapping해주어 구현하였습니다.

// transaction.module.ts
import { Module } from '@nestjs/common';
import { DiscoveryModule } from '@nestjs/core';
import { TransactionService } from './transaction.service';

@Module({
  imports: [DiscoveryModule],
  providers: [TransactionService],
})
export class TransactionModule {}

// transaction.decorator.ts
export const TRANSACTIONAL_KEY = Symbol('TRANSACTION');
export type ORM = 'typeorm' | 'mongoose';
export function Transactional(orm: ORM): MethodDecorator {
  return applyDecorators(SetMetadata(TRANSACTIONAL_KEY, orm));
}

프로젝트에서 typeorm과 mongoose를 사용해서 데코레이터 옵션을 통해 orm 타입을 받아주었습니다.

// transaction.service.ts
@Injectable()
export class TransactionService implements OnModuleInit {
  constructor(
    private readonly discoveryService: DiscoveryService,
    private readonly metadataScanner: MetadataScanner,
    private readonly reflector: Reflector,
    private readonly dataSource: DataSource,
    @InjectConnection() private readonly connection: Connection,
  ) {}

  onModuleInit(): any {
    const providers = this.discoveryService.getProviders();

    const instances = providers
      .filter((v) => v.isDependencyTreeStatic())
      .filter(({ metatype, instance }) => {
        return !(!metatype || !instance);
      })
      .map(({ instance }) => instance);

    instances.map((instance) => {
      const names = this.metadataScanner.getAllMethodNames(
        Object.getPrototypeOf(instance),
      );
      for (const name of names) {
        const originalMethod = instance[name];
        const metadata = this.reflector.get<ORM>(
          TRANSACTIONAL_KEY,
          originalMethod,
        );
        switch (metadata) {
          case 'typeorm':
            instance[name] = this.typeormTransaction(originalMethod, instance);
            return;
          case 'mongoose':
            instance[name] = this.mongooseTransaction(originalMethod, instance);
        }
      }
    });
  }
... (생략)
}

데코레이터에서 생성한 쿼리러너를 서비스함수에서도 이용하기

데코레이터로 함수를 조작해줄 때 인자값을 더해주지 않고 동일한 쿼리러너를 서비스함수에서도 이용할 수 있어야합니다.

자바에서는 이를 ThreadLocal로 구현할 수 있습니다. 스레드 영역에 변수를 할당해 특정 스레드에서 실행되는 모든 코드에서 그 변수에 접근이 가능합니다.

하지만 자바스크립트는 싱글스레드인데요?

필요한 모든 건 항상 만들어져있습니다…

AsyncLocalStorage라는 것을 통해 비동기작업이 일어나는 동안에 접근가능한 storage를 만들 수 있습니다.

데코레이터에서 여기에 쿼리러너를 저장하고, 서비스 함수에서 이 storage영역에서 쿼리러너를 읽어와 사용할 수 있습니다.

// transaction.service.ts
typeormTransaction(originalMethod, instance) {
    const dataSource = this.dataSource;
    return async function (...args: any[]) {
      const qr = await dataSource.createQueryRunner();

      await queryRunnerLocalStorage.run({ qr }, async function () {
        try {
          await qr.startTransaction();
          const result = await originalMethod.apply(instance, args);
          await qr.commitTransaction();
          return result;
        } catch (e) {
          await qr.rollbackTransaction();
          this.logger.error(e);
          throw new TransactionRollback();
        } finally {
          await qr.release();
        }
      });
    };
  }

  mongooseTransaction(originalMethod, instance) {
    const connection = this.connection;
    return async function (...args: any[]) {
      const session: ClientSession = await connection.startSession();

      const result = await sessionLocalStorage.run(
        { session },
        async function () {
          try {
            await session.startTransaction();
            const result = await originalMethod.apply(instance, args);
            await session.commitTransaction();
            return result;
          } catch (e) {
            await session.abortTransaction();
            throw new TransactionRollback();
          } finally {
            await session.endSession();
          }
        },
      );
      return result;
    };
  }

// transaction.decorator.ts
export function getLocalStorageRepository<T extends ObjectLiteral>(
  target,
): Repository<T> {
  const queryRunner = queryRunnerLocalStorage.getStore();
  return queryRunner?.qr?.manager.getRepository(target);
}
export function getSession(): ClientSession {
  const session = sessionLocalStorage.getStore();
  return session?.session;
}

Service에 적용

이제 서비스 클래스에서 트랜잭션을 적용하고자 하는 함수에 위의 데코레이터만 작성해주면 트랜잭션이 적용됩니다.

@Transactional('typeorm')
@Transactional('mongoose')

쓰기 작업에만 트랜잭션이 적용되는 모습

코드 길이 비교

트랜잭션 관심사분리를 통해서 코드를 많이 줄일 수 있었습니다.

TypeORM

Before

async addUser(userDTO: UserDto, oauth: OAUTH) {
  const user = new UserEntity();
  user.name = userDTO.name;
  user.authServiceID = userDTO.authServiceID;
  user.oauth = oauth;

  const qr = this.getQueryRunner();
  try {
    await qr.connect();
    await qr.startTransaction();
    await qr.manager.save<UserEntity>(user);
    await qr.commitTransaction();
  } catch (e) {
    console.log(e);
    await qr.rollbackTransaction();
    throw new TransactionRollback();
  } finally {
    await qr.release();
  }
}

After

@Transactional('typeorm')
  async addUser(userDTO: UserDto, oauth: OAUTH) {
    const user = new UserEntity();
    user.name = userDTO.name;
    user.authServiceID = userDTO.authServiceID;
    user.oauth = oauth;
    const repository = getLocalStorageRepository(UserEntity);
    await repository.save<UserEntity>(user);
  }

Mongoose

특히 쓰기 작업이 필요한 곳이 더 많을수록 더 빛을 발했습니다.
지금보다 DB 작업이 많은 프로젝트라면 더 큰 효과를 보일 것입니다.

Before

async save(saveCodeDto: SaveCodeDto): Promise<Code> {
    const session = await this.connection.startSession();
    session.startTransaction();
    try {
      const code = await this.codeModel.create(saveCodeDto);
      await session.commitTransaction();
      return code;
    } catch (e) {
      await session.abortTransaction();
      this.logger.error(e);
      throw new TransactionRollback();
    } finally {
      await session.endSession();
    }
  }

  async update(userID: number, objectID: string, saveCodeDto: SaveCodeDto) {
    const query = { userID: userID, _id: objectID };
    const session = await this.connection.startSession();
    session.startTransaction();
    try {
      const result = await this.codeModel.updateOne(query, saveCodeDto);
      await session.commitTransaction();
      return result;
    } catch (e) {
      await session.abortTransaction();
      this.logger.error(e);
      throw new TransactionRollback();
    } finally {
      await session.endSession();
    }
    return;
  }

  async delete(userID: number, objectID: string) {
    const query = { userID: userID, _id: objectID };
    const session = await this.connection.startSession();
    session.startTransaction();
    try {
      await this.codeModel.deleteOne(query);
      await session.commitTransaction();
    } catch (e) {
      await session.abortTransaction();
      this.logger.error(e);
      throw new TransactionRollback();
    } finally {
      await session.endSession();
    }
  }

After

@Transactional('mongoose')
  async save(saveCodeDto: SaveCodeDto) {
    const session = getSession();
    const code = await this.codeModel.create([saveCodeDto], {
      session: session,
    });
    return code[0];
  }

  @Transactional('mongoose')
  async update(userID: number, objectID: string, saveCodeDto: SaveCodeDto) {
    const query = { userID: userID, _id: objectID };
    const session: ClientSession = getSession();
    const result = await this.codeModel
      .updateOne(query, saveCodeDto)
      .session(session);
    return result;
  }

  @Transactional('mongoose')
  async delete(userID: number, objectID: string): Promise<DeleteResult> {
    const query = { userID: userID, _id: objectID };
    const session = getSession();
    return this.codeModel.deleteOne(query).session(session);
  }

codes.service.ts에서 import문을 제외했을때 67 lines이던 코드가 40 lines로 줄어들었습니다. (40% 감소)

---

참고

https://tecoble.techcourse.co.kr/post/2021-06-25-aop-transaction/

https://toss.tech/article/nestjs-custom-decorator

https://www.freecodecamp.org/news/async-local-storage-nodejs/

https://nodejs.org/api/async_context.html#class-asynclocalstorage