Add SourceRef system to keep track of used Sources (fixes #332)

src/aes.ts

@@ -236,7 +236,11 @@ export class Aes128CbcContext {
 	}
 }
 
-export const createAesDecryptStream = (reader: Reader, getInit: () => MaybePromise<Aes128CbcContextInit>) => {
+export const createAes128CbcDecryptStream = (
+	reader: Reader,
+	getInit: () => MaybePromise<Aes128CbcContextInit>,
+	close: () => unknown,
+) => {
 	let initted = false;
 	let pos = 0;
 	const CHUNK_SIZE = 2 ** 16;
@@ -288,7 +292,12 @@ export const createAesDecryptStream = (reader: Reader, getInit: () => MaybePromi
 
 				controller.enqueue(trimmedOutput);
 				controller.close();
+
+				close();
 			}
 		},
+		cancel: () => {
+			close();
+		},
 	});
 };

src/hls/hls-segmented-input.ts

@@ -1,9 +1,9 @@
-import { AES_128_BLOCK_SIZE, createAesDecryptStream } from '../aes';
+import { AES_128_BLOCK_SIZE, createAes128CbcDecryptStream } from '../aes';
 import { ENCRYPTION_KEY_CACHE_GROUP, Input } from '../input';
 import { Segment, SegmentedInput, SegmentRetrievalOptions } from '../segmented-input';
 import { toDataView, joinPaths, last, assert, binarySearchLessOrEqual, arrayArgmin, wait } from '../misc';
 import { readAllLines, readBytes, Reader } from '../reader';
-import { ReadableStreamSource, Source } from '../source';
+import { ReadableStreamSource, SourceRef } from '../source';
 import { HlsDemuxer } from './hls-demuxer';
 import { AttributeList, canIgnoreLine } from './hls-misc';
 
@@ -89,18 +89,12 @@ export class HlsSegmentedInput extends SegmentedInput {
 		this.nextLines = null;
 
 		if (!lines) {
-			const source = await this.demuxer.input._getSourceUncached({ path: this.path });
-			source.ref();
+			using ref = await this.demuxer.input._getSourceUncached({ path: this.path });
+			const reader = new Reader(ref.source);
 
-			try {
-				const reader = new Reader(source);
-
-				const slice = await reader.requestEntireFile();
-				assert(slice);
-				lines = readAllLines(slice, slice.length, { ignore: canIgnoreLine });
-			} finally {
-				source.unref();
-			}
+			const slice = await reader.requestEntireFile();
+			assert(slice);
+			lines = readAllLines(slice, slice.length, { ignore: canIgnoreLine });
 		}
 
 		let headerRead = false;
@@ -514,34 +508,44 @@ export class HlsSegmentedInput extends SegmentedInput {
 					return this.input._getSourceUncached(request);
 				}
 
-				let source: Source;
+				let ref: SourceRef;
 				const needsSlice = hlsSegment.location.offset > 0 || hlsSegment.location.length !== null;
 
 				if (!hlsSegment.encryption) {
-					source = await this.input._getSourceCached(request);
+					ref = await this.input._getSourceCached(request);
+
 					if (needsSlice) {
-						source = source.slice(hlsSegment.location.offset, hlsSegment.location.length ?? undefined);
+						const slice = ref.source.slice(
+							hlsSegment.location.offset,
+							hlsSegment.location.length ?? undefined,
+						);
+						const sliceRef = slice.ref();
+						ref.free();
+						ref = sliceRef;
 					}
 				} else {
 					assert(hlsSegment.encryption.iv);
 
-					let ciphertextSource = await this.input._getSourceCached(request);
+					let ciphertextRef = await this.input._getSourceCached(request);
 					if (needsSlice) {
 						// Slice before decrypting
-						ciphertextSource = ciphertextSource.slice(
+						const slice = ciphertextRef.source.slice(
 							hlsSegment.location.offset,
 							hlsSegment.location.length ?? undefined,
 						);
+						const sliceRef = slice.ref();
+						ciphertextRef.free();
+						ciphertextRef = sliceRef;
 					}
 
-					const ciphertextReader = new Reader(ciphertextSource);
+					const ciphertextReader = new Reader(ciphertextRef.source);
 
-					const stream = createAesDecryptStream(ciphertextReader, async () => {
-						const keySource = await this.input._getSourceCached(
+					const stream = createAes128CbcDecryptStream(ciphertextReader, async () => {
+						using keyRef = await this.input._getSourceCached(
 							{ path: hlsSegment.encryption!.keyUri },
 							ENCRYPTION_KEY_CACHE_GROUP,
 						);
-						const keyReader = new Reader(keySource);
+						const keyReader = new Reader(keyRef.source);
 						const keySlice = await keyReader.requestSlice(0, AES_128_BLOCK_SIZE);
 						if (!keySlice) {
 							throw new Error('Invalid AES-128 key; expected at least 16 bytes of data.');
@@ -549,12 +553,14 @@ export class HlsSegmentedInput extends SegmentedInput {
 						const key = readBytes(keySlice, AES_128_BLOCK_SIZE);
 
 						return { key, iv: hlsSegment.encryption!.iv! };
+					}, () => {
+						ciphertextRef.free();
 					});
 
-					source = new ReadableStreamSource(stream);
+					ref = new ReadableStreamSource(stream).ref();
 				}
 
-				return source;
+				return ref!;
 			},
 			formats: this.input._formats,
 			initInput: initInput ?? undefined,
@@ -569,7 +575,10 @@ export class HlsSegmentedInput extends SegmentedInput {
 		const MAX_INPUT_CACHE_SIZE = 4;
 		if (this.inputCache.length > MAX_INPUT_CACHE_SIZE) {
 			const minAgeIndex = arrayArgmin(this.inputCache, x => x.age);
+			assert(minAgeIndex !== -1);
 			this.inputCache.splice(minAgeIndex, 1);
+
+			// DON'T dispose here; the Input might still be used! The source disposal will happen with GC logic
 		}
 
 		return input;

src/input.ts

@@ -19,7 +19,7 @@ import {
 import { PacketRetrievalOptions } from './media-sink';
 import { arrayArgmin, arrayCount, assert, desc, MaybePromise, polyfillSymbolDispose, prefer, removeItem } from './misc';
 import { Reader } from './reader';
-import { Source } from './source';
+import { Source, SourceRef } from './source';
 
 polyfillSymbolDispose();
 
@@ -34,11 +34,11 @@ const sourceRequestsAreEqual = (a: SourceRequest, b: SourceRequest) => {
 	return a.path === b.path;
 };
 
-let inputFinalizationRegistry: FinalizationRegistry<Source[]> | null = null;
+let inputFinalizationRegistry: FinalizationRegistry<SourceRef[]> | null = null;
 if (typeof FinalizationRegistry !== 'undefined') {
-	inputFinalizationRegistry = new FinalizationRegistry((sources) => {
-		for (const source of sources) {
-			source.unref();
+	inputFinalizationRegistry = new FinalizationRegistry((refs) => {
+		for (const ref of refs) {
+			ref.free();
 		}
 	});
 }
@@ -52,11 +52,18 @@ export type InputOptions<S extends Source = Source> = {
 	/** A list of supported formats. If the source file is not of one of these formats, then it cannot be read. */
 	formats: InputFormat[];
 	/** The source from which data will be read. */
-	source: S | ((request: SourceRequest) => MaybePromise<S>);
+	source: S | SourceRef<S> | ((request: SourceRequest) => MaybePromise<S | SourceRef<S>>);
 	entryPath?: string;
 	initInput?: Input;
 };
 
+type SourceCacheEntry<S extends Source> = {
+	request: SourceRequest;
+	sourceRef: SourceRef<S>;
+	age: number;
+	cacheGroup: number;
+};
+
 /**
  * Represents an input media file. This is the root object from which all media read operations start.
  * @group Input files & tracks
@@ -64,7 +71,7 @@ export type InputOptions<S extends Source = Source> = {
  */
 export class Input<S extends Source = Source> implements Disposable {
 	/** @internal */
-	_source: InputOptions<S>['source'];
+	_source: SourceRef<S> | ((request: SourceRequest) => MaybePromise<S | SourceRef<S>>);
 	/** @internal */
 	_formats: InputFormat[];
 	/** @internal */
@@ -84,15 +91,14 @@ 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[] = [];
+	_sourceRefs: SourceRef[] = [];
 	/** @internal */
-	_sourceCache: {
+	_sourceCache: SourceCacheEntry<S>[] = [];
+	/** @internal */
+	_sourceCachePromises: {
 		request: SourceRequest;
-		sourcePromise: Promise<S>;
-		age: number;
 		cacheGroup: number;
+		promise: Promise<SourceCacheEntry<S>>;
 	}[] = [];
 
 	/**
@@ -119,9 +125,6 @@ export class Input<S extends Source = Source> implements Disposable {
 		if (!(options.source instanceof Source) && typeof options.source !== 'function') {
 			throw new TypeError('options.source must be a Source or a function that returns a Source.');
 		}
-		if (options.source instanceof Source && options.source._disposed) {
-			throw new TypeError('options.source must not be disposed.');
-		}
 		if (typeof options.source === 'function' && options.entryPath === undefined) {
 			throw new TypeError('options.entryPath must be provided when options.source is a function.');
 		}
@@ -133,86 +136,123 @@ export class Input<S extends Source = Source> implements Disposable {
 		}
 
 		this._formats = options.formats;
-		this._source = options.source;
 		this._initInput = options.initInput ?? null;
 		this._entryPath = options.entryPath ?? null;
 
-		inputFinalizationRegistry?.register(this, this._reffedSources, this);
+		if (options.source instanceof Source) {
+			this._source = options.source.ref();
+		} else {
+			this._source = options.source;
+		}
+
+		if (this._source instanceof SourceRef) {
+			this._sourceRefs.push(this._source);
+		}
+
+		inputFinalizationRegistry?.register(this, this._sourceRefs, this);
 	}
 
 	async _getSourceUncached(request: SourceRequest) {
 		assert(typeof this._source === 'function');
 
 		const source = await this._source(request);
-		if (!(source instanceof Source)) {
-			throw new TypeError('The source function must return a Source.');
+		if (!(source instanceof Source || source instanceof SourceRef)) {
+			throw new TypeError('The source function must return a Source or a SourceRef.');
 		}
-		if (source._disposed) {
+		if (source instanceof Source && source._disposed) {
 			throw new TypeError('The returned Source must not be disposed.');
 		}
 
-		this.onSource?.(source, request);
+		let ref: SourceRef<S>;
+		if (source instanceof Source) {
+			ref = source.ref();
+		} else {
+			ref = source;
+		}
 
-		return source;
+		this.onSource?.(ref.source, request);
+		return ref;
 	}
 
-	_getSourceCached(request: SourceRequest, cacheGroup = DEFAULT_SOURCE_CACHE_GROUP) {
+	_getSourceCached(request: SourceRequest, cacheGroup = DEFAULT_SOURCE_CACHE_GROUP): Promise<SourceRef<S>> {
 		const cachedEntry = this._sourceCache.find(x =>
 			x.cacheGroup === cacheGroup && sourceRequestsAreEqual(x.request, request),
 		);
 		if (cachedEntry) {
 			cachedEntry.age++;
-			return cachedEntry.sourcePromise;
+			return Promise.resolve(cachedEntry.sourceRef.source.ref());
 		}
 
-		const sourcePromise = this._getSourceUncached(request);
-		this._sourceCache.push({
+		const cachedPromiseEntry = this._sourceCachePromises.find(x =>
+			x.cacheGroup === cacheGroup && sourceRequestsAreEqual(x.request, request),
+		);
+		if (cachedPromiseEntry) {
+			return cachedPromiseEntry.promise.then(x => x.sourceRef.source.ref());
+		}
+
+		const promise = (async () => {
+			const sourceRef = await this._getSourceUncached(request);
+
+			const cacheEntry: SourceCacheEntry<S> = {
+				request,
+				sourceRef,
+				age: this._nextSourceCacheAge++,
+				cacheGroup,
+			};
+
+			this._sourceCache.push(cacheEntry);
+
+			const MAX_SOURCE_CACHE_SIZE = 4;
+			const count = arrayCount(
+				this._sourceCache,
+				x => x.cacheGroup === cacheGroup && x.sourceRef.source._refCount === 1,
+			);
+
+			if (count > MAX_SOURCE_CACHE_SIZE) {
+				const minAgeIndex = arrayArgmin(
+					this._sourceCache,
+					x => x.cacheGroup === cacheGroup && x.sourceRef.source._refCount === 1 ? x.age : Infinity,
+				);
+				assert(minAgeIndex !== -1);
+				const entry = this._sourceCache[minAgeIndex]!;
+				this._sourceCache.splice(minAgeIndex, 1);
+
+				entry.sourceRef.free();
+				removeItem(this._sourceRefs, sourceRef);
+			}
+
+			this._sourceRefs.push(sourceRef);
+
+			const promiseIndex = this._sourceCachePromises.findIndex(x => x.request === request);
+			assert(promiseIndex !== -1);
+			this._sourceCachePromises.splice(promiseIndex, 1);
+
+			return cacheEntry;
+		})();
+
+		this._sourceCachePromises.push({
 			request,
-			sourcePromise,
-			age: this._nextSourceCacheAge++,
 			cacheGroup,
+			promise,
 		});
 
-		const MAX_SOURCE_CACHE_SIZE = 4;
-		const count = arrayCount(this._sourceCache, x => x.cacheGroup === cacheGroup);
-
-		if (count > MAX_SOURCE_CACHE_SIZE) {
-			const minAgeIndex = arrayArgmin(this._sourceCache, x => x.cacheGroup === cacheGroup ? x.age : Infinity);
-			const entry = this._sourceCache[minAgeIndex]!;
-			this._sourceCache.splice(minAgeIndex, 1);
-
-			void entry.sourcePromise
-				.then((source) => {
-					source.unref();
-					removeItem(this._reffedSources, source);
-				});
-		}
-
-		void sourcePromise
-			.then((source) => {
-				source.ref();
-				this._reffedSources.push(source);
-			});
-
-		return sourcePromise;
+		return promise.then(x => x.sourceRef.source.ref());
 	}
 
 	/** @internal */
 	_getDemuxer() {
 		return this._demuxerPromise ??= (async () => {
-			let source: Source;
-			if (this._source instanceof Source) {
-				source = this._source;
-				this.onSource?.(source, null);
+			let ref: SourceRef;
+			if (this._source instanceof SourceRef) {
+				ref = this._source;
+				this.onSource?.(ref.source, null);
 			} else {
 				assert(this._entryPath !== null);
-				source = await this._getSourceUncached({ path: this._entryPath });
+				ref = await this._getSourceUncached({ path: this._entryPath });
+				this._sourceRefs.push(ref);
 			}
 
-			source.ref();
-			this._reffedSources.push(source);
-
-			this._reader = new Reader(source);
+			this._reader = new Reader(ref.source);
 
 			for (const format of this._formats) {
 				const canRead = await format._canReadInput(this);
@@ -227,29 +267,15 @@ export class Input<S extends Source = Source> implements Disposable {
 	}
 
 	/**
-	 * Returns a source for the given request.
-	 *
-	 * If this input was created with a direct {@link Source}, that source is always returned. If this input was created
-	 * with a source function, this method resolves it using the provided request or the entry path.
-	 */
-	getSource(request?: SourceRequest): MaybePromise<S> {
-		if (this._source instanceof Source) {
-			return this._source;
-		}
-
-		assert(this._entryPath !== null);
-		return this._getSourceCached(request ?? { path: this._entryPath });
-	}
-
-	/**
-	 * @deprecated Use {@link getSource} instead.
+	 * @deprecated Prefer not using this getter since it is ill-defined for files driven by multiple sources. The
+	 * {@link Input.onSource} callback provides an alternative.
 	 *
 	 * Returns the source from which this input file reads data for the entry path. Throws if the source-resolving
 	 * function returns a Promise.
 	 */
-	get source() {
-		if (this._source instanceof Source) {
-			return this._source;
+	get source(): S {
+		if (this._source instanceof SourceRef) {
+			return this._source.source;
 		}
 
 		assert(this._entryPath !== null);
@@ -262,7 +288,11 @@ export class Input<S extends Source = Source> implements Disposable {
 			);
 		}
 
-		return source;
+		if (source instanceof Source) {
+			return source;
+		} else {
+			return source.source;
+		}
 	}
 
 	/**
@@ -432,10 +462,10 @@ export class Input<S extends Source = Source> implements Disposable {
 
 		this._disposed = true;
 
-		for (const source of this._reffedSources) {
-			source.unref();
+		for (const ref of this._sourceRefs) {
+			ref.free();
 		}
-		this._reffedSources.length = 0;
+		this._sourceRefs.length = 0;
 
 		inputFinalizationRegistry?.unregister(this);
 

src/isobmff/isobmff-demuxer.ts

@@ -2824,7 +2824,9 @@ abstract class IsobmffTrackBacking implements InputTrackBacking {
 				sampleInfo.sampleSize,
 			);
 			if (slice instanceof Promise) slice = await slice;
-			assert(slice);
+			if (!slice) {
+				return null; // Data is outside
+			}
 
 			data = readBytes(slice, sampleInfo.sampleSize);
 		}
@@ -2864,7 +2866,9 @@ abstract class IsobmffTrackBacking implements InputTrackBacking {
 				fragmentSample.byteSize,
 			);
 			if (slice instanceof Promise) slice = await slice;
-			assert(slice);
+			if (!slice) {
+				return null; // Data is outside
+			}
 
 			data = readBytes(slice, fragmentSample.byteSize);
 		}

src/segmented-input.ts

@@ -99,6 +99,8 @@ export abstract class SegmentedInput {
 			entry.input.dispose();
 		}
 		this.inputCache.length = 0;
+
+		this.virtualInput?.dispose();
 	}
 }
 

src/source.ts

@@ -16,6 +16,7 @@ import {
 	isWebKit,
 	MaybePromise,
 	mergeRequestInit,
+	polyfillSymbolDispose,
 	promiseWithResolvers,
 	retriedFetch,
 	toDataView,
@@ -25,6 +26,8 @@ import {
 import * as nodeAlias from './node';
 import { InputDisposedError } from './input';
 
+polyfillSymbolDispose();
+
 const node = typeof nodeAlias !== 'undefined'
 	? nodeAlias // Aliasing it prevents some bundler warnings
 	: undefined!;
@@ -123,30 +126,56 @@ export abstract class Source {
 	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.
+	 * Creates a new `SourceRef` pointing to this source. You are expected to call `.free()` on said `SourceRef` when
+	 * you're done with it.
 	 */
 	ref() {
-		if (this._disposed) {
+		return new SourceRef(this);
+	}
+}
+
+export class SourceRef<S extends Source = Source> implements Disposable {
+	private _source: S | null;
+	freed = false;
+
+	constructor(source: S) {
+		if (source._disposed) {
 			throw new Error('Cannot ref a disposed source.');
 		}
 
-		this._refCount++;
+		source._refCount++;
+		this._source = source;
 	}
 
-	/**
-	 * 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;
-			}
+	get source() {
+		if (!this._source) {
+			throw new Error('Can\'t get source; ref has already been freed.');
 		}
+
+		return this._source;
+	}
+
+	free() {
+		if (this.freed) {
+			return;
+		}
+
+		const source = this.source;
+		assert(source._refCount > 0);
+
+		source._refCount--;
+
+		if (source._refCount === 0) {
+			source._dispose();
+			source._disposed = true;
+		}
+
+		this.freed = true;
+		this._source = null;
+	}
+
+	[Symbol.dispose]() {
+		this.free();
 	}
 }
 
@@ -1224,6 +1253,7 @@ export class ReadableStreamSource extends Source {
 	_dispose() {
 		this._pendingSlices.length = 0;
 		this._cache.length = 0;
+		void this._reader?.cancel();
 	}
 }
 
@@ -1994,6 +2024,8 @@ export class RangedSource extends Source {
 	/** @internal */
 	_baseSource: Source;
 	/** @internal */
+	_ref: SourceRef | null = null;
+	/** @internal */
 	_offset: number;
 	/** @internal */
 	_length: number | null;
@@ -2001,6 +2033,10 @@ export class RangedSource extends Source {
 	constructor(baseSource: Source, offset: number, length?: number) {
 		super();
 
+		if (baseSource._disposed) {
+			throw new Error('Cannot create a slice of a disposed source.');
+		}
+
 		this._baseSource = baseSource;
 		this._offset = offset;
 		this._length = length ?? null;
@@ -2062,16 +2098,11 @@ export class RangedSource extends Source {
 	}
 
 	override _dispose(): void {
-		// Nada
+		this._ref?.free();
 	}
 
 	override ref() {
-		super.ref();
-		this._baseSource.ref();
-	}
-
-	override unref() {
-		super.unref();
-		this._baseSource.unref();
+		this._ref ??= this._baseSource.ref();
+		return super.ref();
 	}
 }

test/node/aes.test.ts

@@ -1,7 +1,7 @@
 import { expect, test } from 'vitest';
 import { Reader } from '../../src/reader.js';
 import { BufferSource } from '../../src/source.js';
-import { createAesDecryptStream } from '../../src/aes.js';
+import { createAes128CbcDecryptStream } from '../../src/aes.js';
 
 // getRandomValues is length-limited, so let's just do this
 export const fillRandom = <T extends Uint8Array>(buffer: T) => {
@@ -27,7 +27,7 @@ test('createAesDecryptStream', async () => {
 		const source = new BufferSource(ciphertext);
 		const reader = new Reader(source);
 
-		const stream = createAesDecryptStream(reader, () => ({ key, iv }));
+		const stream = createAes128CbcDecryptStream(reader, () => ({ key, iv }));
 		const streamReader = stream.getReader();
 
 		const chunks: Uint8Array[] = [];

test/node/source-lifetime.test.ts

@@ -0,0 +1,56 @@
+import { expect, test } from 'vitest';
+import { FilePathSource } from '../../src/source.js';
+import path from 'node:path';
+import { Input } from '../../src/input.js';
+import { ALL_FORMATS, MP4 } from '../../src/input-format.js';
+
+const __dirname = new URL('.', import.meta.url).pathname;
+
+test('Direct source disposal', async () => {
+	const filePath = path.join(__dirname, '../public/video.mp4');
+	const source = new FilePathSource(filePath);
+
+	expect(!source._disposed);
+
+	const ref = source.ref();
+	ref.free();
+	expect(source._disposed);
+});
+
+test('Implicit source disposal', async () => {
+	const filePath = path.join(__dirname, '../public/video.mp4');
+	const source = new FilePathSource(filePath);
+
+	const input = new Input({
+		source,
+		formats: ALL_FORMATS,
+	});
+	expect(await input.getFormat()).toBe(MP4);
+
+	expect(!source._disposed);
+	input.dispose();
+	expect(source._disposed);
+});
+
+test('Implicit source disposal, double input', async () => {
+	const filePath = path.join(__dirname, '../public/video.mp4');
+	const source = new FilePathSource(filePath);
+
+	const input1 = new Input({
+		source,
+		formats: ALL_FORMATS,
+	});
+	const input2 = new Input({
+		source,
+		formats: ALL_FORMATS,
+	});
+
+	expect(await input1.getFormat()).toBe(MP4);
+	expect(await input2.getFormat()).toBe(MP4);
+
+	expect(!source._disposed);
+	input1.dispose();
+	expect(!source._disposed);
+	input2.dispose();
+	expect(source._disposed);
+});