diff --git a/src/aes.ts b/src/aes.ts index 7af4370..f8296a2 100644 --- a/src/aes.ts +++ b/src/aes.ts @@ -236,7 +236,11 @@ export class Aes128CbcContext { } } -export const createAesDecryptStream = (reader: Reader, getInit: () => MaybePromise) => { +export const createAes128CbcDecryptStream = ( + reader: Reader, + getInit: () => MaybePromise, + 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(); + }, }); }; diff --git a/src/hls/hls-segmented-input.ts b/src/hls/hls-segmented-input.ts index 9493278..e32350c 100644 --- a/src/hls/hls-segmented-input.ts +++ b/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; diff --git a/src/input.ts b/src/input.ts index e1c6ea8..dce3299 100644 --- a/src/input.ts +++ b/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 | null = null; +let inputFinalizationRegistry: FinalizationRegistry | 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 = { /** 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); + source: S | SourceRef | ((request: SourceRequest) => MaybePromise>); entryPath?: string; initInput?: Input; }; +type SourceCacheEntry = { + request: SourceRequest; + sourceRef: SourceRef; + 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 = { */ export class Input implements Disposable { /** @internal */ - _source: InputOptions['source']; + _source: SourceRef | ((request: SourceRequest) => MaybePromise>); /** @internal */ _formats: InputFormat[]; /** @internal */ @@ -84,15 +91,14 @@ export class Input 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[] = []; + /** @internal */ + _sourceCachePromises: { request: SourceRequest; - sourcePromise: Promise; - age: number; cacheGroup: number; + promise: Promise>; }[] = []; /** @@ -119,9 +125,6 @@ export class Input 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 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; + 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> { 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 = { + 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 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 { - 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 implements Disposable { ); } - return source; + if (source instanceof Source) { + return source; + } else { + return source.source; + } } /** @@ -432,10 +462,10 @@ export class Input 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); diff --git a/src/isobmff/isobmff-demuxer.ts b/src/isobmff/isobmff-demuxer.ts index 811d45a..cc68d88 100644 --- a/src/isobmff/isobmff-demuxer.ts +++ b/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); } diff --git a/src/segmented-input.ts b/src/segmented-input.ts index 665c384..2c4df88 100644 --- a/src/segmented-input.ts +++ b/src/segmented-input.ts @@ -99,6 +99,8 @@ export abstract class SegmentedInput { entry.input.dispose(); } this.inputCache.length = 0; + + this.virtualInput?.dispose(); } } diff --git a/src/source.ts b/src/source.ts index f232af5..28f63a8 100644 --- a/src/source.ts +++ b/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 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(); } } diff --git a/test/node/aes.test.ts b/test/node/aes.test.ts index 0986baf..28661ea 100644 --- a/test/node/aes.test.ts +++ b/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 = (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[] = []; diff --git a/test/node/source-lifetime.test.ts b/test/node/source-lifetime.test.ts new file mode 100644 index 0000000..f83b1b6 --- /dev/null +++ b/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); +});