ai-chatbot/lib/canvas/download-zip-stream.ts

38 lines
2.0 KiB
TypeScript

import "server-only";
import { createReadStream, createWriteStream } from "node:fs";
import { mkdir, stat, statfs, unlink, readdir } from "node:fs/promises";
import { join } from "node:path";
import { tmpdir } from "node:os";
import { pipeline } from "node:stream/promises";
import { Readable, Transform } from "node:stream";
import JSZip from "jszip";
const directory = join(tmpdir(), "canvas-download-zips-v1");
export async function writeDownloadZip(zip: JSZip, id: string, progress: (percent: number) => void) {
await mkdir(directory, { recursive: true, mode: 0o700 });
// Remove abandoned files after restart as well as expired in-process tasks.
for (const name of await readdir(directory)) {
const path = join(directory, name);
if ((await stat(path)).mtimeMs < Date.now() - 60 * 60_000) await unlink(path).catch(() => undefined);
}
const path = join(directory, `${id}.zip`);
try {
const disk = await statfs(directory);
const budget = disk.bavail * disk.bsize - 1024 * 1024 * 1024;
let written = 0;
const guard = new Transform({ transform(chunk, _encoding, callback) {
written += chunk.length;
callback(written > budget ? new Error("服务器下载空间不足,请减少本次图片数量") : null, chunk);
} });
await pipeline(zip.generateNodeStream({ type: "nodebuffer", streamFiles: true, compression: "STORE" },
metadata => progress(metadata.percent)), guard, createWriteStream(path, { mode: 0o600 }));
return { path, bytes: (await stat(path)).size };
} catch (error) { await unlink(path).catch(() => undefined); throw error; }
}
export function readDownloadZip(path: string) {
return Readable.toWeb(createReadStream(path)) as ReadableStream<Uint8Array>;
}
export function streamDownloadZip(zip: JSZip) {
return Readable.toWeb(zip.generateNodeStream({ type: "nodebuffer", streamFiles: true, compression: "STORE" }) as Readable) as ReadableStream<Uint8Array>;
}
export async function removeDownloadZip(path: string) { await unlink(path).catch(() => undefined); }