Merge pull request #823 from NapNeko/refactor-worker

refactor: 即刻起逐出piscina

package.json

@@ -66,7 +66,6 @@
   "dependencies": {
     "@ffmpeg.wasm/core-mt": "^0.13.2",
     "express": "^5.0.0",
-    "piscina": "^4.7.0",
     "silk-wasm": "^3.6.1",
     "ws": "^8.18.0"
   }

src/common/audio-worker.ts

@@ -1,9 +1,20 @@
 import { encode } from 'silk-wasm';
+import { parentPort } from 'worker_threads';
 
 export interface EncodeArgs {
     input: ArrayBufferView | ArrayBuffer
     sampleRate: number
 }
-export default async ({ input, sampleRate }: EncodeArgs) => {
+export function recvTask<T>(cb: (taskData: T) => Promise<unknown>) {
+    parentPort?.on('message', async (taskData: T) => {
+        try {
+            let ret = await cb(taskData);
+            parentPort?.postMessage(ret);
+        } catch (error: unknown) {
+            parentPort?.postMessage({ error: (error as Error).message });
+        }
+    });
+}
+recvTask<EncodeArgs>(async ({ input, sampleRate }) => {
     return await encode(input, sampleRate);
-};
+});

src/common/audio.ts

@@ -1,4 +1,3 @@
-import Piscina from 'piscina';
 import fsPromise from 'fs/promises';
 import path from 'node:path';
 import { randomUUID } from 'crypto';
@@ -6,16 +5,16 @@ import { EncodeResult, getDuration, getWavFileInfo, isSilk, isWav } from 'silk-w
 import { LogWrapper } from '@/common/log';
 import { EncodeArgs } from '@/common/audio-worker';
 import { FFmpegService } from '@/common/ffmpeg';
+import { runTask } from './worker';
+import { fileURLToPath } from 'node:url';
 
 const ALLOW_SAMPLE_RATE = [8000, 12000, 16000, 24000, 32000, 44100, 48000];
 
-async function getWorkerPath() {
+function getWorkerPath() {
-    return new URL(/* @vite-ignore */ './audio-worker.mjs', import.meta.url).href;
+    //return new URL(/* @vite-ignore */ './audio-worker.mjs', import.meta.url).href;
+    return path.join(path.dirname(fileURLToPath(import.meta.url)), 'audio-worker.mjs');
 }
 
-const piscina = new Piscina<EncodeArgs, EncodeResult>({
-    filename: await getWorkerPath(),
-});
 
 async function guessDuration(pttPath: string, logger: LogWrapper) {
     const pttFileInfo = await fsPromise.stat(pttPath);
@@ -46,7 +45,7 @@ export async function encodeSilk(filePath: string, TEMP_DIR: string, logger: Log
             const { input, sampleRate } = isWav(file)
                 ? await handleWavFile(file, filePath, pcmPath)
                 : { input: await FFmpegService.convert(filePath, pcmPath), sampleRate: 24000 };
-            const silk = await piscina.run({ input: input, sampleRate: sampleRate });
+            const silk = await runTask<EncodeArgs, EncodeResult>(getWorkerPath(), { input: input, sampleRate: sampleRate });
             fsPromise.unlink(pcmPath).catch((e) => logger.logError('删除临时文件失败', pcmPath, e));
             await fsPromise.writeFile(pttPath, Buffer.from(silk.data));
             logger.log(`语音文件${filePath}转换成功!`, pttPath, '时长:', silk.duration);

src/common/ffmpeg-worker.ts

@@ -5,6 +5,17 @@ import { readFileSync, statSync, writeFileSync } from 'fs';
 import type { VideoInfo } from './video';
 import { fileTypeFromFile } from 'file-type';
 import imageSize from 'image-size';
+import { parentPort } from 'worker_threads';
+export function recvTask<T>(cb: (taskData: T) => Promise<unknown>) {
+    parentPort?.on('message', async (taskData: T) => {
+        try {
+            let ret = await cb(taskData);
+            parentPort?.postMessage(ret);
+        } catch (error: unknown) {
+            parentPort?.postMessage({ error: (error as Error).message });
+        }
+    });
+}
 class FFmpegService {
     public static async extractThumbnail(videoPath: string, thumbnailPath: string): Promise<void> {
         const ffmpegInstance = await FFmpeg.create({ core: '@ffmpeg.wasm/core-mt' });
@@ -137,15 +148,18 @@ interface FFmpegTask {
 }
 export default async function handleFFmpegTask({ method, args }: FFmpegTask): Promise<any> {
     switch (method) {
-    case 'extractThumbnail':
+        case 'extractThumbnail':
-        return await FFmpegService.extractThumbnail(...args as [string, string]);
+            return await FFmpegService.extractThumbnail(...args as [string, string]);
-    case 'convertFile':
+        case 'convertFile':
-        return await FFmpegService.convertFile(...args as [string, string, string]);
+            return await FFmpegService.convertFile(...args as [string, string, string]);
-    case 'convert':
+        case 'convert':
-        return await FFmpegService.convert(...args as [string, string]);
+            return await FFmpegService.convert(...args as [string, string]);
-    case 'getVideoInfo':
+        case 'getVideoInfo':
-        return await FFmpegService.getVideoInfo(...args as [string, string]);
+            return await FFmpegService.getVideoInfo(...args as [string, string]);
-    default:
+        default:
-        throw new Error(`Unknown method: ${method}`);
+            throw new Error(`Unknown method: ${method}`);
     }
-}
+}
+recvTask<FFmpegTask>(async ({ method, args }: FFmpegTask) => {
+    return await handleFFmpegTask({ method, args });
+});

src/common/ffmpeg.ts

@@ -1,6 +1,8 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
-import Piscina from 'piscina';
 import { VideoInfo } from './video';
+import path from 'path';
+import { fileURLToPath } from 'url';
+import { runTask } from './worker';
 
 type EncodeArgs = {
     method: 'extractThumbnail' | 'convertFile' | 'convert' | 'getVideoInfo';
@@ -9,42 +11,26 @@ type EncodeArgs = {
 
 type EncodeResult = any;
 
-async function getWorkerPath() {
+function getWorkerPath() {
-    return new URL(/* @vite-ignore */ './ffmpeg-worker.mjs', import.meta.url).href;
+    return path.join(path.dirname(fileURLToPath(import.meta.url)), './ffmpeg-worker.mjs');
 }
 
 export class FFmpegService {
     public static async extractThumbnail(videoPath: string, thumbnailPath: string): Promise<void> {
-        const piscina = new Piscina<EncodeArgs, EncodeResult>({
+        await runTask<EncodeArgs, EncodeResult>(getWorkerPath(), { method: 'extractThumbnail', args: [videoPath, thumbnailPath] });
-            filename: await getWorkerPath(),
-        });
-        await piscina.run({ method: 'extractThumbnail', args: [videoPath, thumbnailPath] });
-        await piscina.destroy();
     }
 
     public static async convertFile(inputFile: string, outputFile: string, format: string): Promise<void> {
-        const piscina = new Piscina<EncodeArgs, EncodeResult>({
+        await runTask<EncodeArgs, EncodeResult>(getWorkerPath(), { method: 'convertFile', args: [inputFile, outputFile, format] });
-            filename: await getWorkerPath(),
-        });
-        await piscina.run({ method: 'convertFile', args: [inputFile, outputFile, format] });
-        await piscina.destroy();
     }
 
     public static async convert(filePath: string, pcmPath: string): Promise<Buffer> {
-        const piscina = new Piscina<EncodeArgs, EncodeResult>({
+        const result = await runTask<EncodeArgs, EncodeResult>(getWorkerPath(), { method: 'convert', args: [filePath, pcmPath] });
-            filename: await getWorkerPath(),
-        });
-        const result = await piscina.run({ method: 'convert', args: [filePath, pcmPath] });
-        await piscina.destroy();
         return result;
     }
 
     public static async getVideoInfo(videoPath: string, thumbnailPath: string): Promise<VideoInfo> {
-        const piscina = new Piscina<EncodeArgs, EncodeResult>({
+        const result = await await runTask<EncodeArgs, EncodeResult>(getWorkerPath(), { method: 'getVideoInfo', args: [videoPath, thumbnailPath] });
-            filename: await getWorkerPath(),
-        });
-        const result = await piscina.run({ method: 'getVideoInfo', args: [videoPath, thumbnailPath] });
-        await piscina.destroy();
         return result;
     }
 }

src/common/worker.ts

@@ -0,0 +1,29 @@
+import { Worker } from 'worker_threads';
+
+export async function runTask<T, R>(workerScript: string, taskData: T): Promise<R> {
+    let worker = new Worker(workerScript);
+    try {
+        return await new Promise<R>((resolve, reject) => {
+            worker.on('message', (result: R) => {
+                resolve(result);
+            });
+
+            worker.on('error', (error) => {
+                reject(new Error(`Worker error: ${error.message}`));
+            });
+
+            worker.on('exit', (code) => {
+                if (code !== 0) {
+                    reject(new Error(`Worker stopped with exit code ${code}`));
+                }
+            });
+            worker.postMessage(taskData);
+        });
+    } catch (error: unknown) {
+        throw new Error(`Failed to run task: ${(error as Error).message}`);
+    } finally {
+        // Ensure the worker is terminated after the promise is settled
+        worker.terminate();
+    }
+}
+

vite.config.ts

@@ -8,8 +8,7 @@ const external = [
     'silk-wasm',
     'ws',
     'express',
-    '@ffmpeg.wasm/core-mt',
+    '@ffmpeg.wasm/core-mt'
-    'piscina'
 ];
 const nodeModules = [...builtinModules, builtinModules.map((m) => `node:${m}`)].flat();