mirror of
https://github.com/wukko/cobalt.git
synced 2024-11-15 04:39:58 +00:00
stream: fix some memory leaks in internal stream handling (#581)
This commit is contained in:
parent
21d5b4b8d4
commit
ef97ff06af
2 changed files with 24 additions and 10 deletions
|
@ -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';
|
||||
|
||||
|
@ -36,21 +35,17 @@ async function* readChunks(streamInfo, size) {
|
|||
|
||||
read += received;
|
||||
}
|
||||
}
|
||||
|
||||
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();
|
||||
}
|
||||
}
|
||||
|
|
|
@ -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();
|
||||
}
|
||||
|
||||
|
|
Loading…
Reference in a new issue