yuban-server/src/job-store.mjs

import { mkdir, readFile, readdir, rename, unlink, writeFile } from 'node:fs/promises';
import { join } from 'node:path';

export class JobStore {
  constructor(directory, retentionMs) {
    this.directory = directory;
    this.retentionMs = retentionMs;
    this.jobs = new Map();
    this.idempotency = new Map();
  }

  async init() {
    await mkdir(this.directory, { recursive: true });
    for (const name of await readdir(this.directory)) {
      if (!name.endsWith('.json')) continue;
      try {
        const job = JSON.parse(await readFile(join(this.directory, name), 'utf8'));
        if (!job?.id) continue;
        this.jobs.set(job.id, job);
        if (job.idempotencyKey) {
          const current = this.get(this.idempotency.get(job.idempotencyKey));
          if (!current || Date.parse(job.createdAt || '') >= Date.parse(current.createdAt || '')) {
            this.idempotency.set(job.idempotencyKey, job.id);
          }
        }
      } catch {
        // 损坏的单个临时任务不会阻止服务启动。
      }
    }
    await this.cleanupExpired();
  }

  get(id) {
    return this.jobs.get(id) || null;
  }

  findByIdempotency(key) {
    const id = this.idempotency.get(key);
    return id ? this.get(id) : null;
  }

  listRecoverable() {
    return [...this.jobs.values()].filter(job => ['queued', 'running'].includes(job.status));
  }

  listByOwner(ownerKey) {
    return [...this.jobs.values()].filter(job => job.ownerKey === ownerKey);
  }

  async create(job) {
    this.jobs.set(job.id, job);
    if (job.idempotencyKey) this.idempotency.set(job.idempotencyKey, job.id);
    await this.persist(job);
    return job;
  }

  async update(id, patch) {
    const current = this.get(id);
    if (!current) return null;
    const next = {
      ...current,
      ...patch,
      updatedAt: new Date().toISOString(),
    };
    this.jobs.set(id, next);
    await this.persist(next);
    return next;
  }

  async persist(job) {
    const finalPath = join(this.directory, `${job.id}.json`);
    const temporaryPath = join(this.directory, `.${job.id}.${process.pid}.tmp`);
    await writeFile(temporaryPath, JSON.stringify(job), { mode: 0o600 });
    await rename(temporaryPath, finalPath);
  }

  async cleanupExpired(now = Date.now()) {
    for (const job of [...this.jobs.values()]) {
      const expiresAt = Date.parse(job.expiresAt || '');
      if (!Number.isFinite(expiresAt) || expiresAt > now) continue;
      this.jobs.delete(job.id);
      if (job.idempotencyKey && this.idempotency.get(job.idempotencyKey) === job.id) {
        this.idempotency.delete(job.idempotencyKey);
      }
      await unlink(join(this.directory, `${job.id}.json`)).catch(() => undefined);
    }
  }
}

export function publicJob(job) {
  return {
    id: job.id,
    status: job.status,
    stage: job.stage,
    progress: job.progress,
    createdAt: job.createdAt,
    updatedAt: job.updatedAt,
    heartbeatAt: job.heartbeatAt || null,
    expiresAt: job.expiresAt,
    result: job.status === 'completed' ? job.result : undefined,
    error: job.status === 'failed' ? job.error : undefined,
  };
}