Add proper reference count-based source disposal system

src/hls/hls-demuxer.ts

@@ -576,6 +576,13 @@ export class HlsDemuxer extends Demuxer {
 	async getMimeType(): Promise<string> {
 		return 'application/vnd.apple.mpegurl';
 	}
+
+	override dispose(): void {
+		for (const segInput of this.segmentedInputs) {
+			segInput.dispose();
+		}
+		this.segmentedInputs.length = 0;
+	}
 }
 
 abstract class HlsInputTrackBacking implements InputTrackBacking {

src/hls/hls-segmented-input.ts

@@ -54,14 +54,17 @@ export class HlsSegmentedInput extends SegmentedInput {
 
 	runUpdateSegments() {
 		return this.currentUpdateSegmentsPromise ??= (async () => {
-			const remainingWaitTimeMs = this.getRemainingWaitTimeMs();
-			if (remainingWaitTimeMs > 0) {
-				await wait(remainingWaitTimeMs);
-			}
+			try {
+				const remainingWaitTimeMs = this.getRemainingWaitTimeMs();
+				if (remainingWaitTimeMs > 0) {
+					await wait(remainingWaitTimeMs);
+				}
 
-			this.lastSegmentUpdateTime = performance.now();
-			await this.updateSegments();
-			this.currentUpdateSegmentsPromise = null;
+				this.lastSegmentUpdateTime = performance.now();
+				await this.updateSegments();
+			} finally {
+				this.currentUpdateSegmentsPromise = null;
+			}
 		})();
 	}
 
@@ -88,11 +91,17 @@ export class HlsSegmentedInput extends SegmentedInput {
 
 		if (!lines) {
 			const source = await this.demuxer.input._getSourceUncached({ path: this.path });
-			const reader = new Reader(source);
+			source.ref();
 
-			const slice = await reader.requestEntireFile();
-			assert(slice);
-			lines = readAllLines(slice, slice.length, { ignore: canIgnoreLine });
+			try {
+				const reader = new Reader(source);
+
+				const slice = await reader.requestEntireFile();
+				assert(slice);
+				lines = readAllLines(slice, slice.length, { ignore: canIgnoreLine });
+			} finally {
+				source.unref();
+			}
 		}
 
 		let headerRead = false;
@@ -106,16 +115,10 @@ export class HlsSegmentedInput extends SegmentedInput {
 		let nextByteRange: { offset: number; length: number } | null = null;
 		let lastProgramDateTimeSeconds: number | null = null;
 
+		// Used for repeated parses where our job it is to only add the new segments
 		let prevLastSegment = last(this.segments) ?? null;
 
-		if (Math.PI === 3) {
-			// Stupid hack needed to prevent TypeScript from incorrectly narrowing the local variables; not sure if
-			// there is a better workaround
-			nextByteRange = { offset: 6, length: 7 };
-			lastByteRangeEnd = 1337;
-		}
-
-		const parseAndUpdateByteRange = (content: string) => {
+		const parseByteRange = (content: string) => {
 			const atIndex = content.indexOf('@');
 
 			const length = Number(atIndex === -1 ? content : content.slice(0, atIndex));
@@ -123,23 +126,15 @@ export class HlsSegmentedInput extends SegmentedInput {
 				throw new Error(`Invalid #EXT-X-BYTERANGE length '${content}'.`);
 			}
 
-			let offset: number;
+			let offset: number | null = null;
 			if (atIndex !== -1) {
 				offset = Number(content.slice(atIndex + 1));
 				if (!Number.isInteger(offset) || offset < 0) {
 					throw new Error(`Invalid #EXT-X-BYTERANGE offset '${content}'.`);
 				}
-			} else {
-				if (lastByteRangeEnd === null) {
-					throw new Error(
-						'Invalid M3U8 file; #EXT-X-BYTERANGE without offset requires a previous byte range.',
-					);
-				}
-				offset = lastByteRangeEnd;
 			}
 
-			nextByteRange = { offset, length };
-			lastByteRangeEnd = offset + length;
+			return { length, offset };
 		};
 
 		const setNextSequenceNumber = (number: number) => {
@@ -252,16 +247,21 @@ export class HlsSegmentedInput extends SegmentedInput {
 				}
 
 				const byteRange = attributes.get('byterange');
+				let parsedByteRange: ReturnType<typeof parseByteRange> | null = null;
 				if (byteRange !== null) {
-					parseAndUpdateByteRange(byteRange);
+					parsedByteRange = parseByteRange(byteRange);
+				}
+
+				if (parsedByteRange && parsedByteRange.offset === null) {
+					throw new Error('Invalid #EXT-X-MAP tag; BYTERANGE attribute must have a specified offset.');
 				}
 
 				if (!prevLastSegment) {
 					const fullPath = joinPaths(this.path, uri);
 					const location: HlsSegmentLocation = {
 						path: fullPath,
-						offset: nextByteRange?.offset ?? 0,
-						length: nextByteRange?.length ?? null,
+						offset: parsedByteRange?.offset ?? 0,
+						length: parsedByteRange?.length ?? null,
 					};
 
 					if (currentKey?.method === 'AES-128' && !currentKey.iv) {
@@ -345,10 +345,23 @@ export class HlsSegmentedInput extends SegmentedInput {
 
 				setNextSequenceNumber(number);
 			} else if (line.startsWith('#EXT-X-BYTERANGE:')) {
-				parseAndUpdateByteRange(line.slice(17));
+				const parsed = parseByteRange(line.slice(17));
+				if (parsed.offset === null) {
+					if (lastByteRangeEnd === null) {
+						throw new Error(
+							'Invalid M3U8 file; #EXT-X-BYTERANGE without offset requires a previous byte range.',
+						);
+					}
+					parsed.offset = lastByteRangeEnd;
+				}
+
+				nextByteRange = parsed as { length: number; offset: number };
+				lastByteRangeEnd = parsed.offset + parsed.length;
 			} else if (line.startsWith('#EXT-X-PROGRAM-DATE-TIME:')) {
 				if (prevLastSegment) {
-					continue; // No need to spend effort parsing dates if we're gonna discard it anyway
+					// No need to spend effort parsing dates if we're gonna discard it anyway. Also would be wrong to do
+					// the segment shifting!
+					continue;
 				}
 
 				const dateTime = line.slice(25);

src/input.ts

@@ -17,7 +17,7 @@ import {
 	TrackQuery,
 } from './input-track';
 import { PacketRetrievalOptions } from './media-sink';
-import { arrayArgmin, arrayCount, assert, desc, MaybePromise, polyfillSymbolDispose, prefer } from './misc';
+import { arrayArgmin, arrayCount, assert, desc, MaybePromise, polyfillSymbolDispose, prefer, removeItem } from './misc';
 import { Reader } from './reader';
 import { Source } from './source';
 
@@ -34,6 +34,15 @@ const sourceRequestsAreEqual = (a: SourceRequest, b: SourceRequest) => {
 	return a.path === b.path;
 };
 
+let inputFinalizationRegistry: FinalizationRegistry<Source[]> | null = null;
+if (typeof FinalizationRegistry !== 'undefined') {
+	inputFinalizationRegistry = new FinalizationRegistry((sources) => {
+		for (const source of sources) {
+			source.unref();
+		}
+	});
+}
+
 /**
  * The options for creating an Input object.
  * @group Input files & tracks
@@ -75,6 +84,10 @@ export class Input<S extends Source = Source> implements Disposable {
 	/** @internal */
 	_nextSourceCacheAge = 0;
 	/** @internal */
+	// This is an array, not a set, because the same source may be reffed multiple times and therefore also needs to be
+	// unreffed multiple times.
+	_reffedSources: Source[] = [];
+	/** @internal */
 	_sourceCache: {
 		request: SourceRequest;
 		sourcePromise: Promise<S>;
@@ -123,6 +136,8 @@ export class Input<S extends Source = Source> implements Disposable {
 		this._source = options.source;
 		this._initInput = options.initInput ?? null;
 		this._entryPath = options.entryPath ?? null;
+
+		inputFinalizationRegistry?.register(this, this._reffedSources, this);
 	}
 
 	async _getSourceUncached(request: SourceRequest) {
@@ -166,12 +181,19 @@ export class Input<S extends Source = Source> implements Disposable {
 			const entry = this._sourceCache[minAgeIndex]!;
 			this._sourceCache.splice(minAgeIndex, 1);
 
-			/*
 			void entry.sourcePromise
-				.then(source => source._dispose());
-			*/
+				.then((source) => {
+					source.unref();
+					removeItem(this._reffedSources, source);
+				});
 		}
 
+		void sourcePromise
+			.then((source) => {
+				source.ref();
+				this._reffedSources.push(source);
+			});
+
 		return sourcePromise;
 	}
 
@@ -187,6 +209,9 @@ export class Input<S extends Source = Source> implements Disposable {
 				source = await this._getSourceUncached({ path: this._entryPath });
 			}
 
+			source.ref();
+			this._reffedSources.push(source);
+
 			this._reader = new Reader(source);
 
 			for (const format of this._formats) {
@@ -392,14 +417,15 @@ export class Input<S extends Source = Source> implements Disposable {
 
 		this._disposed = true;
 
-		if (this._source instanceof Source) {
-			this._source._disposed = true;
-			this._source._dispose();
-		} else {
-			// TODO
-			// TODO
-			// throw new Error('TODO');
+		for (const source of this._reffedSources) {
+			source.unref();
 		}
+		this._reffedSources.length = 0;
+
+		inputFinalizationRegistry?.unregister(this);
+
+		void this._demuxerPromise
+			?.then(demuxer => demuxer.dispose());
 	}
 
 	/**

src/sample.ts

@@ -45,7 +45,7 @@ let lastAudioGcErrorLog = -Infinity;
 let finalizationRegistry: FinalizationRegistry<FinalizationRegistryValue> | null = null;
 if (typeof FinalizationRegistry !== 'undefined') {
 	finalizationRegistry = new FinalizationRegistry<FinalizationRegistryValue>((value) => {
-		const now = Date.now();
+		const now = performance.now();
 
 		if (value.type === 'video') {
 			if (now - lastVideoGcErrorLog >= 1000) {

src/segmented-input.ts

@@ -81,6 +81,13 @@ export abstract class SegmentedInput {
 			formats: [new VirtualInputFormat(() => new SegmentedInputDemuxer(this.input, this))],
 		});
 	}
+
+	dispose() {
+		for (const entry of this.inputCache) {
+			entry.input.dispose();
+		}
+		this.inputCache.length = 0;
+	}
 }
 
 class SegmentedInputDemuxer extends Demuxer {

src/source.ts

@@ -58,6 +58,8 @@ export abstract class Source {
 	abstract _dispose(): void;
 	/** @internal */
 	_disposed = false;
+	/** @internal */
+	_refCount = 0;
 
 	/** @internal */
 	private _sizePromise: Promise<number | null> | null = null;
@@ -119,6 +121,33 @@ export abstract class Source {
 
 	/** Called each time data is retrieved from the source. Will be called with the retrieved range (end exclusive). */
 	onread: ((start: number, end: number) => unknown) | null = null;
+
+	/**
+	 * Increases the internal reference count of this source. Call this method when you don't want the source to be
+	 * disposed.
+	 */
+	ref() {
+		if (this._disposed) {
+			throw new Error('Cannot ref a disposed source.');
+		}
+
+		this._refCount++;
+	}
+
+	/**
+	 * Decreases the internal reference count of this source, signalling a lost of interest in this source. If the
+	 * internal count reaches zero, meaning nobody is interested in the source anymore, its resources get disposed.
+	 */
+	unref() {
+		if (this._refCount > 0) {
+			this._refCount--;
+
+			if (this._refCount === 0) {
+				this._dispose();
+				this._disposed = true;
+			}
+		}
+	}
 }
 
 /**
@@ -2033,6 +2062,16 @@ export class RangedSource extends Source {
 	}
 
 	override _dispose(): void {
-		this._baseSource._dispose();
+		// Nada
+	}
+
+	override ref() {
+		super.ref();
+		this._baseSource.ref();
+	}
+
+	override unref() {
+		super.unref();
+		this._baseSource.unref();
 	}
 }