stream: fix some memory leaks in internal stream handling (#581)

src/modules/stream/internal.js

@@ -1,6 +1,5 @@
 import { request } from 'undici';
 import { Readable } from 'node:stream';
-import { assert } from 'console';
 import { getHeaders, pipe } from './shared.js';
 import { handleHlsPlaylist, isHlsRequest } from './internal-hls.js';
 
@@ -38,19 +37,15 @@ async function* readChunks(streamInfo, size) {
     }
 }
 
-function chunkedStream(streamInfo, size) {
-    assert(streamInfo.controller instanceof AbortController);
-    const stream = Readable.from(readChunks(streamInfo, size));
-    return stream;
-}
-
 async function handleYoutubeStream(streamInfo, res) {
+    const { signal } = streamInfo.controller;
+
     try {
         const req = await fetch(streamInfo.url, {
             headers: getHeaders('youtube'),
             method: 'HEAD',
             dispatcher: streamInfo.dispatcher,
-            signal: streamInfo.controller.signal
+            signal
         });
 
         streamInfo.url = req.url;
@@ -60,7 +55,16 @@ async function handleYoutubeStream(streamInfo, res) {
             return res.end();
         }
 
-        const stream = chunkedStream(streamInfo, size);
+        const generator = readChunks(streamInfo, size);
+
+        const abortGenerator = () => {
+            generator.return();
+            signal.removeEventListener('abort', abortGenerator);
+        }
+
+        signal.addEventListener('abort', abortGenerator);
+    
+        const stream = Readable.from(generator);
 
         for (const headerName of ['content-type', 'content-length']) {
             const headerValue = req.headers.get(headerName);
@@ -69,6 +73,7 @@ async function handleYoutubeStream(streamInfo, res) {
 
         pipe(stream, res, () => res.end());
     } catch {
+        signal.abort();
         res.end();
     }
 }

src/modules/stream/manage.js

@@ -78,16 +78,25 @@ export function createInternalStream(url, obj = {}) {
     }
 
     const streamID = nanoid();
+    const controller = new AbortController();
     internalStreamCache[streamID] = {
         url,
         service: obj.service,
         headers: obj.headers,
-        controller: new AbortController(),
+        controller,
         dispatcher
     };
 
     let streamLink = new URL('/api/istream', `http://127.0.0.1:${env.apiPort}`);
     streamLink.searchParams.set('id', streamID);
+
+    const cleanup = () => {
+        destroyInternalStream(streamLink);
+        controller.signal.removeEventListener('abort', cleanup);
+    }
+
+    controller.signal.addEventListener('abort', cleanup);
+
     return streamLink.toString();
 }