536 lines
20 KiB
TypeScript
536 lines
20 KiB
TypeScript
/*
|
|
* The scrape engine: queue, walk, pacing, budgets, backoff, resume.
|
|
* Rules: docs/media-scraper.md.
|
|
*/
|
|
|
|
import { copyFileSync, mkdirSync, rmSync } from "node:fs";
|
|
import { dirname, join } from "node:path";
|
|
import {
|
|
downloadGapMs,
|
|
downloadOne,
|
|
downloadPreview,
|
|
gifPoster,
|
|
previewGapMs,
|
|
probeDimensions,
|
|
remoteFrame,
|
|
videoPoster,
|
|
type DownloadOptions
|
|
} from "./download.js";
|
|
import { resolveEmbedPreview } from "./embed-preview.js";
|
|
import {
|
|
createScrapeDirs,
|
|
listScrapes,
|
|
readScrapeMeta,
|
|
ScrapeStore,
|
|
uniqueScrapeId,
|
|
writeScrapeMeta
|
|
} from "./store.js";
|
|
import {
|
|
emptyMediaCounts,
|
|
MEDIA_KINDS,
|
|
type BackendName,
|
|
type CoreMessage,
|
|
type ScrapeMeta,
|
|
type ScrapeMode,
|
|
type ScrapeState
|
|
} from "./types.js";
|
|
|
|
export interface FetchPage {
|
|
messages: CoreMessage[];
|
|
/** Set by a backend that saw a rate-limit hint. The engine sleeps before the next page. */
|
|
rateLimitedMs?: number;
|
|
}
|
|
|
|
export interface ScrapeBackend {
|
|
readonly name: BackendName;
|
|
fetchPage(channelId: string, before: string | null, limit: number): Promise<FetchPage>;
|
|
}
|
|
|
|
export interface ScrapeEngineOptions {
|
|
rootDir: string;
|
|
pageSize?: number;
|
|
pageMinMs?: number;
|
|
pageMaxMs?: number;
|
|
pagesPerHour?: number;
|
|
pagesPerDay?: number;
|
|
maxRetries?: number;
|
|
sleep?: (ms: number) => Promise<void>;
|
|
random?: () => number;
|
|
now?: () => number;
|
|
log?: (line: string) => void;
|
|
/** Download settings for `mode: "download"`. */
|
|
download?: DownloadOptions;
|
|
/** Cache Discord-hosted image previews after every walk. Default true. */
|
|
previews?: boolean;
|
|
/** Read dimensions and duration from a downloaded file. Defaults to ffprobe. */
|
|
probe?: (
|
|
path: string
|
|
) => Promise<{ width: number | null; height: number | null; durationSecs: number | null } | null>;
|
|
/** Called when a scrape fails. */
|
|
onError?: (scrapeId: string, message: string) => void;
|
|
}
|
|
|
|
interface QueuedJob {
|
|
meta: ScrapeMeta;
|
|
dir: string;
|
|
backend: ScrapeBackend;
|
|
}
|
|
|
|
const defaultSleep = (ms: number) => new Promise<void>(resolve => setTimeout(resolve, ms));
|
|
|
|
function oldestId(messages: CoreMessage[]): string | null {
|
|
let oldest: string | null = null;
|
|
for (const message of messages) {
|
|
if (oldest === null || BigInt(message.id) < BigInt(oldest)) oldest = message.id;
|
|
}
|
|
return oldest;
|
|
}
|
|
|
|
export class ScrapeEngine {
|
|
private readonly queue: QueuedJob[] = [];
|
|
private readonly active = new Map<string, ScrapeMeta>();
|
|
private readonly stopRequested = new Set<string>();
|
|
private readonly pageTimes: number[] = [];
|
|
private processingPromise: Promise<void> | null = null;
|
|
|
|
constructor(private readonly options: ScrapeEngineOptions) {}
|
|
|
|
/** Create a scrape and queue the walk. Returns immediately. */
|
|
async start(input: {
|
|
channelId: string;
|
|
guildId: string | null;
|
|
mode: ScrapeMode;
|
|
backend: ScrapeBackend;
|
|
downloadDir?: string | null;
|
|
afterId?: string | null;
|
|
beforeId?: string | null;
|
|
fromDate?: string | null;
|
|
toDate?: string | null;
|
|
estimatedMessages?: number | null;
|
|
}): Promise<ScrapeMeta> {
|
|
const now = this.now();
|
|
const id = uniqueScrapeId(this.options.rootDir, input.channelId, new Date(now));
|
|
const dir = createScrapeDirs(this.options.rootDir, id);
|
|
const timestamp = new Date(now).toISOString();
|
|
|
|
const meta: ScrapeMeta = {
|
|
id,
|
|
channelId: input.channelId,
|
|
guildId: input.guildId,
|
|
backend: input.backend.name,
|
|
mode: input.mode,
|
|
state: "queued",
|
|
cursor: input.beforeId ?? null,
|
|
createdAt: timestamp,
|
|
updatedAt: timestamp,
|
|
pages: 0,
|
|
messages: 0,
|
|
media: emptyMediaCounts(),
|
|
errors: 0,
|
|
lastError: null,
|
|
downloadDir: input.downloadDir ?? null,
|
|
afterId: input.afterId ?? null,
|
|
beforeId: input.beforeId ?? null,
|
|
fromDate: input.fromDate ?? null,
|
|
toDate: input.toDate ?? null,
|
|
estimatedMessages: input.estimatedMessages ?? null
|
|
};
|
|
|
|
writeScrapeMeta(dir, meta);
|
|
this.active.set(id, meta);
|
|
this.queue.push({ meta, dir, backend: input.backend });
|
|
void this.processQueue();
|
|
|
|
return { ...meta };
|
|
}
|
|
|
|
/** Run queued jobs until the queue is empty. Tests await this; the daemon lets it run. */
|
|
processQueue(): Promise<void> {
|
|
if (!this.processingPromise) {
|
|
this.processingPromise = this.drainQueue().finally(() => {
|
|
this.processingPromise = null;
|
|
});
|
|
}
|
|
return this.processingPromise;
|
|
}
|
|
|
|
private async drainQueue(): Promise<void> {
|
|
while (this.queue.length > 0) {
|
|
const job = this.queue.shift();
|
|
if (!job) break;
|
|
await this.runJob(job);
|
|
}
|
|
}
|
|
|
|
get(id: string): ScrapeMeta | null {
|
|
const active = this.active.get(id);
|
|
if (active) return { ...active };
|
|
return readScrapeMeta(join(this.options.rootDir, id));
|
|
}
|
|
|
|
/** Every scrape on disk, merged with the jobs this engine still holds in memory. */
|
|
list(): ScrapeMeta[] {
|
|
const merged = new Map<string, ScrapeMeta>();
|
|
for (const meta of listScrapes(this.options.rootDir)) merged.set(meta.id, meta);
|
|
for (const [id, meta] of this.active) merged.set(id, meta);
|
|
|
|
return [...merged.values()].sort((a, b) => (a.createdAt < b.createdAt ? 1 : -1));
|
|
}
|
|
|
|
/** Ask a running or queued scrape to stop. It pauses after the current page. */
|
|
stop(id: string): void {
|
|
this.stopRequested.add(id);
|
|
}
|
|
|
|
/** Cache previews for rows a finished or paused scrape never reached. Returns the count. */
|
|
async backfillPreviews(scrapeId: string): Promise<number> {
|
|
const meta = this.active.get(scrapeId) ?? readScrapeMeta(join(this.options.rootDir, scrapeId));
|
|
if (!meta || meta.state === "running" || meta.state === "queued") return 0;
|
|
|
|
const dir = join(this.options.rootDir, scrapeId);
|
|
const store = new ScrapeStore(dir);
|
|
const before = meta.previewed ?? 0;
|
|
|
|
try {
|
|
await this.runPreviews(store, dir, meta);
|
|
} finally {
|
|
store.close();
|
|
}
|
|
|
|
return (meta.previewed ?? 0) - before;
|
|
}
|
|
|
|
/** Persist a state change made outside a job, such as retention marking. */
|
|
setState(id: string, state: ScrapeState): ScrapeMeta | null {
|
|
const meta = this.active.get(id) ?? readScrapeMeta(join(this.options.rootDir, id));
|
|
if (!meta) return null;
|
|
|
|
meta.state = state;
|
|
meta.updatedAt = new Date(this.now()).toISOString();
|
|
writeScrapeMeta(join(this.options.rootDir, id), meta);
|
|
this.active.set(id, meta);
|
|
return { ...meta };
|
|
}
|
|
|
|
/** Drop a deleted scrape from the in-memory list. */
|
|
forget(id: string): void {
|
|
this.active.delete(id);
|
|
this.stopRequested.delete(id);
|
|
}
|
|
|
|
private async runJob(job: QueuedJob): Promise<void> {
|
|
const { meta, dir, backend } = job;
|
|
const store = new ScrapeStore(dir);
|
|
const pageSize = this.options.pageSize ?? 100;
|
|
|
|
try {
|
|
if (this.stopRequested.has(meta.id)) {
|
|
meta.state = "paused";
|
|
return;
|
|
}
|
|
|
|
meta.state = "running";
|
|
this.persist(dir, meta);
|
|
|
|
for (;;) {
|
|
if (this.stopRequested.has(meta.id)) {
|
|
meta.state = "paused";
|
|
break;
|
|
}
|
|
|
|
if (!this.allowPage()) {
|
|
meta.state = "paused";
|
|
this.log(`${meta.id}: page budget reached, paused`);
|
|
break;
|
|
}
|
|
|
|
const page = await this.fetchWithBackoff(meta, backend);
|
|
if (page.rateLimitedMs) await this.sleep(page.rateLimitedMs);
|
|
|
|
if (page.messages.length === 0) {
|
|
meta.state = "done";
|
|
break;
|
|
}
|
|
|
|
// The from-date is a lower ID boundary: trim anything older and finish.
|
|
const kept = meta.afterId
|
|
? page.messages.filter(message => BigInt(message.id) > BigInt(meta.afterId as string))
|
|
: page.messages;
|
|
|
|
if (kept.length === 0) {
|
|
meta.state = "done";
|
|
break;
|
|
}
|
|
|
|
const added = store.insertPage(kept);
|
|
for (const kind of MEDIA_KINDS) meta.media[kind] += added[kind];
|
|
meta.pages++;
|
|
meta.messages += kept.length;
|
|
meta.cursor = oldestId(kept) ?? meta.cursor;
|
|
this.pageTimes.push(this.now());
|
|
this.persist(dir, meta);
|
|
|
|
if (kept.length < page.messages.length || page.messages.length < pageSize) {
|
|
meta.state = "done";
|
|
break;
|
|
}
|
|
|
|
await this.sleep(this.gapMs());
|
|
}
|
|
|
|
const wasPaused = meta.state === "paused";
|
|
|
|
if (
|
|
(meta.state === "done" || wasPaused) &&
|
|
(meta.mode === "download" || this.options.previews !== false)
|
|
) {
|
|
// Keep the job running until post-processing finishes, so a client never
|
|
// treats a half-archived scrape as complete. A paused walk still caches
|
|
// previews and downloads for the pages it did collect.
|
|
meta.state = "running";
|
|
this.persist(dir, meta);
|
|
|
|
if (meta.mode === "download") {
|
|
await this.runDownloads(store, dir, meta);
|
|
}
|
|
|
|
if (this.options.previews !== false && !this.stopRequested.has(meta.id)) {
|
|
await this.runPreviews(store, dir, meta);
|
|
}
|
|
|
|
if (meta.state === "running") {
|
|
meta.state = wasPaused ? "paused" : "done";
|
|
}
|
|
this.persist(dir, meta);
|
|
}
|
|
} catch (error) {
|
|
meta.state = "failed";
|
|
meta.lastError = error instanceof Error ? error.message : String(error);
|
|
this.log(`${meta.id}: failed: ${meta.lastError}`);
|
|
this.options.onError?.(meta.id, meta.lastError);
|
|
} finally {
|
|
if ((meta.state === "done" || meta.state === "failed") && !meta.finishedAt) {
|
|
meta.finishedAt = new Date(this.now()).toISOString();
|
|
}
|
|
this.persist(dir, meta);
|
|
store.close();
|
|
this.active.set(meta.id, meta);
|
|
}
|
|
}
|
|
|
|
/** Fetch attachment bytes for a finished walk. Serial, hash-addressed, resumable. */
|
|
private async runDownloads(store: ScrapeStore, dir: string, meta: ScrapeMeta): Promise<void> {
|
|
const customDir = meta.downloadDir ?? null;
|
|
const options: DownloadOptions = {
|
|
...(this.options.download ?? {}),
|
|
destinationDir: customDir
|
|
};
|
|
const targetDir = customDir ? join(customDir, meta.id, "media") : join(dir, "media");
|
|
mkdirSync(targetDir, { recursive: true });
|
|
const pending = store.pendingDownloads();
|
|
|
|
for (const row of pending) {
|
|
if (this.stopRequested.has(meta.id)) {
|
|
meta.state = "paused";
|
|
return;
|
|
}
|
|
|
|
const outcome = await downloadOne(row, targetDir, {
|
|
...options,
|
|
log: line => this.log(`${meta.id}: ${line}`)
|
|
});
|
|
|
|
if (outcome.state === "downloaded" && outcome.path && outcome.hash) {
|
|
const filePath = customDir ? join(targetDir, outcome.path) : join("media", outcome.path);
|
|
store.markDownloaded(row.id, {
|
|
path: filePath,
|
|
hash: outcome.hash,
|
|
bytes: outcome.bytes ?? 0
|
|
});
|
|
meta.downloaded = (meta.downloaded ?? 0) + 1;
|
|
|
|
if (row.width === null || row.height === null || row.duration_secs === null) {
|
|
const info = await (this.options.probe ?? probeDimensions)(
|
|
join(targetDir, outcome.path)
|
|
);
|
|
if (info) store.updateMediaInfo(row.id, info);
|
|
}
|
|
|
|
if (row.kind === "video" && !row.preview_hash) {
|
|
const poster = await videoPoster(join(targetDir, outcome.path), join(dir, "thumbs"));
|
|
if (poster) {
|
|
store.markPreview(row.id, poster);
|
|
meta.previewed = (meta.previewed ?? 0) + 1;
|
|
}
|
|
}
|
|
} else {
|
|
store.markDownloadError(row.id, outcome.error ?? outcome.state);
|
|
if (outcome.state === "failed") meta.downloadErrors = (meta.downloadErrors ?? 0) + 1;
|
|
}
|
|
|
|
this.persist(dir, meta);
|
|
await this.sleep(downloadGapMs(options));
|
|
}
|
|
}
|
|
|
|
/** Cache small previews. Embed thumbnails are cached in both modes. */
|
|
private async runPreviews(store: ScrapeStore, dir: string, meta: ScrapeMeta): Promise<void> {
|
|
const options = {
|
|
...(this.options.download ?? {}),
|
|
embedPreviews: true,
|
|
destinationDir: meta.downloadDir ?? null
|
|
};
|
|
const pending = store.pendingPreviews(5000, true);
|
|
|
|
for (const row of pending) {
|
|
if (this.stopRequested.has(meta.id)) {
|
|
meta.state = "paused";
|
|
return;
|
|
}
|
|
|
|
// Audio has no image proxy: render a waveform strip with ffmpeg.
|
|
if (row.kind === "audio" && row.source !== "embed") {
|
|
const waveform = await remoteFrame(row.url, join(dir, "thumbs"), { waveform: true });
|
|
if (waveform) {
|
|
store.markPreview(row.id, {
|
|
hash: waveform.hash,
|
|
bytes: waveform.bytes,
|
|
ext: "webp"
|
|
});
|
|
meta.previewed = (meta.previewed ?? 0) + 1;
|
|
} else {
|
|
store.markPreviewError(row.id, "waveform failed");
|
|
}
|
|
this.persist(dir, meta);
|
|
await this.sleep(previewGapMs(options));
|
|
continue;
|
|
}
|
|
|
|
// Some providers (fxtwitter, vxReddit, ...) ship no thumbnail in Discord's data;
|
|
// resolve one from the provider or the page before caching it.
|
|
let previewRow = row;
|
|
if (row.source === "embed" && !row.preview_url) {
|
|
const resolved = await resolveEmbedPreview(row, options);
|
|
if (!resolved) {
|
|
store.markPreviewError(row.id, "no preview source");
|
|
this.persist(dir, meta);
|
|
await this.sleep(previewGapMs(options));
|
|
continue;
|
|
}
|
|
previewRow = { ...row, preview_url: resolved };
|
|
}
|
|
|
|
const outcome = await downloadPreview(previewRow, join(dir, "thumbs"), options);
|
|
|
|
if (outcome.state === "downloaded" && outcome.hash) {
|
|
let hash = outcome.hash;
|
|
let bytes = outcome.bytes ?? 0;
|
|
let ext = outcome.ext ?? "webp";
|
|
|
|
// Animated GIFs are kept static in the grid; hovering plays the original.
|
|
if (ext === "gif") {
|
|
const poster = await gifPoster(join(dir, "thumbs", `${hash}.gif`), join(dir, "thumbs"));
|
|
if (poster) {
|
|
rmSync(join(dir, "thumbs", `${hash}.gif`), { force: true });
|
|
hash = poster.hash;
|
|
bytes = poster.bytes;
|
|
ext = "webp";
|
|
}
|
|
}
|
|
|
|
store.markPreview(row.id, { hash, bytes, ext });
|
|
meta.previewed = (meta.previewed ?? 0) + 1;
|
|
|
|
// Embed thumbnails also land in the custom destination, so the download has a
|
|
// viewable poster for embeds while the media tab keeps its local copy.
|
|
if (options.destinationDir && row.source === "embed") {
|
|
const target = join(options.destinationDir, meta.id, "thumbnails", `${hash}.${ext}`);
|
|
mkdirSync(dirname(target), { recursive: true });
|
|
copyFileSync(join(dir, "thumbs", `${hash}.${ext}`), target);
|
|
}
|
|
} else if (
|
|
outcome.state === "skipped" &&
|
|
outcome.error === "too large" &&
|
|
/\.gif(\?|$)/i.test(previewRow.preview_url ?? previewRow.url)
|
|
) {
|
|
// Oversized animated GIF: let ffmpeg read the stream and keep one frame.
|
|
const frame = await remoteFrame(previewRow.preview_url ?? previewRow.url, join(dir, "thumbs"));
|
|
if (frame) {
|
|
store.markPreview(row.id, { hash: frame.hash, bytes: frame.bytes, ext: "webp" });
|
|
meta.previewed = (meta.previewed ?? 0) + 1;
|
|
} else {
|
|
store.markPreviewError(row.id, "too large");
|
|
}
|
|
} else if (outcome.state === "failed") {
|
|
// Transient network failure: leave the row pending for the next pass.
|
|
this.log(`${meta.id}: preview ${row.id} failed: ${outcome.error ?? "unknown"}`);
|
|
} else {
|
|
store.markPreviewError(row.id, outcome.error ?? outcome.state);
|
|
}
|
|
|
|
this.persist(dir, meta);
|
|
await this.sleep(previewGapMs(options));
|
|
}
|
|
}
|
|
|
|
private async fetchWithBackoff(meta: ScrapeMeta, backend: ScrapeBackend): Promise<FetchPage> {
|
|
const pageSize = this.options.pageSize ?? 100;
|
|
const maxRetries = this.options.maxRetries ?? 8;
|
|
|
|
let attempt = 0;
|
|
for (;;) {
|
|
try {
|
|
return await backend.fetchPage(meta.channelId, meta.cursor, pageSize);
|
|
} catch (error) {
|
|
meta.errors++;
|
|
if (attempt >= maxRetries) throw error;
|
|
|
|
const retryAfter = (error as { retryAfterMs?: number }).retryAfterMs;
|
|
const delay = retryAfter ?? Math.pow(2, attempt) * 1000 + 1000;
|
|
this.log(`${meta.id}: page failed, retrying in ${delay} ms`);
|
|
await this.sleep(delay);
|
|
attempt++;
|
|
}
|
|
}
|
|
}
|
|
|
|
private allowPage(): boolean {
|
|
const now = this.now();
|
|
while (this.pageTimes.length > 0 && this.pageTimes[0] < now - 86_400_000) {
|
|
this.pageTimes.shift();
|
|
}
|
|
|
|
const lastHour = this.pageTimes.filter(time => time > now - 3_600_000).length;
|
|
if (lastHour >= (this.options.pagesPerHour ?? 300)) return false;
|
|
if (this.pageTimes.length >= (this.options.pagesPerDay ?? 3000)) return false;
|
|
return true;
|
|
}
|
|
|
|
private gapMs(): number {
|
|
const min = this.options.pageMinMs ?? 1200;
|
|
const max = Math.max(min, this.options.pageMaxMs ?? 3500);
|
|
return min + Math.floor(this.random() * (max - min + 1));
|
|
}
|
|
|
|
private persist(dir: string, meta: ScrapeMeta): void {
|
|
meta.updatedAt = new Date(this.now()).toISOString();
|
|
writeScrapeMeta(dir, meta);
|
|
}
|
|
|
|
private now(): number {
|
|
return (this.options.now ?? Date.now)();
|
|
}
|
|
|
|
private random(): number {
|
|
return (this.options.random ?? Math.random)();
|
|
}
|
|
|
|
private sleep(ms: number): Promise<void> {
|
|
return (this.options.sleep ?? defaultSleep)(ms);
|
|
}
|
|
|
|
private log(line: string): void {
|
|
this.options.log?.(line);
|
|
}
|
|
}
|