tests/worker.integration.test.ts

import { afterAll, beforeAll, describe, expect, it } from 'vitest';
import { randomUUID } from 'node:crypto';
import { Queue, QueueEvents, type WorkerOptions } from 'bullmq';
import IORedis from 'ioredis';
import { PrismaPg } from '@prisma/adapter-pg';
import { PrismaClient } from '@/generated/prisma/client';
import { createTaskWorker } from '@/server/jobs/worker';
import { createWorkerLogger } from '@/server/jobs/logger';
import { dispatchTasks } from '@/server/jobs/dispatcher';
import {
  queueName,
  taskJobId,
  type TaskJob,
  type createTaskQueue,
} from '@/server/jobs/queue';

function gate() {
  let open!: () => void;
  const promise = new Promise<void>((resolve) => {
    open = resolve;
  });
  return { promise, open };
}
async function within<T>(promise: Promise<T>, ms = 8000): Promise<T> {
  let timer!: ReturnType<typeof setTimeout>;
  try {
    return await Promise.race([
      promise,
      new Promise<never>((_, reject) => {
        timer = setTimeout(
          () => reject(new Error('Worker 集成同步点超时')),
          ms,
        );
      }),
    ]);
  } finally {
    clearTimeout(timer);
  }
}
function logs() {
  const records: { event: string; exhausted?: boolean }[] = [];
  const signals = new Map<string, ReturnType<typeof gate>>();
  const logger = createWorkerLogger({
    write(line) {
      const record = JSON.parse(String(line));
      records.push(record);
      signals.get(record.event)?.open();
    },
  });
  logger.level = 'trace';
  return {
    logger,
    records,
    event(name: string) {
      if (records.some((record) => record.event === name))
        return Promise.resolve();
      if (!signals.has(name)) signals.set(name, gate());
      return within(signals.get(name)!.promise);
    },
  };
}

// 两个真实 Worker 共用生产工厂、处理器与事件监听器；每例只清理自己的队列前缀。
describe.skipIf(!process.env.TEST_DATABASE_URL || !process.env.TEST_REDIS_URL)(
  '真实 Worker 的 stalled、失败回调与数据库租约',
  () => {
    const db = new PrismaClient({
      adapter: new PrismaPg({
        connectionString: process.env.TEST_DATABASE_URL ?? '',
        max: 8,
      }),
    });
    const userId = randomUUID();
    let teamId: string;
    beforeAll(async () => {
      await db.user.create({
        data: {
          id: userId,
          name: 'Worker 测试管理员',
          email: `${userId}@example.com`,
        },
      });
      teamId = (
        await db.team.create({
          data: {
            name: 'Worker 隔离团队',
            members: { create: { userId, role: 'ADMIN' } },
          },
        })
      ).id;
    });
    afterAll(async () => {
      await db.ticket.deleteMany({ where: { teamId } });
      await db.team.delete({ where: { id: teamId } });
      await db.user.delete({ where: { id: userId } });
      await db.$disconnect();
    });
    const fixture = async (kind: 'EXPORT' | 'REMINDER' = 'EXPORT') => {
      const item = await db.ticket.create({
        data: {
          teamId,
          creatorId: userId,
          title: `Worker-${randomUUID()}`,
          description: '真实消费者接手',
          dueAt: new Date(0),
        },
      });
      const task = await db.backgroundTask.create({
        data: {
          teamId,
          ownerId: userId,
          kind,
          ...(kind === 'EXPORT'
            ? { filters: { q: item.title } }
            : { ticketId: item.id, reminderRevision: item.reminderRevision }),
        },
      });
      return { item, task };
    };
    async function harness() {
      const prefix = `worker-test-${randomUUID()}`;
      const connection = new IORedis(process.env.TEST_REDIS_URL!, {
        maxRetriesPerRequest: null,
      });
      const queue = new Queue<TaskJob>(queueName, {
        connection,
        prefix,
        defaultJobOptions: { attempts: 1 },
      });
      const events = new QueueEvents(queueName, { connection, prefix });
      const workers: ReturnType<typeof createTaskWorker>[] = [];
      const releases: (() => void)[] = [];
      await events.waitUntilReady();
      return {
        queue,
        events,
        releases,
        async worker(
          client: PrismaClient,
          options: Omit<WorkerOptions, 'connection'> = {},
        ) {
          const capture = logs();
          const worker = createTaskWorker(
            client,
            connection,
            {
              prefix,
              concurrency: 1,
              lockDuration: 10000,
              stalledInterval: 100,
              ...options,
            },
            capture.logger,
          );
          workers.push(worker);
          await worker.waitUntilReady();
          return { worker, ...capture };
        },
        async dispatch() {
          // 测试与其他套件共享数据库，只限制夹具团队；不改变生产恢复条件。
          const scoped = db.$extends({
            query: {
              backgroundTask: {
                findMany({ args, query }) {
                  args.where = { AND: [args.where ?? {}, { teamId }] };
                  return query(args);
                },
                updateMany({ args, query }) {
                  args.where = { AND: [args.where ?? {}, { teamId }] };
                  return query(args);
                },
              },
            },
          });
          await dispatchTasks(
            scoped as unknown as PrismaClient,
            queue as ReturnType<typeof createTaskQueue>,
          );
        },
        async close() {
          releases.forEach((release) => release());
          await Promise.all(workers.map((worker) => worker.close()));
          await events.close();
          await queue.obliterate({ force: true });
          await queue.close();
          connection.disconnect();
        },
      };
    }

    it.each(['EXPORT', 'REMINDER'] as const)(
      '真实 stalled 后 %s 的新执行者完成，旧执行者晚到错误不改变 E2',
      async (kind) => {
        const h = await harness();
        const { task, item } = await fixture(kind);
        const pausedA = gate(),
          resumeA = gate(),
          claimedB = gate(),
          resumeB = gate();
        h.releases.push(resumeA.open, resumeB.open);
        const oldDb = db.$extends({
          query: {
            ticket: {
              async findMany({ args, query }) {
                const result = await query(args);
                pausedA.open();
                await resumeA.promise;
                throw new Error(
                  `测试注入：旧执行者读取失败，已读取 ${result.length} 行`,
                );
              },
            },
            backgroundTask: {
              async findFirst({ args, query }) {
                await query(args);
                pausedA.open();
                await resumeA.promise;
                throw new Error('测试注入：旧提醒执行者读取失败');
              },
            },
          },
        });
        const newDb = db.$extends({
          query: {
            backgroundTask: {
              async updateMany({ args, query }) {
                const result = await query(args);
                if (args.data.status === 'RUNNING') {
                  claimedB.open();
                  await resumeB.promise;
                }
                return result;
              },
            },
          },
        });
        try {
          const a = await h.worker(oldDb as unknown as PrismaClient, {
            skipLockRenewal: true,
            skipStalledCheck: true,
            lockDuration: 1000,
          });
          const job = await h.queue.add(
            'execute',
            { taskId: task.id, run: task.run },
            { jobId: taskJobId(task.id, task.run) },
          );
          await within(pausedA.promise);
          const first = await db.backgroundTask.findUniqueOrThrow({
            where: { id: task.id },
          });
          const b = await h.worker(newDb as unknown as PrismaClient);
          // 不删锁、不手工 emit stalled，等待 Redis 锁自然过期与真实检测器恢复。
          await b.event('task.stalled');
          await within(claimedB.promise);
          const replacement = await db.backgroundTask.findUniqueOrThrow({
            where: { id: task.id },
          });
          expect(replacement.executionId).not.toBe(first.executionId);
          expect(replacement.executionId).not.toBeNull();
          expect(replacement).toMatchObject({
            status: 'RUNNING',
            run: task.run,
          });
          resumeA.open();
          await a.event('task.skipped');
          expect(
            await db.backgroundTask.findUnique({ where: { id: task.id } }),
          ).toEqual(replacement);
          expect(
            a.records.some((record) => record.event === 'task.attempt_failed'),
          ).toBe(false);
          const finished = job.waitUntilFinished(h.events, 8000);
          resumeB.open();
          await finished;
          expect(await job.getState()).toBe('completed');
          expect((await h.queue.getJob(job.id!))?.progress).toBe(100);
          const completed = await db.backgroundTask.findUniqueOrThrow({
            where: { id: task.id },
          });
          expect(completed).toMatchObject({
            status: 'COMPLETED',
            progress: 100,
            executionId: null,
            lastError: null,
            run: task.run,
          });
          if (kind === 'EXPORT') {
            expect(completed.resultSummary).toBe('已导出 1 条工单');
            expect(completed.resultCsv).toContain(item.title);
            expect(completed.expiresAt!.getTime()).toBeGreaterThan(Date.now());
          } else {
            expect(completed.resultSummary).toBe('已生成 1 条站内提醒');
            expect(
              await db.reminderNotification.count({
                where: { ticketId: item.id },
              }),
            ).toBe(1);
            const records = await db.ticketEvent.findMany({
              where: { ticketId: item.id, kind: 'TIMEOUT_REMINDED' },
            });
            expect(records).toHaveLength(1);
            expect(records[0].version).toBe(item.version + 1);
            expect(
              await db.ticket.findUnique({ where: { id: item.id } }),
            ).toMatchObject({ version: item.version + 1 });
          }
        } finally {
          await h.close();
        }
      },
    );

    it('真实 failed 回调迟到，恢复投递后的同 run 新租约不能被清空', async () => {
      const h = await harness();
      const { task, item } = await fixture();
      const releaseCallback = gate(),
        claimedB = gate(),
        resumeB = gate();
      h.releases.push(releaseCallback.open, resumeB.open);
      const lateWrites: Promise<unknown>[] = [];
      const oldDb = db.$extends({
        query: {
          ticket: {
            async findMany() {
              throw new Error('测试注入：导出读取失败');
            },
          },
          backgroundTask: {
            async updateMany({ args, query }) {
              if (args.data.status === 'WAITING')
                throw new Error('测试注入：失败状态暂时无法写入');
              // 对旧回调的 SQL 模拟迟到。修复后回调不再发 SQL，因此列表为空。
              if (args.data.lastError === '任务多次执行失败，请手动重试') {
                const pending = releaseCallback.promise.then(() => query(args));
                lateWrites.push(pending);
                return pending;
              }
              return query(args);
            },
          },
        },
      });
      const newDb = db.$extends({
        query: {
          backgroundTask: {
            async updateMany({ args, query }) {
              const result = await query(args);
              if (args.data.status === 'RUNNING') {
                claimedB.open();
                await resumeB.promise;
              }
              return result;
            },
          },
        },
      });
      try {
        const a = await h.worker(oldDb as unknown as PrismaClient);
        const job = await h.queue.add(
          'execute',
          { taskId: task.id, run: task.run },
          { jobId: taskJobId(task.id, task.run) },
        );
        await a.event('task.failed');
        await a.worker.pause();
        expect(await job.getState()).toBe('failed');
        expect(
          a.records.find((record) => record.event === 'task.failed'),
        ).toMatchObject({ exhausted: true });
        const first = await db.backgroundTask.findUniqueOrThrow({
          where: { id: task.id },
        });
        expect(first).toMatchObject({ status: 'RUNNING' });
        // 只移除本例单个失败 job，模拟队列记录丢失；不清开发 Redis。
        await job.remove();
        await db.backgroundTask.update({
          where: { id: task.id },
          data: { updatedAt: new Date(0) },
        });
        await h.dispatch();
        expect(
          await db.backgroundTask.findUnique({ where: { id: task.id } }),
        ).toMatchObject({ status: 'WAITING', executionId: null });
        await h.worker(newDb as unknown as PrismaClient);
        await h.dispatch();
        await within(claimedB.promise);
        const replacement = await db.backgroundTask.findUniqueOrThrow({
          where: { id: task.id },
        });
        expect(replacement.executionId).not.toBe(first.executionId);
        expect(replacement.executionId).not.toBeNull();
        releaseCallback.open();
        await Promise.all(lateWrites);
        expect(
          await db.backgroundTask.findUnique({ where: { id: task.id } }),
        ).toEqual(replacement);
        const newJob = (await h.queue.getJob(taskJobId(task.id, task.run)))!;
        const finished = newJob.waitUntilFinished(h.events, 8000);
        resumeB.open();
        await finished;
        expect(await newJob.getState()).toBe('completed');
        const completed = await db.backgroundTask.findUniqueOrThrow({
          where: { id: task.id },
        });
        expect(completed).toMatchObject({
          status: 'COMPLETED',
          progress: 100,
          executionId: null,
          lastError: null,
          run: task.run,
        });
        expect(completed.resultCsv).toContain(item.title);
        expect(completed.resultSummary).toBe('已导出 1 条工单');
        expect(completed.expiresAt!.getTime()).toBeGreaterThan(Date.now());
        expect((await h.queue.getJob(newJob.id!))?.progress).toBe(100);
      } finally {
        releaseCallback.open();
        await Promise.allSettled(lateWrites);
        await h.close();
      }
    });

    it('真实 Worker 临时失败自动重试，failed 日志不提前终止下一次执行', async () => {
      const h = await harness();
      const { task, item } = await fixture();
      let reads = 0;
      const client = db.$extends({
        query: {
          ticket: {
            async findMany({ args, query }) {
              if (++reads === 1) throw new Error('测试注入：首次读取暂时失败');
              return query(args);
            },
          },
        },
      });
      try {
        const consumer = await h.worker(client as unknown as PrismaClient);
        const job = await h.queue.add(
          'execute',
          { taskId: task.id, run: task.run },
          {
            jobId: taskJobId(task.id, task.run),
            attempts: 3,
            backoff: { type: 'exponential', delay: 30 },
          },
        );
        await job.waitUntilFinished(h.events, 8000);
        expect(
          consumer.records.filter((record) => record.event === 'task.failed'),
        ).toEqual([expect.objectContaining({ exhausted: false })]);
        expect((await h.queue.getJob(job.id!))?.attemptsMade).toBe(2);
        const completed = await db.backgroundTask.findUniqueOrThrow({
          where: { id: task.id },
        });
        expect(completed).toMatchObject({
          status: 'COMPLETED',
          attempts: 2,
          run: 0,
          progress: 100,
          executionId: null,
          lastError: null,
        });
        expect(completed.resultCsv).toContain(item.title);
      } finally {
        await h.close();
      }
    });

    it('真实 Worker 的永久拒绝同时保留数据库与队列 FAILED，不依赖回调代写状态', async () => {
      const h = await harness();
      const { task } = await fixture();
      await db.backgroundTask.update({
        where: { id: task.id },
        data: { filters: { q: 123 } },
      });
      try {
        const consumer = await h.worker(db);
        const job = await h.queue.add(
          'execute',
          { taskId: task.id, run: task.run },
          {
            jobId: taskJobId(task.id, task.run),
            attempts: 3,
          },
        );
        await consumer.event('task.failed');
        expect(
          consumer.records.find((record) => record.event === 'task.failed'),
        ).toMatchObject({ exhausted: true });
        expect(await job.getState()).toBe('failed');
        expect((await h.queue.getJob(job.id!))?.attemptsMade).toBe(1);
        expect(
          await db.backgroundTask.findUnique({ where: { id: task.id } }),
        ).toMatchObject({
          status: 'FAILED',
          attempts: 1,
          executionId: null,
          resultCsv: null,
          lastError: '导出筛选条件无效',
        });
      } finally {
        await h.close();
      }
    });

    it('stalled 超过恢复上限后，由扫描核对真实队列 FAILED，旧执行者返回也不能改写终态', async () => {
      const h = await harness();
      const { task } = await fixture();
      const pausedA = gate(),
        resumeA = gate();
      h.releases.push(resumeA.open);
      const oldDb = db.$extends({
        query: {
          ticket: {
            async findMany({ args, query }) {
              await query(args);
              pausedA.open();
              await resumeA.promise;
              throw new Error('测试注入：旧执行者晚到失败');
            },
          },
        },
      });
      try {
        const a = await h.worker(oldDb as unknown as PrismaClient, {
          skipLockRenewal: true,
          skipStalledCheck: true,
          lockDuration: 1000,
        });
        const job = await h.queue.add(
          'execute',
          { taskId: task.id, run: task.run },
          { jobId: taskJobId(task.id, task.run) },
        );
        await within(pausedA.promise);
        const b = await h.worker(db, { maxStalledCount: 0 });
        await b.event('task.failed');
        expect(await job.getState()).toBe('failed');
        expect((await h.queue.getJob(job.id!))?.failedReason).toContain(
          'stalled',
        );
        expect(
          await db.backgroundTask.findUnique({ where: { id: task.id } }),
        ).toMatchObject({ status: 'RUNNING' });
        await db.backgroundTask.update({
          where: { id: task.id },
          data: { updatedAt: new Date(0) },
        });
        await h.dispatch();
        expect(
          await db.backgroundTask.findUnique({ where: { id: task.id } }),
        ).toMatchObject({ status: 'WAITING', executionId: null });
        await h.dispatch();
        const terminal = await db.backgroundTask.findUniqueOrThrow({
          where: { id: task.id },
        });
        expect(terminal).toMatchObject({
          status: 'FAILED',
          executionId: null,
          resultCsv: null,
          resultSummary: null,
          expiresAt: null,
          lastError: '队列执行失败，请手动重试',
        });
        resumeA.open();
        await a.event('task.skipped');
        expect(
          await db.backgroundTask.findUnique({ where: { id: task.id } }),
        ).toEqual(terminal);
        expect(await job.getState()).toBe('failed');
      } finally {
        await h.close();
      }
    });
  },
);