feat: TransformStream implement

WindRunnerMax · WindRunnerMax · commit cd00cfb462ce · 2024-08-18T00:02:22.000+08:00
diff --git a/packages/webrtc/client/app/modal.tsx b/packages/webrtc/client/app/modal.tsx
@@ -17,9 +17,9 @@ import {
   FILE_STATE,
   ID_SIZE,
   STEAM_TYPE,
-  destructureChunk,
+  deserializeChunk,
   getMaxMessageSize,
-  getNextChunk,
+  serializeNextChunk,
   sendChunkMessage,
 } from "../utils/binary";
 import { WorkerEvent } from "../worker/event";
@@ -96,7 +96,7 @@ export const TransferModal: FC<{
         const { id, series, total } = data;
         const progress = Math.floor((series / total) * 100);
         updateFileProgress(id, progress);
-        const nextChunk = getNextChunk(rtc, id, series);
+        const nextChunk = serializeNextChunk(rtc, id, series);
         // 通知 接收方 发送块数据
         sendChunkMessage(rtc, nextChunk);
       } else if (data.key === MESSAGE_TYPE.FILE_FINISH) {
@@ -108,7 +108,7 @@ export const TransferModal: FC<{
     } else {
       // Binary - 接收 发送方 ArrayBuffer 数据
       const blob = event.data;
-      destructureChunk(blob).then(({ id, series, data }) => {
+      deserializeChunk(blob).then(({ id, series, data }) => {
         const state = FILE_STATE.get(id);
         if (!state) return void 0;
         const { size, total } = state;
@@ -121,7 +121,7 @@ export const TransferModal: FC<{
         } else {
           // 数据块序列号 [0, TOTAL)
           if (stream) {
-            WorkerEvent.post(blob);
+            WorkerEvent.post(id, data);
           } else {
             // 在内存中存储块数据
             const mapper = FILE_MAPPER.get(id) || [];
diff --git a/packages/webrtc/client/utils/binary.ts b/packages/webrtc/client/utils/binary.ts
@@ -23,7 +23,7 @@ export const getMaxMessageSize = (
   return maxSize - (ID_SIZE + CHUNK_SIZE);
 };
 
-export const getNextChunk = (
+export const serializeNextChunk = (
   instance: React.MutableRefObject<WebRTCApi | null>,
   id: string,
   series: number
@@ -75,7 +75,7 @@ export const sendChunkMessage = async (
   !isSending && start(rtc);
 };
 
-export const destructureChunk = async (chunk: BufferType) => {
+export const deserializeChunk = async (chunk: BufferType) => {
   const buffer = chunk instanceof Blob ? await chunk.arrayBuffer() : chunk;
   const id = new Uint8Array(buffer.slice(0, ID_SIZE));
   const series = new Uint8Array(buffer.slice(ID_SIZE, ID_SIZE + CHUNK_SIZE));
diff --git a/packages/webrtc/client/worker/event.ts b/packages/webrtc/client/worker/event.ts
@@ -1,8 +1,10 @@
-import type { BufferType } from "../../types/client";
-import { HEADER_KEY } from "../../types/worker";
+import type { MessageType } from "../../types/worker";
+import { HEADER_KEY, MESSAGE_TYPE } from "../../types/worker";
 
 export class WorkerEvent {
+  public static channel: MessageChannel | null = null;
   public static worker: ServiceWorkerRegistration | null = null;
+  public static writer: Map<string, WritableStreamDefaultWriter<Uint8Array>> = new Map();
 
   public static async register(): Promise<ServiceWorkerRegistration> {
     try {
@@ -28,6 +30,29 @@ export class WorkerEvent {
   }
 
   public static start(fileId: string, fileName: string, fileSize: number, fileTotal: number) {
+    if (!WorkerEvent.channel) {
+      WorkerEvent.channel = new MessageChannel();
+      WorkerEvent.channel.port1.onmessage = event => {
+        console.log("WorkerEvent", event.data);
+      };
+      WorkerEvent.worker?.active?.postMessage({ type: MESSAGE_TYPE.INIT_CHANNEL }, [
+        WorkerEvent.channel.port2,
+      ]);
+    }
+    const ts = new TransformStream({
+      transform(chunk, controller) {
+        controller.enqueue(chunk);
+      },
+    });
+    WorkerEvent.channel.port1.postMessage(
+      <MessageType>{
+        key: MESSAGE_TYPE.TRANSFER_START,
+        id: fileId,
+        readable: ts.readable,
+      },
+      [ts.readable]
+    );
+    WorkerEvent.writer.set(fileId, ts.writable.getWriter());
     const newFileName = encodeURIComponent(fileName)
       .replace(/['()]/g, escape)
       .replace(/\*/g, "%2A");
@@ -44,13 +69,20 @@ export class WorkerEvent {
     document.body.appendChild(iframe);
   }
 
-  public static post(data: BufferType) {
-    if (!WorkerEvent.worker) return;
-    WorkerEvent.worker.active?.postMessage(data);
+  public static post(fileId: string, data: ArrayBuffer) {
+    const ts = WorkerEvent.writer.get(fileId);
+    ts?.write(new Uint8Array(data));
   }
 
   public static close(fileId: string) {
     const iframe = document.getElementById(fileId);
     iframe && iframe.remove();
+    WorkerEvent.channel?.port1.postMessage(<MessageType>{
+      key: MESSAGE_TYPE.TRANSFER_CLOSE,
+      id: fileId,
+    });
+    const ts = WorkerEvent.writer.get(fileId);
+    ts?.close();
+    WorkerEvent.writer.delete(fileId);
   }
 }
diff --git a/packages/webrtc/client/worker/index.ts b/packages/webrtc/client/worker/index.ts
@@ -1,7 +1,6 @@
 declare let self: ServiceWorkerGlobalScope;
-import type { BufferType } from "../../types/client";
-import { HEADER_KEY } from "../../types/worker";
-import { destructureChunk } from "../utils/binary";
+import type { MessageType } from "../../types/worker";
+import { HEADER_KEY, MESSAGE_TYPE } from "../../types/worker";
 
 self.addEventListener("install", () => {
   // 跳过等待 直接激活
@@ -15,24 +14,26 @@ self.addEventListener("activate", event => {
   console.log("Service Worker Activate");
 });
 
-type StreamTuple = [ReadableStream, ReadableStreamDefaultController<BufferType>, number];
+type StreamTuple = [ReadableStream<Uint8Array>];
 const map = new Map<string, StreamTuple>();
 
 self.onmessage = event => {
-  const data = <BufferType>event.data;
-  destructureChunk(data).then(({ id, series, data }) => {
-    const stream = map.get(id);
-    if (!stream) return void 0;
-    const [, controller, size] = stream;
-    // 需要处理 BackPressure 而 TransformStream 解决了问题
-    // 不能直接写入 ArrayBuffer 必须要写入 TypedArray 类型
-    controller.enqueue(new Uint8Array(data));
-    // 数据块序列号 [0, TOTAL)
-    if (series === size - 1) {
-      controller.close();
+  const port = event.ports[0];
+  if (!port) return void 0;
+  port.onmessage = event => {
+    const payload = event.data as MessageType;
+    if (!payload) return void 0;
+    if (payload.key === MESSAGE_TYPE.TRANSFER_START) {
+      // 直接使用 ReadableStream 需要处理 BackPressure 而 TransformStream 解决了问题
+      // controller.enqueue 不能直接写入 ArrayBuffer 必须要写入 TypedArray 类型
+      const { id, readable } = payload;
+      map.set(id, [readable]);
+    }
+    if (payload.key === MESSAGE_TYPE.TRANSFER_CLOSE) {
+      const { id } = payload;
       map.delete(id);
     }
-  });
+  };
 };
 
 self.onfetch = event => {
@@ -42,35 +43,30 @@ self.onfetch = event => {
   const fileName = search.get(HEADER_KEY.FILE_NAME);
   const fileSize = search.get(HEADER_KEY.FILE_SIZE);
   const fileTotal = search.get(HEADER_KEY.FILE_TOTAL);
-  if (fileId && fileName && fileSize && fileTotal) {
-    const newFileName = decodeURIComponent(fileName);
-    let controller: ReadableStreamDefaultController | null = null;
-    const readable = new ReadableStream({
-      start(ctr) {
-        controller = ctr;
-      },
-      cancel(reason) {
-        console.log("ReadableStream Aborted", reason);
-      },
-    });
-    map.set(fileId, [readable, controller!, Number(fileTotal)]);
-    const responseHeader = new Headers({
-      [HEADER_KEY.FILE_ID]: fileId,
-      [HEADER_KEY.FILE_SIZE]: fileSize,
-      [HEADER_KEY.FILE_NAME]: newFileName,
-      "Content-Type": "application/octet-stream; charset=utf-8",
-      "Content-Security-Policy": "default-src 'none'",
-      "X-Content-Security-Policy": "default-src 'none'",
-      "X-WebKit-CSP": "default-src 'none'",
-      "X-XSS-Protection": "1; mode=block",
-      "Cross-Origin-Embedder-Policy": "require-corp",
-      "Content-Disposition": "attachment; filename*=UTF-8''" + newFileName,
-      "Content-Length": fileSize,
-    });
-    const response = new Response(readable, {
-      headers: responseHeader,
-    });
-    return event.respondWith(response);
+  if (!fileId || !fileName || !fileSize || !fileTotal) {
+    return fetch(event.request);
+  }
+  const transfer = map.get(fileId);
+  if (!transfer) {
+    return null;
   }
-  return fetch(event.request);
+  const [readable] = transfer;
+  const newFileName = decodeURIComponent(fileName);
+  const responseHeader = new Headers({
+    [HEADER_KEY.FILE_ID]: fileId,
+    [HEADER_KEY.FILE_SIZE]: fileSize,
+    [HEADER_KEY.FILE_NAME]: newFileName,
+    "Content-Type": "application/octet-stream; charset=utf-8",
+    "Content-Security-Policy": "default-src 'none'",
+    "X-Content-Security-Policy": "default-src 'none'",
+    "X-WebKit-CSP": "default-src 'none'",
+    "X-XSS-Protection": "1; mode=block",
+    "Cross-Origin-Embedder-Policy": "require-corp",
+    "Content-Disposition": "attachment; filename*=UTF-8''" + newFileName,
+    "Content-Length": fileSize,
+  });
+  const response = new Response(readable, {
+    headers: responseHeader,
+  });
+  return event.respondWith(response);
 };
diff --git a/packages/webrtc/types/worker.ts b/packages/webrtc/types/worker.ts
@@ -1,6 +1,22 @@
+import type { Spread } from "./client";
+
 export const HEADER_KEY = {
   FILE_ID: "X-File-Id",
   FILE_NAME: "X-File-Name",
   FILE_SIZE: "X-File-Size",
   FILE_TOTAL: "X-File-Total",
 };
+
+export const MESSAGE_TYPE = {
+  INIT_CHANNEL: "INIT_CHANNEL",
+  TRANSFER_START: "TRANSFER_START",
+  TRANSFER_CLOSE: "TRANSFER_CLOSE",
+} as const;
+
+export type MessageTypeMap = {
+  [MESSAGE_TYPE.INIT_CHANNEL]: Record<string, never>;
+  [MESSAGE_TYPE.TRANSFER_START]: { id: string; readable: ReadableStream<Uint8Array> };
+  [MESSAGE_TYPE.TRANSFER_CLOSE]: { id: string };
+};
+
+export type MessageType = Spread<MessageTypeMap>;
