tests/worker-process.integration.test.ts

import { afterAll, beforeAll, describe, expect, it } from 'vitest';
import { fork, spawn, type ChildProcess } from 'node:child_process';
import { randomUUID } from 'node:crypto';
import { EventEmitter } from 'node:events';
import { createServer, connect, type Socket } from 'node:net';
import { Queue } from 'bullmq';
import IORedis from 'ioredis';
import { PrismaPg } from '@prisma/adapter-pg';
import { PrismaClient } from '@/generated/prisma/client';
import { heartbeatRegistryKey, hasHeartbeat } from '@/server/jobs/heartbeat';
import { queueName, taskJobId, type TaskJob } from '@/server/jobs/queue';

type ProcessInstance = {
  child: ChildProcess;
  scope: string;
  records: { event: string; taskId?: string; instanceId?: string }[];
  exited: Promise<{ code: number | null; signal: string | null }>;
  event(name: string): Promise<void>;
  resume(label: string): void;
  stop(): Promise<{ code: number | null; signal: string | null }>;
};

async function within<T>(promise: Promise<T>, ms = 10000): Promise<T> {
  let timer!: ReturnType<typeof setTimeout>;
  try {
    return await Promise.race([
      promise,
      new Promise<never>((_, reject) => {
        timer = setTimeout(() => reject(new Error('进程集成测试等待超时')), ms);
      }),
    ]);
  } finally {
    clearTimeout(timer);
  }
}
async function until(check: () => Promise<boolean>, ms = 10000) {
  const deadline = Date.now() + ms;
  while (!(await check())) {
    if (Date.now() >= deadline) throw new Error('进程状态条件未满足');
    // 这里只轮询可观察状态；谁先执行由 IPC、日志和数据库锁确定。
    await new Promise((resolve) => setTimeout(resolve, 20));
  }
}

describe.skipIf(!process.env.TEST_DATABASE_URL || !process.env.TEST_REDIS_URL)(
  '真实进程的 Worker 心跳、SIGTERM 与恢复',
  () => {
    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: '进程测试', email: `${userId}@example.com` },
      });
      teamId = (
        await db.team.create({
          data: {
            name: '进程隔离团队',
            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: `Process-${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 = `process-test-${randomUUID()}`;
      const redis = new IORedis(process.env.TEST_REDIS_URL!, {
        maxRetriesPerRequest: null,
      });
      const queue = new Queue<TaskJob>(queueName, {
        connection: redis,
        prefix,
      });
      const children: ProcessInstance[] = [];
      const baseEnv = {
        ...process.env,
        DATABASE_URL: process.env.TEST_DATABASE_URL!,
        REDIS_URL: process.env.TEST_REDIS_URL!,
        QUEUE_PREFIX: prefix,
        LOG_LEVEL: 'info',
      };
      function start(
        config: Record<string, unknown> = {},
        scope = randomUUID(),
      ): ProcessInstance {
        const child = fork('tests/fixtures/worker-process.ts', [], {
          execArgv: ['--import', 'tsx'],
          silent: true,
          env: {
            ...baseEnv,
            WORKER_INSTANCE_ID: scope,
            WORKER_PROCESS_FIXTURE: JSON.stringify({ teamId, ...config }),
          },
        });
        const bus = new EventEmitter();
        const checkpoints: string[] = [];
        const records: {
          event: string;
          taskId?: string;
          instanceId?: string;
        }[] = [];
        let output = '',
          stderr = '';
        child.stdout!.on('data', (data) => {
          output += String(data);
          let end;
          while ((end = output.indexOf('\n')) >= 0) {
            const line = output.slice(0, end);
            output = output.slice(end + 1);
            try {
              const record = JSON.parse(line);
              records.push(record);
              bus.emit(record.event);
            } catch {
              /* 忽略非 JSON 工具输出。 */
            }
          }
        });
        child.stderr!.on('data', (data) => {
          stderr += String(data).slice(0, 2000);
        });
        child.on('message', (message) => {
          const label = (message as { checkpoint?: string }).checkpoint;
          if (label) {
            checkpoints.push(label);
            bus.emit(`checkpoint:${label}`);
          }
        });
        const exited = new Promise<{
          code: number | null;
          signal: string | null;
        }>((resolve, reject) => {
          child.once('error', reject);
          child.once('close', (code, signal) => resolve({ code, signal }));
        });
        const result = {
          child,
          scope,
          records,
          exited,
          async event(name: string) {
            if (
              records.some((record) => record.event === name) ||
              (name.startsWith('checkpoint:') &&
                checkpoints.includes(name.slice(11)))
            )
              return;
            await within(
              Promise.race([
                new Promise<void>((resolve) => bus.once(name, resolve)),
                exited.then((state) => {
                  throw new Error(
                    `等待 ${name} 时进程已退出：${state.code}/${state.signal} ${stderr}`,
                  );
                }),
              ]),
            );
          },
          resume(label: string) {
            if (child.connected) child.send({ resume: label });
          },
          async stop() {
            if (child.exitCode === null && child.signalCode === null)
              child.kill('SIGTERM');
            return within(exited);
          },
        };
        children.push(result);
        return result;
      }
      return {
        prefix,
        redis,
        queue,
        start,
        healthy: (scope?: string) => hasHeartbeat(redis, scope, prefix),
        async probe(scope: string, redisUrl = process.env.TEST_REDIS_URL!) {
          const child = spawn(
            process.execPath,
            ['--import', 'tsx', 'scripts/worker-health.ts'],
            {
              env: {
                ...baseEnv,
                REDIS_URL: redisUrl,
                WORKER_INSTANCE_ID: scope,
              },
              stdio: 'ignore',
            },
          );
          try {
            return await within(
              new Promise<number | null>((resolve, reject) => {
                child.once('error', reject);
                child.once('close', resolve);
              }),
              7000,
            );
          } finally {
            if (child.exitCode === null && child.signalCode === null)
              child.kill('SIGKILL');
          }
        },
        async close() {
          for (const current of children) {
            if (
              current.child.exitCode === null &&
              current.child.signalCode === null
            )
              current.child.kill('SIGKILL');
          }
          await Promise.all(children.map((current) => current.exited));
          await queue.obliterate({ force: true });
          for (const key of await redis.keys(`${prefix}:worker-heartbeat:*`))
            await redis.del(key);
          await queue.close();
          redis.disconnect();
        },
      };
    }
    async function completed(id: string) {
      await until(
        async () =>
          (await db.backgroundTask.findUnique({ where: { id } }))?.status ===
          'COMPLETED',
      );
      return db.backgroundTask.findUniqueOrThrow({ where: { id } });
    }
    const checkResult = async (
      kind: 'EXPORT' | 'REMINDER',
      data: Awaited<ReturnType<typeof fixture>>,
    ) => {
      const result = await completed(data.task.id);
      expect(result).toMatchObject({
        status: 'COMPLETED',
        progress: 100,
        executionId: null,
        lastError: null,
        run: data.task.run,
      });
      if (kind === 'EXPORT') {
        expect(result.resultCsv).toContain(data.item.title);
        expect(result.resultSummary).toBe('已导出 1 条工单');
        expect(result.expiresAt!.getTime()).toBeGreaterThan(Date.now());
      } else {
        expect(result.resultSummary).toBe('已生成 1 条站内提醒');
        expect(
          await db.reminderNotification.count({
            where: { ticketId: data.item.id },
          }),
        ).toBe(1);
        const events = await db.ticketEvent.findMany({
          where: { ticketId: data.item.id, kind: 'TIMEOUT_REMINDED' },
        });
        expect(events).toHaveLength(1);
        expect(events[0].version).toBe(data.item.version + 1);
        expect(
          await db.ticket.findUnique({ where: { id: data.item.id } }),
        ).toMatchObject({ version: data.item.version + 1 });
      }
      return result;
    };

    it.each(['EXPORT', 'REMINDER'] as const)(
      'SIGTERM 排空在途 %s、持续续锁，不再认领新任务',
      async (kind) => {
        const h = await harness();
        const data = await fixture(kind);
        // 提前准备待投递夹具，避免提醒事务持有团队锁时，新建行的外键检查等待该锁。
        const waiting = await fixture();
        await db.backgroundTask.update({
          where: { id: waiting.task.id },
          data: { status: 'FAILED' },
        });
        try {
          const a = h.start({
            pause: 'read',
            taskId: data.task.id,
            ticketId: data.item.id,
            heartbeatTtlMs: 5000,
          });
          await a.event('checkpoint:read');
          await until(() => h.healthy(a.scope));
          expect(await h.probe(a.scope)).toBe(0);
          const leased = await db.backgroundTask.findUniqueOrThrow({
            where: { id: data.task.id },
          });
          expect(leased.executionId).not.toBeNull();
          a.child.kill('SIGTERM');
          await a.event('worker.stopping');
          await until(async () => !(await h.healthy(a.scope)), 2000);
          expect(await h.probe(a.scope)).toBe(1);
          a.child.kill('SIGTERM');
          await db.backgroundTask.update({
            where: { id: waiting.task.id },
            data: { status: 'WAITING' },
          });
          const deadline = Date.now() + 1100;
          await until(
            async () =>
              Date.now() > deadline &&
              (await h.redis.pttl(
                `${h.queue.toKey(taskJobId(data.task.id, data.task.run))}:lock`,
              )) > 0,
          );
          expect(
            await db.backgroundTask.findUnique({ where: { id: data.task.id } }),
          ).toEqual(leased);
          a.resume('read');
          expect(await within(a.exited)).toEqual({ code: 0, signal: null });
          expect(
            a.records.filter((record) => record.event === 'worker.stopping'),
          ).toHaveLength(1);
          const original = await checkResult(kind, data);
          expect(
            await db.backgroundTask.findUnique({
              where: { id: waiting.task.id },
            }),
          ).toMatchObject({ status: 'WAITING', executionId: null });
          const b = h.start();
          await completed(waiting.task.id);
          expect(await b.stop()).toEqual({ code: 0, signal: null });
          expect(
            await db.backgroundTask.findUnique({ where: { id: data.task.id } }),
          ).toEqual(original);
          expect(await h.healthy()).toBe(false);
        } finally {
          await h.close();
        }
      },
      30000,
    );

    it.each(['EXPORT', 'REMINDER'] as const)(
      'SIGTERM 超时退出后 %s 保留未完成租约，新进程 stalled 接手并只提交一次',
      async (kind) => {
        const h = await harness();
        const data = await fixture(kind);
        try {
          const a = h.start({
            pause: 'read',
            taskId: data.task.id,
            ticketId: data.item.id,
            shutdownTimeoutMs: 1000,
          });
          await a.event('checkpoint:read');
          const first = await db.backgroundTask.findUniqueOrThrow({
            where: { id: data.task.id },
          });
          a.child.kill('SIGTERM');
          expect(await within(a.exited)).toEqual({ code: 1, signal: null });
          expect(
            a.records.some(
              (record) => record.event === 'worker.shutdown_timeout',
            ),
          ).toBe(true);
          expect(
            await db.backgroundTask.findUnique({ where: { id: data.task.id } }),
          ).toEqual(first);
          expect(
            await db.reminderNotification.count({
              where: { ticketId: data.item.id },
            }),
          ).toBe(0);
          expect(
            await db.ticket.findUnique({ where: { id: data.item.id } }),
          ).toMatchObject({ version: data.item.version });
          expect(await h.healthy()).toBe(false);
          const b = h.start({ pause: 'claim' });
          await b.event('task.stalled');
          await b.event('checkpoint:claim');
          const replacement = await db.backgroundTask.findUniqueOrThrow({
            where: { id: data.task.id },
          });
          expect(replacement).toMatchObject({
            status: 'RUNNING',
            run: first.run,
          });
          expect(replacement.executionId).not.toBeNull();
          expect(replacement.executionId).not.toBe(first.executionId);
          b.resume('claim');
          await checkResult(kind, data);
          await until(
            async () =>
              (await h.queue.getJob(taskJobId(data.task.id, data.task.run)))
                ?.progress === 100,
          );
          expect(
            await (await h.queue.getJob(
              taskJobId(data.task.id, data.task.run),
            ))!.getState(),
          ).toBe('completed');
          expect(await b.stop()).toEqual({ code: 0, signal: null });
        } finally {
          await h.close();
        }
      },
      30000,
    );

    it('同作用域两个进程退出一个，不删除同伴心跳；最后退出立即不健康', async () => {
      const h = await harness();
      const scope = randomUUID();
      try {
        const a = h.start({}, scope),
          b = h.start({ heartbeatTtlMs: 5000 }, scope);
        await until(
          async () =>
            (await h.redis.zcard(heartbeatRegistryKey(scope, h.prefix))) === 2,
        );
        expect(await h.probe(scope)).toBe(0);
        // 冻结续期，防止错误删键后恰好被同伴的新心跳掩盖。
        b.child.kill('SIGSTOP');
        expect(await a.stop()).toEqual({ code: 0, signal: null });
        expect(await h.healthy(scope)).toBe(true);
        expect(await h.redis.zcard(heartbeatRegistryKey(scope, h.prefix))).toBe(
          1,
        );
        b.child.kill('SIGCONT');
        expect(await h.probe(scope)).toBe(0);
        expect(await b.stop()).toEqual({ code: 0, signal: null });
        expect(await h.healthy()).toBe(false);
        expect(await h.probe(scope)).toBe(1);
      } finally {
        await h.close();
      }
    }, 30000);

    it('冻结一个进程后其作用域探针过期失败，其他进程仍维持全局健康', async () => {
      const h = await harness();
      try {
        const a = h.start(),
          b = h.start();
        await until(
          async () => (await h.healthy(a.scope)) && (await h.healthy(b.scope)),
        );
        a.child.kill('SIGSTOP');
        await until(async () => !(await h.healthy(a.scope)));
        expect(await h.healthy()).toBe(true);
        expect(await h.probe(a.scope)).toBe(1);
        expect(await h.probe(b.scope)).toBe(0);
        a.child.kill('SIGCONT');
        expect(await a.stop()).toEqual({ code: 0, signal: null });
        expect(await h.healthy(b.scope)).toBe(true);
        expect(await b.stop()).toEqual({ code: 0, signal: null });
      } finally {
        await h.close();
      }
    }, 30000);

    it('SIGTERM 等待在途扫描，扫描返回后不复活心跳，新进程接手它留下的排队任务', async () => {
      const h = await harness();
      try {
        const a = h.start({ pauseDispatch: true, heartbeatTtlMs: 5000 });
        await a.event('checkpoint:dispatch');
        expect(await h.healthy(a.scope)).toBe(true);
        a.child.kill('SIGTERM');
        await a.event('worker.stopping');
        await until(async () => !(await h.healthy(a.scope)), 2000);
        const data = await fixture();
        a.resume('dispatch');
        expect(await within(a.exited)).toEqual({ code: 0, signal: null });
        expect(await h.healthy()).toBe(false);
        expect(
          await db.backgroundTask.findUnique({ where: { id: data.task.id } }),
        ).toMatchObject({ status: 'WAITING', executionId: null });
        const b = h.start();
        await checkResult('EXPORT', data);
        expect(await b.stop()).toEqual({ code: 0, signal: null });
      } finally {
        await h.close();
      }
    }, 30000);

    it('真实 Redis 连接成功后读取无响应，健康探针仍在总时限内失败退出', async () => {
      const h = await harness();
      const sockets = new Set<Socket>();
      let reached!: () => void;
      const blocked = new Promise<void>((resolve) => {
        reached = resolve;
      });
      // 仅代理本次探针连接：转发握手，完整解析 RESP 命令后截住 EVAL 的响应。
      const target = new URL(process.env.TEST_REDIS_URL!);
      const server = createServer((client) => {
        const upstream = connect(Number(target.port || 6379), target.hostname);
        sockets.add(client);
        sockets.add(upstream);
        upstream.on('error', () => client.destroy());
        client.on('error', () => upstream.destroy());
        upstream.on('data', (data) => client.write(data));
        let buffered = Buffer.alloc(0);
        client.on('data', (data) => {
          buffered = Buffer.concat([buffered, data]);
          while (buffered.length) {
            const line = buffered.indexOf('\r\n');
            if (line < 0) return;
            const count = Number(buffered.subarray(1, line).toString());
            let cursor = line + 2,
              command = '';
            let complete = true;
            for (let i = 0; i < count; i++) {
              const end = buffered.indexOf('\r\n', cursor);
              if (end < 0) {
                complete = false;
                break;
              }
              const size = Number(
                buffered.subarray(cursor + 1, end).toString(),
              );
              if (buffered.length < end + 2 + size + 2) {
                complete = false;
                break;
              }
              if (i === 0)
                command = buffered
                  .subarray(end + 2, end + 2 + size)
                  .toString()
                  .toLowerCase();
              cursor = end + 2 + size + 2;
            }
            if (!complete) return;
            const raw = buffered.subarray(0, cursor);
            buffered = buffered.subarray(cursor);
            if (command === 'eval') reached();
            else upstream.write(raw);
          }
        });
      });
      try {
        await new Promise<void>((resolve) =>
          server.listen(0, '127.0.0.1', resolve),
        );
        const proxy = new URL(target);
        proxy.hostname = '127.0.0.1';
        proxy.port = String((server.address() as { port: number }).port);
        const probing = h.probe(randomUUID(), proxy.toString());
        await within(blocked);
        const readingSince = Date.now();
        expect(await probing).toBe(1);
        // 证明是共享健康查询的 2 秒读取上限，而非短脚本外层 4 秒兜底。
        expect(Date.now() - readingSince).toBeLessThan(3000);
      } finally {
        sockets.forEach((socket) => socket.destroy());
        await new Promise<void>((resolve) => server.close(() => resolve()));
        await h.close();
      }
    }, 30000);
  },
);