diff --git a/client.ts b/client.ts index ba06b1c..f15de52 100644 --- a/client.ts +++ b/client.ts @@ -10,11 +10,13 @@ export class Client { cookies; host; ua; - constructor(cfg: Config) { + signal; + constructor(cfg: Config, signal?: AbortSignal) { this.cookies = cfg.cookies; this.ua = cfg.ua || "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/108.0.0.0 Safari/537.36"; this.host = cfg.ex ? "exhentai.org" : "e-hentai.org"; + this.signal = signal; } get( url: string | Request | URL, @@ -81,6 +83,9 @@ export class Client { } else { d.headers = headers; } + if (!d.signal && this.signal) { + d.signal = this.signal; + } return fetch(url, d); } } diff --git a/signal_handler.ts b/signal_handler.ts index e58972f..b72e859 100644 --- a/signal_handler.ts +++ b/signal_handler.ts @@ -1,7 +1,22 @@ import { TaskManager } from "./task_manager.ts"; export function add_exit_handler(m: TaskManager) { - const handler = () => { + let first_aborted = true; + let ignore_signal = false; + const handler = async () => { + if (ignore_signal) return; + if (first_aborted) { + m.abort(); + console.log( + "Already abort all tasks. Please wait for a while. You can press Ctrl + C again to force abort.", + ); + first_aborted = false; + } else { + m.force_abort(); + ignore_signal = true; + return; + } + await m.waiting_unfinished_task(); m.close(); }; Deno.addSignalListener("SIGINT", handler); diff --git a/task_manager.ts b/task_manager.ts index 1773643..aed226f 100644 --- a/task_manager.ts +++ b/task_manager.ts @@ -17,9 +17,13 @@ export class TaskManager { db; running_tasks: Map>; max_task_count; + #abort; + #force_abort; constructor(cfg: Config) { this.cfg = cfg; - this.client = new Client(cfg); + this.#abort = new AbortController(); + this.#force_abort = new AbortController(); + this.client = new Client(cfg, this.#force_abort.signal); this.db = new EhDb(cfg.db_path || cfg.base); this.running_tasks = new Map(); this.max_task_count = cfg.max_task_count; @@ -28,6 +32,12 @@ export class TaskManager { #check_closed() { if (this.#closed) throw new AlreadyClosedError(); } + abort(reason?: unknown) { + this.#abort.abort(reason); + } + get aborted() { + return this.#abort.signal.aborted; + } async add_download_task(gid: number, token: string) { this.#check_closed(); const otask = await this.db.check_download_task(gid, token); @@ -74,7 +84,7 @@ export class TaskManager { removed_task.push(id); await this.db.delete_task(status.value); } else if (status.status == PromiseStatus.Rejected) { - if (status.reason) console.log(status.reason); + if (status.reason && !this.aborted) console.log(status.reason); removed_task.push(id); } } @@ -87,7 +97,14 @@ export class TaskManager { this.#closed = true; this.db.close(); } + force_abort(reason?: unknown) { + this.#force_abort.abort(reason); + } + get force_aborted() { + return this.#force_abort.signal.aborted; + } async run() { + if (this.aborted || this.force_aborted) throw new AlreadyClosedError(); this.#check_closed(); while (1) { await this.check_running_tasks(); @@ -116,8 +133,22 @@ export class TaskManager { if (task.type == TaskType.Download) { this.running_tasks.set( task.id, - download_task(task, this.client, this.db, this.cfg), + download_task( + task, + this.client, + this.db, + this.cfg, + this.#abort.signal, + this.#force_abort.signal, + ), ); } } + async waiting_unfinished_task() { + while (1) { + await this.check_running_tasks(); + if (this.running_tasks.size == 0) break; + await sleep(10); + } + } } diff --git a/tasks/download.ts b/tasks/download.ts index 62cbf84..28bfea9 100644 --- a/tasks/download.ts +++ b/tasks/download.ts @@ -14,12 +14,16 @@ import { join, resolve } from "std/path/mod.ts"; import { exists } from "std/fs/exists.ts"; class DownloadManager { + #abort: AbortSignal; + #force_abort: AbortSignal; #max_download_count; #running_tasks: Promise[]; #has_failed_task = false; - constructor(cfg: Config) { + constructor(cfg: Config, abort: AbortSignal, force_abort: AbortSignal) { this.#max_download_count = cfg.max_download_img_count; this.#running_tasks = []; + this.#abort = abort; + this.#force_abort = force_abort; } async #check_tasks() { this.#running_tasks = await asyncFilter( @@ -27,7 +31,7 @@ class DownloadManager { async (t) => { const s = await promiseState(t); if (s.status === PromiseStatus.Rejected) { - console.log(s.reason); + if (!this.#force_abort.aborted) console.log(s.reason); this.#has_failed_task = true; } return s.status === PromiseStatus.Pending; @@ -36,6 +40,7 @@ class DownloadManager { } async add_new_task(f: () => Promise) { while (1) { + if (this.#abort.aborted) break; await this.#check_tasks(); if (this.#running_tasks.length < this.#max_download_count) { this.#running_tasks.push(f()); @@ -61,6 +66,8 @@ export async function download_task( client: Client, db: EhDb, cfg: Config, + abort: AbortSignal, + force_abort: AbortSignal, ) { console.log("Started to download gallery", task.gid); const gdatas = await client.fetchGalleryMetadataByAPI([ @@ -75,10 +82,11 @@ export async function download_task( await db.add_gtag(task.gid, new Set(gdata.tags)); const base_path = join(cfg.base, task.gid.toString()); await sure_dir(base_path); - const m = new DownloadManager(cfg); + const m = new DownloadManager(cfg, abort, force_abort); if (cfg.mpv) { const mpv = await client.fetchMPVPage(task.gid, task.token); for (const i of mpv.imagelist) { + if (abort.aborted) break; await m.add_new_task(async () => { const ofiles = db.get_files(task.gid, i.page_token); if (ofiles.length) { @@ -97,6 +105,9 @@ export async function download_task( function try_load(a: number) { if (a >= cfg.max_retry_count) reject(errors); i.load().then(resolve).catch((e) => { + if (force_abort.aborted) { + throw Error("aborted."); + } errors.push(e); try_load(a + 1); }); @@ -147,6 +158,9 @@ export async function download_task( reject(errors); } download().then(resolve).catch((e) => { + if (force_abort.aborted) { + throw Error("aborted."); + } errors.push(e); try_download(a + 1); }); @@ -165,5 +179,6 @@ export async function download_task( } } await m.join(); + if (abort.aborted || force_abort.aborted) throw Error("aborted"); return task; }