From 5d75c211a7b174b0ca66d3d2d58dcef73856fb34 Mon Sep 17 00:00:00 2001
From: lifegpc <g1710431395@gmail.com>
Date: Wed, 31 May 2023 09:44:14 +0800
Subject: [PATCH] Update

---
 routes/api/task.ts | 35 ++++++++++++++++++++---------------
 task.ts            | 19 +++++++++++++++++++
 task_manager.ts    | 16 ++++++++++++++--
 utils.ts           |  9 +++++++++
 4 files changed, 62 insertions(+), 17 deletions(-)

diff --git a/routes/api/task.ts b/routes/api/task.ts
index 703d03a..e2ef666 100644
--- a/routes/api/task.ts
+++ b/routes/api/task.ts
@@ -1,18 +1,27 @@
 import { Handlers } from "$fresh/server.ts";
 import { get_task_manager } from "../../server.ts";
-import { Task } from "../../task.ts";
+import { Task, TaskProgress } from "../../task.ts";
+import { DiscriminatedUnion } from "../../utils.ts";
+
+type EventMap = {
+    close: Record<PropertyKey, never>;
+    new_download_task: { gid: number; token: string };
+};
+
+type EventData = DiscriminatedUnion<"type", EventMap>;
 
 export const handler: Handlers<Task[]> = {
     GET(req, _ctx) {
         const t = get_task_manager();
         const { socket, response } = Deno.upgradeWebSocket(req);
-        const task_handle = (e: CustomEvent<Task>) => {
+        const handle = (e: CustomEvent<Task | TaskProgress>) => {
             socket.send(JSON.stringify({ type: e.type, detail: e.detail }));
         };
         const removeListener = () => {
-            t.removeEventListener("new_task", task_handle);
-            t.removeEventListener("task_started", task_handle);
-            t.removeEventListener("task_finished", task_handle);
+            t.removeEventListener("new_task", handle);
+            t.removeEventListener("task_started", handle);
+            t.removeEventListener("task_finished", handle);
+            t.removeEventListener("task_progress", handle);
         };
         socket.onclose = () => {
             removeListener();
@@ -23,25 +32,21 @@ export const handler: Handlers<Task[]> = {
         };
         socket.onmessage = (e) => {
             try {
-                const d = JSON.parse(e.data);
+                const d: EventData = JSON.parse(e.data);
                 if (d.type == "close") {
                     socket.close();
                 } else if (d.type == "new_download_task") {
-                    const gid = d.gid;
-                    const token = d.token;
-                    if (typeof gid !== "number" || typeof token !== "string") {
-                        return;
-                    }
-                    t.add_download_task(gid, token);
+                    t.add_download_task(d.gid, d.token);
                 }
             } catch (_) {
                 null;
             }
         };
         socket.onopen = () => {
-            t.addEventListener("new_task", task_handle);
-            t.addEventListener("task_started", task_handle);
-            t.addEventListener("task_finished", task_handle);
+            t.addEventListener("new_task", handle);
+            t.addEventListener("task_started", handle);
+            t.addEventListener("task_finished", handle);
+            t.addEventListener("task_progress", handle);
         };
         return response;
     },
diff --git a/task.ts b/task.ts
index 89eec5b..6edfe4f 100644
--- a/task.ts
+++ b/task.ts
@@ -1,3 +1,5 @@
+import { DiscriminatedUnion } from "./utils.ts";
+
 export enum TaskType {
     Download,
     ExportZip,
@@ -11,3 +13,20 @@ export type Task = {
     pid: number;
     details: string | null;
 };
+
+export type TaskDownloadProgess = {
+    downloaded_page: number;
+    total_page: number;
+};
+
+export type TaskExportZipProgress = {
+    added_page: number;
+    total_page: number;
+};
+
+export type TaskProgressType = {
+    [TaskType.Download]: TaskDownloadProgess;
+    [TaskType.ExportZip]: TaskExportZipProgress;
+};
+
+export type TaskProgress = DiscriminatedUnion<"type", TaskProgressType>;
diff --git a/task_manager.ts b/task_manager.ts
index d76d665..880fab9 100644
--- a/task_manager.ts
+++ b/task_manager.ts
@@ -3,14 +3,19 @@ import { Config } from "./config.ts";
 import { EhDb } from "./db.ts";
 import { check_running } from "./pid_check.ts";
 import { add_exit_handler } from "./signal_handler.ts";
-import { Task, TaskType } from "./task.ts";
+import { Task, TaskProgress, TaskType } from "./task.ts";
 import { download_task } from "./tasks/download.ts";
 import {
     DEFAULT_EXPORT_ZIP_CONFIG,
     export_zip,
     ExportZipConfig,
 } from "./tasks/export_zip.ts";
-import { promiseState, PromiseStatus, sleep } from "./utils.ts";
+import {
+    DiscriminatedUnion,
+    promiseState,
+    PromiseStatus,
+    sleep,
+} from "./utils.ts";
 
 export class AlreadyClosedError extends Error {
 }
@@ -19,8 +24,15 @@ type EventMap = {
     new_task: Task;
     task_started: Task;
     task_finished: Task;
+    task_progress: TaskProgress;
 };
 
+type Detail<T extends Record<PropertyKey, unknown>> = {
+    [P in keyof T]: { detail: T[P] };
+};
+
+export type TaskEventData = DiscriminatedUnion<"type", Detail<EventMap>>;
+
 export class TaskManager extends EventTarget {
     #closed = false;
     cfg;
diff --git a/utils.ts b/utils.ts
index 963d954..a5033f7 100644
--- a/utils.ts
+++ b/utils.ts
@@ -134,3 +134,12 @@ export function filterFilename(p: string, maxLength = 256) {
     }
     return p;
 }
+
+export type DiscriminatedUnion<
+    K extends PropertyKey,
+    T extends Record<PropertyKey, unknown>,
+> = {
+    [P in keyof T]: ({ [Q in K]: P } & T[P]) extends infer U
+        ? { [Q in keyof U]: U[Q] }
+        : never;
+}[keyof T];