Async Local Storage

AsyncLocalStorage là một Node.js API (dựa trên API async_hooks) cung cấp một cách khác để lan truyền state cục bộ (local state) xuyên suốt ứng dụng mà không cần truyền nó tường minh qua tham số hàm. Nó tương tự thread-local storage ở các ngôn ngữ khác.

Ý tưởng chính của Async Local Storage là chúng ta có thể bọc (wrap) một lời gọi hàm nào đó bằng lời gọi AsyncLocalStorage#run. Mọi code được gọi bên trong lời gọi đã bọc đó đều truy cập được cùng một store, và store này là duy nhất cho từng chuỗi lời gọi (call chain).

Trong ngữ cảnh của NestJS, điều đó nghĩa là nếu chúng ta tìm được một điểm trong vòng đời request mà tại đó ta có thể bọc phần code còn lại của request, chúng ta sẽ có thể truy cập và thay đổi state chỉ hiển thị với riêng request đó - đây có thể là một giải pháp thay thế cho provider có scope REQUEST cùng một số hạn chế của chúng.

Ngoài ra, chúng ta cũng có thể dùng ALS để lan truyền context cho chỉ một phần của hệ thống (ví dụ đối tượng transaction) mà không cần truyền nó tường minh qua các service, giúp tăng tính cô lập (isolation) và đóng gói (encapsulation).

Nếu ứng dụng của bạn đã được instrument bằng NestJS Observe, bạn đã có sẵn một store theo scope request và không cần tự xây dựng nó - xem bên dưới.

Triển khai tùy chỉnh

Bản thân NestJS không cung cấp bất kỳ abstraction có sẵn nào cho AsyncLocalStorage, vì vậy hãy cùng xem qua cách chúng ta có thể tự triển khai nó cho trường hợp HTTP đơn giản nhất, để hiểu rõ hơn toàn bộ khái niệm này:

1. Đầu tiên, tạo một instance mới của AsyncLocalStorage trong một file nguồn dùng chung. Vì chúng ta đang dùng NestJS, hãy biến nó thành một module với một custom provider.

@Module({
  providers: [
    {
      provide: AsyncLocalStorage,
      useValue: new AsyncLocalStorage(),
    },
  ],
  exports: [AsyncLocalStorage],
})
export class AlsModule {}

2. Chúng ta chỉ quan tâm đến HTTP, vì vậy hãy dùng một middleware để bọc hàm next bằng AsyncLocalStorage#run. Vì middleware là thứ đầu tiên mà request chạm vào, điều này sẽ khiến store khả dụng ở mọi enhancer và phần còn lại của hệ thống.

@Module({
  imports: [AlsModule],
  providers: [CatsService],
  controllers: [CatsController],
})
export class AppModule implements NestModule {
  constructor(
    // inject AsyncLocalStorage trong constructor của module,
    private readonly als: AsyncLocalStorage
  ) {}

  configure(consumer: MiddlewareConsumer) {
    // gắn (bind) middleware,
    consumer
      .apply((req, res, next) => {
        // khởi tạo store với một vài giá trị mặc định
        // dựa trên request,
        const store = {
          userId: req.headers['x-user-id'],
        };
        // và truyền hàm "next" như một callback
        // cho method "als.run" cùng với store.
        this.als.run(store, () => next());
      })
      .forRoutes('*path');
  }
}

@Module({
  imports: [AlsModule],
  providers: [CatsService],
  controllers: [CatsController],
})
@Dependencies(AsyncLocalStorage)
export class AppModule {
  constructor(als) {
    // inject AsyncLocalStorage trong constructor của module,
    this.als = als
  }

  configure(consumer) {
    // gắn (bind) middleware,
    consumer
      .apply((req, res, next) => {
        // khởi tạo store với một vài giá trị mặc định
        // dựa trên request,
        const store = {
          userId: req.headers['x-user-id'],
        };
        // và truyền hàm "next" như một callback
        // cho method "als.run" cùng với store.
        this.als.run(store, () => next());
      })
      .forRoutes('*path');
  }
}

3. Giờ đây, ở bất kỳ đâu trong vòng đời của một request, chúng ta có thể truy cập instance store cục bộ.

@Injectable()
export class CatsService {
  constructor(
    // Chúng ta có thể inject instance ALS đã được cung cấp.
    private readonly als: AsyncLocalStorage,
    private readonly catsRepository: CatsRepository,
  ) {}

  getCatForUser() {
    // Method "getStore" sẽ luôn trả về
    // instance store gắn với request tương ứng.
    const userId = this.als.getStore()["userId"] as number;
    return this.catsRepository.getForUser(userId);
  }
}

@Injectable()
@Dependencies(AsyncLocalStorage, CatsRepository)
export class CatsService {
  constructor(als, catsRepository) {
    // Chúng ta có thể inject instance ALS đã được cung cấp.
    this.als = als
    this.catsRepository = catsRepository
  }

  getCatForUser() {
    // Method "getStore" sẽ luôn trả về
    // instance store gắn với request tương ứng.
    const userId = this.als.getStore()["userId"] as number;
    return this.catsRepository.getForUser(userId);
  }
}

4. Vậy là xong. Giờ chúng ta đã có cách chia sẻ state liên quan đến request mà không cần inject toàn bộ đối tượng REQUEST.

NestJS Observe

Nếu bạn dùng NestJS Observe, SDK @nestjs/observe đã tự duy trì sẵn một store AsyncLocalStorage cho mọi request, job và message mà nó instrument - đây chính là thứ mang trace context xuyên suốt call stack của bạn. TracerService expose store đó, nên state theo scope request trở thành thứ bạn đọc/ghi thay vì phải tự thiết lập: không cần viết module, không cần gắn middleware, và không cần cấu hình lại cho từng transport.

Thiết lập

Không có gì cần thêm ngoài việc tích hợp SDK tiêu chuẩn đã mô tả trong Observability → SDK: cài package, import ObserveModule.forRoot(), và truyền ObserveInstrument vào NestFactory.create().

$ npm i @nestjs/observe

ObserveModule export TracerService, nên nó có thể được inject ở bất kỳ đâu trong ứng dụng của bạn:

import { Injectable } from '@nestjs/common';
import { TracerService } from '@nestjs/observe';

@Injectable()
export class CatsService {
  constructor(private readonly tracerService: TracerService) {}
}

Đọc và ghi store

setAttribute(key, value) ghi vào store của context hiện tại, và getAttribute(key) đọc lại giá trị đó từ bất kỳ đâu ở downstream trong cùng một request - một service khác, một guard, một interceptor - mà không cần truyền giá trị đó qua từng chữ ký hàm (function signature):

@Get()
findAll(@Req() req: Request) {
  this.tracerService.setAttribute('userId', req.headers['x-user-id']);
  return this.catsService.getCatForUser();
}

@Injectable()
export class CatsService {
  constructor(
    private readonly tracerService: TracerService,
    private readonly catsRepository: CatsRepository,
  ) {}

  getCatForUser() {
    const userId = this.tracerService.getAttribute('userId');
    return this.catsRepository.getForUser(userId);
  }
}

getAttribute() trả về undefined với một key chưa từng được set. Cả hai method đều ném lỗi (throw) khi được gọi ngoài một traced context, vì lúc đó không có store nào để đọc hay ghi - nếu bạn cần một giá trị có thể hợp lệ được đọc trước khi bất kỳ request nào tồn tại, hãy tự bảo vệ (guard) lời gọi đó cho phù hợp.

Định kiểu cho store

Truyền shape của store dưới dạng type argument đầu tiên của TracerService để kiểm tra key và value, kể cả các đường dẫn lồng nhau (nested path):

interface RequestStore {
  userId: number;
  flags: { betaCheckout: boolean };
}

@Injectable()
export class CatsService {
  constructor(private readonly tracerService: TracerService<RequestStore>) {}

  enableBeta() {
    this.tracerService.setAttribute('flags.betaCheckout', true);
  }
}

Vượt ra ngoài HTTP

Vì store được tạo bởi lớp instrumentation chứ không phải bởi middleware, nó tồn tại ở bất cứ đâu mà SDK trace một operation - request HTTP và GraphQL, message gRPC và @nestjs/microservices, và cả công việc chạy nền như consumer BullMQ và cron job. Pattern ở trên hoàn toàn giống nhau trong tất cả các trường hợp đó, đây chính là khác biệt thực tế so với một triển khai dựa trên middleware tự viết tay - vốn chỉ bao phủ được HTTP.

Store cũng lưu trace id hiện tại, nên nó đồng thời đóng vai trò là khóa tương quan (correlation key) cho log và các service downstream. currentTraceId() trả về giá trị đó, và khác với getAttribute(), nó trả về null khi ở ngoài traced context thay vì ném lỗi:

const traceId = this.tracerService.currentTraceId();

NestJS CLS

Package nestjs-cls cung cấp một số cải thiện về trải nghiệm dev (DX) so với việc dùng AsyncLocalStorage thuần (CLS là viết tắt của thuật ngữ continuation-local storage). Nó trừu tượng hóa phần triển khai thành một ClsModule, cung cấp nhiều cách khởi tạo store cho các transport khác nhau (không chỉ HTTP), cùng với hỗ trợ strong-typing.

Sau đó, store có thể được truy cập bằng một ClsService có thể inject, hoặc được trừu tượng hóa hoàn toàn khỏi business logic bằng cách dùng Proxy Providers.

Cài đặt

Ngoài một peer dependency với các thư viện @nestjs, nó chỉ dùng API có sẵn của Node.js. Cài nó như bất kỳ package nào khác.

npm i nestjs-cls

Sử dụng

Một chức năng tương tự như đã mô tả ở trên có thể được triển khai bằng nestjs-cls như sau:

1. Import ClsModule vào module gốc.

@Module({
  imports: [
    // Đăng ký ClsModule,
    ClsModule.forRoot({
      middleware: {
        // tự động gắn
        // ClsMiddleware cho tất cả các route
        mount: true,
        // và dùng method setup để
        // cung cấp giá trị mặc định cho store.
        setup: (cls, req) => {
          cls.set('userId', req.headers['x-user-id']);
        },
      },
    }),
  ],
  providers: [CatsService],
  controllers: [CatsController],
})
export class AppModule {}

2. Sau đó có thể dùng ClsService để truy cập các giá trị trong store.

@Injectable()
export class CatsService {
  constructor(
    // Chúng ta có thể inject instance ClsService đã được cung cấp,
    private readonly cls: ClsService,
    private readonly catsRepository: CatsRepository,
  ) {}

  getCatForUser() {
    // và dùng method "get" để lấy bất kỳ giá trị đã lưu nào.
    const userId = this.cls.get('userId');
    return this.catsRepository.getForUser(userId);
  }
}

@Injectable()
@Dependencies(AsyncLocalStorage, CatsRepository)
export class CatsService {
  constructor(cls, catsRepository) {
    // Chúng ta có thể inject instance ClsService đã được cung cấp,
    this.cls = cls
    this.catsRepository = catsRepository
  }

  getCatForUser() {
    // và dùng method "get" để lấy bất kỳ giá trị đã lưu nào.
    const userId = this.cls.get('userId');
    return this.catsRepository.getForUser(userId);
  }
}

3. Để có strong typing cho các giá trị store được quản lý bởi ClsService (và cũng để có auto-suggestion cho các key dạng string), chúng ta có thể dùng type parameter tùy chọn ClsService<MyClsStore> khi inject nó.

export interface MyClsStore extends ClsStore {
  userId: number;
}

Testing

Vì ClsService cũng chỉ là một provider có thể inject như bao provider khác, nó hoàn toàn có thể được mock trong unit test.

Tuy nhiên, trong một số integration test, chúng ta vẫn có thể muốn dùng implementation thật của ClsService. Trong trường hợp đó, chúng ta sẽ cần bọc đoạn code phụ thuộc context bằng một lời gọi đến ClsService#run hoặc ClsService#runWith.

describe('CatsService', () => {
  let service: CatsService
  let cls: ClsService
  const mockCatsRepository = createMock<CatsRepository>()

  beforeEach(async () => {
    const module = await Test.createTestingModule({
      // Thiết lập phần lớn testing module như thông thường.
      providers: [
        CatsService,
        {
          provide: CatsRepository
          useValue: mockCatsRepository
        }
      ],
      imports: [
        // Import phiên bản static của ClsModule, phiên bản này chỉ cung cấp
        // ClsService, nhưng không thiết lập store theo bất kỳ cách nào.
        ClsModule
      ],
    }).compile()

    service = module.get(CatsService)

    // Cũng lấy ClsService để dùng sau.
    cls = module.get(ClsService)
  })

  describe('getCatForUser', () => {
    it('retrieves cat based on user id', async () => {
      const expectedUserId = 42
      mocksCatsRepository.getForUser.mockImplementationOnce(
        (id) => ({ userId: id })
      )

      // Bọc lời gọi test trong method "runWith"
      // nơi chúng ta có thể truyền giá trị store tự tạo.
      const cat = await cls.runWith(
        { userId: expectedUserId },
        () => service.getCatForUser()
      )

      expect(cat.userId).toEqual(expectedUserId)
    })
  })
})

Thông tin thêm

Truy cập trang GitHub của NestJS CLS để xem tài liệu API đầy đủ và thêm ví dụ code.