329 lines
10 KiB
TypeScript
329 lines
10 KiB
TypeScript
import http from "node:http";
|
|
import { createReadStream } from "node:fs";
|
|
import {
|
|
mkdir,
|
|
readFile,
|
|
rename,
|
|
readdir,
|
|
rm,
|
|
stat,
|
|
statfs,
|
|
open,
|
|
} from "node:fs/promises";
|
|
import path from "node:path";
|
|
import { createHash, timingSafeEqual } from "node:crypto";
|
|
import { pipeline } from "node:stream/promises";
|
|
import { z } from "zod";
|
|
import {
|
|
videoUpscaleSchema,
|
|
type VideoUpscaleInput,
|
|
} from "../../lib/video-canvas/video-upscale-contract";
|
|
import { upscaleJob } from "./job";
|
|
|
|
const root = path.resolve(
|
|
process.env.UPSCALE_DATA_DIR || "/var/lib/video-upscale"
|
|
);
|
|
const token = process.env.UPSCALE_API_TOKEN || "";
|
|
if (token.length < 32)
|
|
throw Error("UPSCALE_API_TOKEN must have at least 32 characters");
|
|
const jobs = new Map<string, Job>();
|
|
type Job = {
|
|
id: string;
|
|
digest: string;
|
|
tenant: string;
|
|
plan: VideoUpscaleInput;
|
|
status: "queued" | "running" | "succeeded" | "failed";
|
|
stage: string;
|
|
createdAt: string;
|
|
updatedAt: string;
|
|
attempts: number;
|
|
error?: string;
|
|
result?: Record<string, any>;
|
|
};
|
|
const inputSchema = z.object({
|
|
id: z.string().regex(/^[a-f0-9]{64}$/),
|
|
tenant: z.string().regex(/^[a-f0-9]{64}$/),
|
|
plan: videoUpscaleSchema,
|
|
});
|
|
let busy = false,
|
|
stopping = false,
|
|
accepting = false;
|
|
const file = (id: string) => path.join(root, id + ".json");
|
|
async function persist(job: Job) {
|
|
job.updatedAt = new Date().toISOString();
|
|
const temp = file(job.id) + ".tmp";
|
|
const handle = await open(temp, "w", 0o600);
|
|
try {
|
|
await handle.writeFile(JSON.stringify(job));
|
|
await handle.sync();
|
|
} finally {
|
|
await handle.close();
|
|
}
|
|
await rename(temp, file(job.id));
|
|
const directory = await open(root, "r");
|
|
try {
|
|
await directory.sync();
|
|
} finally {
|
|
await directory.close();
|
|
}
|
|
jobs.set(job.id, job);
|
|
}
|
|
function publicJob(job: Job) {
|
|
return {
|
|
id: job.id,
|
|
status: job.status,
|
|
stage: job.stage,
|
|
attempts: job.attempts,
|
|
createdAt: job.createdAt,
|
|
updatedAt: job.updatedAt,
|
|
error: job.error,
|
|
result: job.result,
|
|
};
|
|
}
|
|
async function prune() {
|
|
for (const job of jobs.values())
|
|
if (
|
|
["succeeded", "failed"].includes(job.status) &&
|
|
Date.now() - Date.parse(job.updatedAt) > 7 * 86400000
|
|
) {
|
|
await rm(path.join(root, job.id), { recursive: true, force: true });
|
|
await rm(file(job.id), { force: true });
|
|
jobs.delete(job.id);
|
|
}
|
|
}
|
|
async function work() {
|
|
if (busy || stopping) return;
|
|
const job = [...jobs.values()].find((j) => j.status === "queued");
|
|
if (!job) return;
|
|
busy = true;
|
|
const dir = path.join(root, job.id);
|
|
try {
|
|
await prune();
|
|
const disk = await statfs(root);
|
|
if (disk.bavail * disk.bsize < 10 * 1024 ** 3)
|
|
throw Error("超分磁盘剩余空间不足 10 GiB");
|
|
job.status = "running";
|
|
job.attempts++;
|
|
job.stage = "准备素材";
|
|
await persist(job);
|
|
await rm(dir, { recursive: true, force: true });
|
|
await mkdir(dir, { mode: 0o700 });
|
|
const result = await upscaleJob(
|
|
job.plan,
|
|
job.tenant,
|
|
dir,
|
|
async (stage) => {
|
|
job.stage = stage;
|
|
await persist(job);
|
|
}
|
|
);
|
|
job.result = result;
|
|
// Keep only the deliverables. Intermediate files must not accumulate per job.
|
|
for (const name of await readdir(dir))
|
|
if (name !== "result.mp4")
|
|
await rm(path.join(dir, name), { recursive: true, force: true }).catch(
|
|
() => {
|
|
console.warn(
|
|
JSON.stringify({ event: "upscale-cleanup-deferred", id: job.id })
|
|
);
|
|
}
|
|
);
|
|
job.status = "succeeded";
|
|
job.stage = "超分完成";
|
|
await persist(job);
|
|
} catch (error) {
|
|
if (stopping) return; // Preserve running journal; restart will resume this same job.
|
|
job.status = "failed";
|
|
job.stage = "超分失败";
|
|
// Do not expose ffmpeg arguments, private cache paths or signed source URLs.
|
|
job.error =
|
|
error instanceof Error &&
|
|
!/Command failed|https?:|\/var\/|\/tmp\//.test(error.message)
|
|
? error.message.slice(0, 300)
|
|
: "素材处理失败，请检查格式、参数或服务日志";
|
|
await rm(dir, { recursive: true, force: true }).catch(() => {});
|
|
await persist(job);
|
|
console.error(
|
|
JSON.stringify({ event: "upscale-failed", id: job.id, error: job.error })
|
|
);
|
|
} finally {
|
|
busy = false;
|
|
}
|
|
setImmediate(
|
|
() =>
|
|
void work().catch((e) => {
|
|
console.error("Render queue persistence failure");
|
|
process.exit(1);
|
|
})
|
|
);
|
|
}
|
|
function json(res: http.ServerResponse, status: number, value: unknown) {
|
|
res.writeHead(status, {
|
|
"content-type": "application/json; charset=utf-8",
|
|
"cache-control": "no-store",
|
|
});
|
|
res.end(JSON.stringify(value));
|
|
}
|
|
async function body(req: http.IncomingMessage) {
|
|
let size = 0;
|
|
const chunks: Buffer[] = [];
|
|
for await (const chunk of req) {
|
|
size += chunk.length;
|
|
if (size > 1024 * 1024) throw Error("请求超过 1 MiB");
|
|
chunks.push(chunk);
|
|
}
|
|
return JSON.parse(Buffer.concat(chunks).toString("utf8"));
|
|
}
|
|
function authenticated(req: http.IncomingMessage) {
|
|
const expected = Buffer.from("Bearer " + token),
|
|
actual = Buffer.from(req.headers.authorization || "");
|
|
return expected.length === actual.length && timingSafeEqual(expected, actual);
|
|
}
|
|
async function main() {
|
|
await mkdir(root, { recursive: true, mode: 0o700 });
|
|
// This service runs as a single systemd process. A restart requeues local work at most twice.
|
|
for (const name of await readdir(root))
|
|
if (/^[a-f0-9]{64}\.json$/.test(name)) {
|
|
const job: Job = JSON.parse(
|
|
await readFile(path.join(root, name), "utf8")
|
|
);
|
|
if (job.id + ".json" !== name) throw Error("Invalid upscale record");
|
|
if (job.status === "running") {
|
|
job.status = job.attempts < 3 ? "queued" : "failed";
|
|
job.stage =
|
|
job.status === "queued"
|
|
? "服务重启，等待恢复超分"
|
|
: "多次中断，已停止";
|
|
if (job.status === "failed") job.error = job.stage;
|
|
await persist(job);
|
|
}
|
|
jobs.set(job.id, job);
|
|
}
|
|
await prune();
|
|
const server = http.createServer(async (req, res) => {
|
|
if (!authenticated(req)) return json(res, 401, { error: "Unauthorized" });
|
|
const url = new URL(req.url || "/", "http://localhost");
|
|
try {
|
|
if (req.method === "GET" && url.pathname === "/health")
|
|
return json(res, 200, {
|
|
status: "ok",
|
|
version: 1,
|
|
active: busy ? 1 : 0,
|
|
queued: [...jobs.values()].filter((j) => j.status === "queued")
|
|
.length,
|
|
render: "realesr-general-x4v3",
|
|
gpu: process.env.CUDA_VISIBLE_DEVICES,
|
|
});
|
|
if (req.method === "POST" && url.pathname === "/v1/upscales") {
|
|
const input = inputSchema.parse(await body(req));
|
|
// Serialize admission across the async filesystem write to preserve idempotency.
|
|
if (accepting) return json(res, 503, { error: "请重试提交" });
|
|
accepting = true;
|
|
try {
|
|
const digest = createHash("sha256")
|
|
.update(JSON.stringify({ tenant: input.tenant, plan: input.plan }))
|
|
.digest("hex");
|
|
const existing = jobs.get(input.id);
|
|
if (existing)
|
|
return json(
|
|
res,
|
|
existing.digest === digest ? 200 : 409,
|
|
existing.digest === digest
|
|
? publicJob(existing)
|
|
: { error: "任务 ID 已用于不同参数" }
|
|
);
|
|
if (
|
|
jobs.size >= 200 ||
|
|
[...jobs.values()].filter((j) =>
|
|
["queued", "running"].includes(j.status)
|
|
).length >= 30
|
|
)
|
|
return json(res, 429, { error: "超分队列已满，请稍后重试" });
|
|
const job: Job = {
|
|
...input,
|
|
digest,
|
|
status: "queued",
|
|
stage: "排队中",
|
|
attempts: 0,
|
|
createdAt: new Date().toISOString(),
|
|
updatedAt: new Date().toISOString(),
|
|
};
|
|
await persist(job);
|
|
json(res, 202, publicJob(job));
|
|
void work().catch(() => {
|
|
console.error("Render queue failed");
|
|
process.exit(1);
|
|
});
|
|
return;
|
|
} finally {
|
|
accepting = false;
|
|
}
|
|
}
|
|
const match = url.pathname.match(
|
|
/^\/v1\/upscales\/([a-f0-9]{64})(?:\/(video))?$/
|
|
);
|
|
if (req.method === "GET" && match) {
|
|
const job = jobs.get(match[1]);
|
|
if (!job) return json(res, 404, { error: "任务不存在或已过期" });
|
|
if (!match[2]) return json(res, 200, publicJob(job));
|
|
if (job.status !== "succeeded")
|
|
return json(res, 409, { error: "成片尚未就绪" });
|
|
const output = path.join(root, job.id, "result.mp4");
|
|
const info = await stat(output);
|
|
res.writeHead(200, {
|
|
"content-type": "video/mp4",
|
|
"content-length": info.size,
|
|
"cache-control": "no-store",
|
|
});
|
|
await pipeline(createReadStream(output), res);
|
|
return;
|
|
}
|
|
json(res, 404, { error: "Not found" });
|
|
} catch (error) {
|
|
if (res.headersSent) {
|
|
res.destroy();
|
|
return;
|
|
}
|
|
json(res, 400, {
|
|
error:
|
|
error instanceof z.ZodError
|
|
? "超分参数无效"
|
|
: error instanceof Error && !/https?:|\/var\//.test(error.message)
|
|
? error.message.slice(0, 300)
|
|
: "请求处理失败",
|
|
});
|
|
}
|
|
});
|
|
server.requestTimeout = 120000;
|
|
server.headersTimeout = 10000;
|
|
server.listen(
|
|
Number(process.env.UPSCALE_PORT || 18341),
|
|
process.env.UPSCALE_HOST || "127.0.0.1"
|
|
);
|
|
const fatal = () => {
|
|
console.error("Render queue persistence failure");
|
|
process.exit(1);
|
|
};
|
|
const timer = setInterval(() => {
|
|
void work().catch(fatal);
|
|
}, 3000);
|
|
const cleanupTimer = setInterval(() => {
|
|
if (!busy && !accepting) void prune().catch(fatal);
|
|
}, 3600000);
|
|
cleanupTimer.unref();
|
|
for (const signal of ["SIGTERM", "SIGINT"])
|
|
process.on(signal, () => {
|
|
stopping = true;
|
|
clearInterval(timer);
|
|
clearInterval(cleanupTimer);
|
|
server.close();
|
|
// systemd KillMode=mixed also terminates FFmpeg; persisted running jobs resume on restart.
|
|
setTimeout(() => process.exit(0), 1000).unref();
|
|
});
|
|
console.log("Video upscale service ready");
|
|
}
|
|
void main().catch(() => {
|
|
console.error("Video upscale service startup failed");
|
|
process.exitCode = 1;
|
|
});
|