feat(server): job system (#10134)

packages/backend/server/src/base/job/queue/__tests__/queue.spec.ts

@@ -0,0 +1,243 @@
+import { Injectable } from '@nestjs/common';
+import { TestingModule } from '@nestjs/testing';
+import test from 'ava';
+import { CLS_ID, ClsServiceManager } from 'nestjs-cls';
+import Sinon from 'sinon';
+
+import { createTestingModule } from '../../../../__tests__/utils';
+import { ConfigModule } from '../../../config';
+import { metrics } from '../../../metrics';
+import { genRequestId } from '../../../utils';
+import { JobModule, JobQueue, OnJob } from '..';
+import { JobExecutor } from '../executor';
+import { JobHandlerScanner } from '../scanner';
+
+let module: TestingModule;
+let queue: JobQueue;
+let executor: JobExecutor;
+
+declare global {
+  interface Jobs {
+    'nightly.__test__job': {
+      name: string;
+    };
+    'nightly.__test__job2': {
+      name: string;
+    };
+    'nightly.__test__throw': any;
+    'nightly.__test__requestId': any;
+  }
+}
+
+@Injectable()
+class JobHandlers {
+  @OnJob('nightly.__test__job')
+  @OnJob('nightly.__test__job2')
+  async handleJob(job: Jobs['nightly.__test__job']) {
+    return job.name;
+  }
+
+  @OnJob('nightly.__test__throw')
+  async throwJob() {
+    throw new Error('Throw in job handler');
+  }
+
+  @OnJob('nightly.__test__requestId')
+  onRequestId() {
+    const cls = ClsServiceManager.getClsService();
+    return cls.getId() ?? genRequestId('job');
+  }
+}
+
+test.before(async () => {
+  module = await createTestingModule({
+    imports: [
+      ConfigModule.forRoot({
+        job: {
+          worker: {
+            // NOTE(@forehalo):
+            //   bullmq will hold the connection to check stalled jobs,
+            //   which will keep the test process alive to timeout.
+            stalledInterval: 100,
+          },
+        },
+      }),
+      JobModule.forRoot(),
+    ],
+    providers: [JobHandlers],
+  });
+
+  queue = module.get(JobQueue);
+  executor = module.get(JobExecutor);
+});
+
+test.afterEach(async () => {
+  // @ts-expect-error private api
+  const inner = queue.getQueue('nightly');
+  await inner.obliterate({ force: true });
+  inner.resume();
+});
+
+test.after.always(async () => {
+  await module.close();
+});
+
+// #region scanner
+test('should register job handler', async t => {
+  const scanner = module.get(JobHandlerScanner);
+
+  const handler = scanner.getHandler('nightly.__test__job');
+
+  t.is(handler!.name, 'JobHandlers.handleJob');
+  t.is(typeof handler!.fn, 'function');
+
+  const result = await handler!.fn({ name: 'test' });
+
+  t.is(result, 'test');
+});
+// #endregion
+
+// #region queue
+test('should add job to queue', async t => {
+  const job = await queue.add('nightly.__test__job', { name: 'test' });
+
+  // @ts-expect-error private api
+  const innerQueue = queue.getQueue('nightly');
+  const queuedJob = await innerQueue.getJob(job.id!);
+
+  t.is(queuedJob.name, job.name);
+});
+
+test('should remove job from queue', async t => {
+  const job = await queue.add('nightly.__test__job', { name: 'test' });
+
+  // @ts-expect-error private api
+  const innerQueue = queue.getQueue('nightly');
+
+  const data = await queue.remove(job.id!, job.name as JobName);
+
+  t.deepEqual(data, { name: 'test' });
+
+  const nullData = await queue.remove(job.id!, job.name as JobName);
+  const nullJob = await innerQueue.getJob(job.id!);
+
+  t.is(nullData, undefined);
+  t.is(nullJob, undefined);
+});
+// #endregion
+
+// #region executor
+test('should start workers', async t => {
+  // @ts-expect-error private api
+  const worker = executor.workers['nightly'];
+
+  t.truthy(worker);
+  t.true(worker.isRunning());
+});
+
+test('should dispatch job handler', async t => {
+  const handlers = module.get(JobHandlers);
+  const spy = Sinon.spy(handlers, 'handleJob');
+
+  await executor.run('nightly.__test__job', { name: 'test executor' });
+
+  t.true(spy.calledOnceWithExactly({ name: 'test executor' }));
+});
+
+test('should be able to record job metrics', async t => {
+  const counterStub = Sinon.stub(metrics.job.counter('function_calls'), 'add');
+  const timerStub = Sinon.stub(
+    metrics.job.histogram('function_timer'),
+    'record'
+  );
+
+  await executor.run('nightly.__test__job', { name: 'test executor' });
+
+  t.deepEqual(counterStub.firstCall.args[1], {
+    name: 'job_handler',
+    job: 'nightly.__test__job',
+    namespace: 'nightly',
+    handler: 'JobHandlers.handleJob',
+    error: false,
+  });
+
+  t.deepEqual(timerStub.firstCall.args[1], {
+    name: 'job_handler',
+    job: 'nightly.__test__job',
+    namespace: 'nightly',
+    handler: 'JobHandlers.handleJob',
+    error: false,
+  });
+
+  counterStub.reset();
+  timerStub.reset();
+
+  await executor.run('nightly.__test__job2', { name: 'test executor' });
+
+  t.deepEqual(counterStub.firstCall.args[1], {
+    name: 'job_handler',
+    job: 'nightly.__test__job2',
+    namespace: 'nightly',
+    handler: 'JobHandlers.handleJob',
+    error: false,
+  });
+
+  t.deepEqual(timerStub.firstCall.args[1], {
+    name: 'job_handler',
+    job: 'nightly.__test__job2',
+    namespace: 'nightly',
+    handler: 'JobHandlers.handleJob',
+    error: false,
+  });
+
+  counterStub.reset();
+  timerStub.reset();
+
+  await t.throwsAsync(
+    executor.run('nightly.__test__throw', { name: 'test executor' }),
+    {
+      message: 'Throw in job handler',
+    }
+  );
+
+  t.deepEqual(counterStub.firstCall.args[1], {
+    name: 'job_handler',
+    job: 'nightly.__test__throw',
+    namespace: 'nightly',
+    handler: 'JobHandlers.throwJob',
+    error: true,
+  });
+
+  t.deepEqual(timerStub.firstCall.args[1], {
+    name: 'job_handler',
+    job: 'nightly.__test__throw',
+    namespace: 'nightly',
+    handler: 'JobHandlers.throwJob',
+    error: true,
+  });
+});
+
+test('should generate request id', async t => {
+  const handlers = module.get(JobHandlers);
+  const spy = Sinon.spy(handlers, 'onRequestId');
+
+  await executor.run('nightly.__test__requestId', {});
+
+  t.true(spy.returnValues.some(v => v.includes(':job/')));
+
+  spy.restore();
+});
+
+test('should continuously use request id', async t => {
+  const handlers = module.get(JobHandlers);
+  const spy = Sinon.spy(handlers, 'onRequestId');
+
+  const cls = ClsServiceManager.getClsService();
+  await cls.run(async () => {
+    cls.set(CLS_ID, 'test-request-id');
+    await executor.run('nightly.__test__requestId', {});
+  });
+  t.true(spy.returned('test-request-id'));
+  spy.restore();
+});
+// #endregion

packages/backend/server/src/base/job/queue/config.ts

@@ -0,0 +1,53 @@
+import { QueueOptions, WorkerOptions } from 'bullmq';
+
+import {
+  defineRuntimeConfig,
+  defineStartupConfig,
+  ModuleConfig,
+} from '../../config';
+import { Queue } from './def';
+
+declare module '../../config' {
+  interface AppConfig {
+    job: ModuleConfig<
+      {
+        queue: Omit<QueueOptions, 'connection'>;
+        worker: Omit<WorkerOptions, 'connection'>;
+      },
+      {
+        queues: {
+          [key in Queue]: {
+            concurrency: number;
+          };
+        };
+      }
+    >;
+  }
+}
+
+defineStartupConfig('job', {
+  queue: {
+    prefix: 'affine_job',
+    defaultJobOptions: {
+      attempts: 3,
+      removeOnComplete: true,
+      removeOnFail: false,
+    },
+  },
+  worker: {},
+});
+
+defineRuntimeConfig('job', {
+  'queues.nightly.concurrency': {
+    default: 1,
+    desc: 'Concurrency of worker consuming of nightly checking job queue',
+  },
+  'queues.notification.concurrency': {
+    default: 10,
+    desc: 'Concurrency of worker consuming of notification job queue',
+  },
+  'queues.doc.concurrency': {
+    default: 1,
+    desc: 'Concurrency of worker consuming of doc job queue',
+  },
+});

packages/backend/server/src/base/job/queue/def.ts

@@ -0,0 +1,64 @@
+import { join } from 'node:path';
+
+import { PushMetadata, sliceMetadata } from '../../nestjs';
+
+declare global {
+  /**
+   * Job definitions can be extended by
+   *
+   * @example
+   *
+   * declare global {
+   *   interface Jobs {
+   *     'nightly.deleteExpiredUserSessions': {}
+   *      ^^^^^^^ first segment must be namespace and a standalone queue will be created for each namespace
+   *   }
+   * }
+   */
+  interface Jobs {}
+
+  type JobName = keyof Jobs;
+}
+
+export const JOB_METADATA = Symbol('JOB');
+
+export enum Queue {
+  NIGHTLY_JOB = 'nightly',
+  NOTIFICATION = 'notification',
+  DOC = 'doc',
+}
+
+export const QUEUES = Object.values(Queue);
+
+export function namespace(job: JobName) {
+  const parts = job.split('.');
+
+  // no namespace
+  if (parts.length === 1) {
+    throw new Error(
+      `Job name must contain at least one namespace like [namespace].[job], get [${job}].`
+    );
+  }
+
+  return parts[0];
+}
+
+export const OnJob = (job: JobName) => {
+  const ns = namespace(job);
+  if (!QUEUES.includes(ns as Queue)) {
+    throw new Error(
+      `Invalid job queue: ${ns}, must be one of [${QUEUES.join(', ')}].
+If you want to introduce new job queue, please modify the Queue enum first in ${join(AFFiNE.projectRoot, 'src/base/job/queue/def.ts')}`
+    );
+  }
+
+  if (job === ns) {
+    throw new Error("The job name must not be the same as it's namespace.");
+  }
+
+  return PushMetadata(JOB_METADATA, job);
+};
+
+export function getJobHandlerMetadata(target: any): JobName[] {
+  return sliceMetadata<JobName>(JOB_METADATA, target);
+}

packages/backend/server/src/base/job/queue/executor.ts

@@ -0,0 +1,149 @@
+import {
+  Injectable,
+  Logger,
+  OnApplicationBootstrap,
+  OnApplicationShutdown,
+} from '@nestjs/common';
+import { Worker } from 'bullmq';
+import { difference } from 'lodash-es';
+import { CLS_ID, ClsServiceManager } from 'nestjs-cls';
+
+import { Config } from '../../config';
+import { metrics, wrapCallMetric } from '../../metrics';
+import { QueueRedis } from '../../redis';
+import { Runtime } from '../../runtime';
+import { genRequestId } from '../../utils';
+import { namespace, Queue, QUEUES } from './def';
+import { JobHandlerScanner } from './scanner';
+
+@Injectable()
+export class JobExecutor
+  implements OnApplicationBootstrap, OnApplicationShutdown
+{
+  private readonly logger = new Logger('job');
+  private readonly workers: Record<string, Worker> = {};
+
+  constructor(
+    private readonly config: Config,
+    private readonly redis: QueueRedis,
+    private readonly scanner: JobHandlerScanner,
+    private readonly runtime: Runtime
+  ) {}
+
+  async onApplicationBootstrap() {
+    const queues = this.config.flavor.graphql
+      ? difference(QUEUES, [Queue.DOC])
+      : [];
+
+    // NOTE(@forehalo): only enable doc queue in doc service
+    if (this.config.flavor.doc) {
+      queues.push(Queue.DOC);
+    }
+
+    await this.startWorkers(queues);
+  }
+
+  async onApplicationShutdown() {
+    await this.stopWorkers();
+  }
+
+  async run(name: JobName, payload: any) {
+    const ns = namespace(name);
+    const handler = this.scanner.getHandler(name);
+
+    if (!handler) {
+      this.logger.warn(`Job handler for [${name}] not found.`);
+      return;
+    }
+
+    const fn = wrapCallMetric(
+      async () => {
+        const cls = ClsServiceManager.getClsService();
+        await cls.run({ ifNested: 'reuse' }, async () => {
+          const requestId = cls.getId();
+          if (!requestId) {
+            cls.set(CLS_ID, genRequestId('job'));
+          }
+
+          const signature = `[${name}] (${handler.name})`;
+          try {
+            this.logger.debug(`Job started: ${signature}`);
+            const result = await handler.fn(payload);
+            this.logger.debug(`Job finished: ${signature}`);
+            return result;
+          } catch (e) {
+            this.logger.error(`Job failed: ${signature}`, e);
+            throw e;
+          }
+        });
+      },
+      'job',
+      'job_handler',
+      {
+        job: name,
+        namespace: ns,
+        handler: handler.name,
+      }
+    );
+    const activeJobs = metrics.job.gauge('queue_active_jobs');
+    activeJobs.record(1, { queue: ns });
+    try {
+      return await fn();
+    } finally {
+      activeJobs.record(-1, { queue: ns });
+    }
+  }
+
+  private async startWorkers(queues: Queue[]) {
+    const configs =
+      (await this.runtime.fetchAll(
+        queues.reduce(
+          (ret, queue) => {
+            ret[`job/queues.${queue}.concurrency`] = true;
+            return ret;
+          },
+          {} as {
+            [key in `job/queues.${Queue}.concurrency`]: true;
+          }
+        )
+        // TODO(@forehalo): fix the override by [payment/service.spec.ts]
+      )) ?? {};
+
+    for (const queue of queues) {
+      const concurrency =
+        (configs[`job/queues.${queue}.concurrency`] as number) ??
+        this.config.job.worker.concurrency ??
+        1;
+
+      const worker = new Worker(
+        queue,
+        async job => {
+          await this.run(job.name as JobName, job.data);
+        },
+        {
+          ...this.config.job.worker,
+          connection: this.redis,
+          concurrency,
+        }
+      );
+
+      worker.on('error', error => {
+        this.logger.error(`Queue Worker [${queue}] error`, error);
+      });
+
+      this.logger.log(
+        `Queue Worker [${queue}] started; concurrency=${concurrency};`
+      );
+
+      this.workers[queue] = worker;
+    }
+  }
+
+  private async stopWorkers() {
+    await Promise.all(
+      Object.values(this.workers).map(async worker => {
+        await worker.close(true);
+      })
+    );
+  }
+}

packages/backend/server/src/base/job/queue/index.ts

@@ -0,0 +1,37 @@
+import './config';
+
+import { BullModule } from '@nestjs/bullmq';
+import { DynamicModule } from '@nestjs/common';
+
+import { Config } from '../../config';
+import { QueueRedis } from '../../redis';
+import { QUEUES } from './def';
+import { JobExecutor } from './executor';
+import { JobQueue } from './queue';
+import { JobHandlerScanner } from './scanner';
+
+export class JobModule {
+  static forRoot(): DynamicModule {
+    return {
+      global: true,
+      module: JobModule,
+      imports: [
+        BullModule.forRootAsync({
+          useFactory: (config: Config, redis: QueueRedis) => {
+            return {
+              ...config.job.queue,
+              connection: redis,
+            };
+          },
+          inject: [Config, QueueRedis],
+        }),
+        BullModule.registerQueue(...QUEUES.map(name => ({ name }))),
+      ],
+      providers: [JobQueue, JobExecutor, JobHandlerScanner],
+      exports: [JobQueue],
+    };
+  }
+}
+
+export { JobQueue };
+export { OnJob } from './def';

packages/backend/server/src/base/job/queue/queue.ts

@@ -0,0 +1,43 @@
+import { getQueueToken } from '@nestjs/bullmq';
+import { Injectable, Logger } from '@nestjs/common';
+import { ModuleRef } from '@nestjs/core';
+import { Job, JobsOptions, Queue } from 'bullmq';
+
+import { namespace } from './def';
+
+@Injectable()
+export class JobQueue {
+  private readonly logger = new Logger(JobQueue.name);
+
+  constructor(private readonly moduleRef: ModuleRef) {}
+
+  async add<T extends JobName>(name: T, payload: Jobs[T], opts?: JobsOptions) {
+    const ns = namespace(name);
+    const queue = this.getQueue(ns);
+    const job = await queue.add(name, payload, opts);
+    this.logger.debug(`Job [${name}] added; id=${job.id}`);
+    return job;
+  }
+
+  async remove<T extends JobName>(jobId: string, jobName: T) {
+    const ns = namespace(jobName);
+    const queue = this.getQueue(ns);
+    const job = (await queue.getJob(jobId)) as Job<Jobs[T]> | undefined;
+
+    if (!job) {
+      return;
+    }
+
+    const removed = await queue.remove(jobId);
+    if (removed) {
+      this.logger.log(`Job ${jobName} removed from queue ${ns}`);
+      return job.data;
+    }
+
+    return undefined;
+  }
+
+  private getQueue(ns: string): Queue {
+    return this.moduleRef.get(getQueueToken(ns), { strict: false });
+  }
+}

packages/backend/server/src/base/job/queue/scanner.ts

@@ -0,0 +1,77 @@
+import { Injectable, OnModuleInit } from '@nestjs/common';
+
+import { ModuleScanner } from '../../nestjs';
+import { getJobHandlerMetadata } from './def';
+
+interface JobHandler {
+  name: string;
+  fn: (payload: any) => any;
+}
+
+@Injectable()
+export class JobHandlerScanner implements OnModuleInit {
+  private readonly handlers: Record<string, JobHandler> = {};
+
+  constructor(private readonly scanner: ModuleScanner) {}
+
+  async onModuleInit() {
+    this.scan();
+  }
+
+  getHandler(jobName: JobName): JobHandler | undefined {
+    return this.handlers[jobName];
+  }
+
+  private scan() {
+    const providers = this.scanner.getAtInjectables();
+
+    providers.forEach(wrapper => {
+      const { instance, name } = wrapper;
+      if (!instance || wrapper.isAlias) {
+        return;
+      }
+
+      const methods = this.scanner.getAllMethodNames(instance);
+
+      methods.forEach(method => {
+        const fn = instance[method];
+
+        let jobNames = getJobHandlerMetadata(instance[method]);
+
+        if (jobNames.length === 0) {
+          return;
+        }
+
+        const signature = `${name}.${method}`;
+
+        if (typeof fn !== 'function') {
+          throw new Error(`Job handler [${signature}] is not a function.`);
+        }
+
+        if (!wrapper.isDependencyTreeStatic()) {
+          throw new Error(
+            `Provider [${name}] could not be RequestScoped or TransientScoped injectable if it contains job handlers.`
+          );
+        }
+
+        jobNames.forEach(jobName => {
+          if (this.handlers[jobName]) {
+            throw new Error(
+              `Job handler ${jobName} already defined in [${this.handlers[jobName].name}].`
+            );
+          }
+
+          this.handlers[jobName] = {
+            name: signature,
+            fn: (payload: any) => {
+              // NOTE(@forehalo):
+              //   we might create spies on the job handlers when testing,
+              //   avoid reusing `fn` variable to fail the spies or stubs
+              return instance[method].bind(instance)(payload);
+            },
+          };
+        });
+      });
+    });
+  }
+}