From 141bed1b378e617b65c6076cc84a4eeac91c0ae2 Mon Sep 17 00:00:00 2001
From: Vanilagy <1696106+Vanilagy@users.noreply.github.com>
Date: Thu, 2 Jan 2025 19:57:32 +0100
Subject: [PATCH] Add Output.cancel

---
 src/media-source.ts | 14 ++++++++++++--
 src/output.ts       | 24 ++++++++++++++++++++++++
 src/writer.ts       | 12 ++++++++++++
 todo.txt            |  3 +--
 4 files changed, 49 insertions(+), 4 deletions(-)

diff --git a/src/media-source.ts b/src/media-source.ts
index 5199ee0..5e1a659 100644
--- a/src/media-source.ts
+++ b/src/media-source.ts
@@ -20,6 +20,8 @@ export abstract class MediaSource {
 	/** @internal */
 	_connectedTrack: OutputTrack | null = null;
 	/** @internal */
+	_closing = false;
+	/** @internal */
 	_closed = false;
 	/** @internal */
 	_offsetTimestamps = false;
@@ -30,6 +32,10 @@ export abstract class MediaSource {
 			throw new Error('Cannot call digest without connecting the source to an output track.');
 		}
 
+		if (this._connectedTrack.output._canceled) {
+			throw new Error('Cannot call digest after output has been canceled.');
+		}
+
 		if (!this._connectedTrack.output._started) {
 			throw new Error('Cannot call digest before output has been started.');
 		}
@@ -48,8 +54,8 @@ export abstract class MediaSource {
 	/** @internal */
 	async _flush() {}
 
-	close() {
-		if (this._closed) {
+	async close() {
+		if (this._closing) {
 			throw new Error('Source already closed.');
 		}
 
@@ -61,6 +67,10 @@ export abstract class MediaSource {
 			throw new Error('Cannot call close before output has been started.');
 		}
 
+		this._closing = true;
+
+		await this._flush();
+
 		this._closed = true;
 
 		if (this._connectedTrack.output._finalizing) {
diff --git a/src/output.ts b/src/output.ts
index 77a76d7..64fadef 100644
--- a/src/output.ts
+++ b/src/output.ts
@@ -53,6 +53,8 @@ export class Output {
 	/** @internal */
 	_started = false;
 	/** @internal */
+	_canceled = false;
+	/** @internal */
 	_finalizing = false;
 	/** @internal */
 	_mutex = new AsyncMutex();
@@ -150,6 +152,9 @@ export class Output {
 	}
 
 	async start() {
+		if (this._canceled) {
+			throw new Error('Output has been canceled.');
+		}
 		if (this._started) {
 			throw new Error('Output already started.');
 		}
@@ -168,6 +173,25 @@ export class Output {
 		release();
 	}
 
+	async cancel() {
+		if (this._finalizing) {
+			throw new Error('Cannot cancel after calling finalize.');
+		}
+		if (this._canceled) {
+			throw new Error('Output already canceled.');
+		}
+		this._canceled = true;
+
+		const release = await this._mutex.acquire();
+
+		const promises = this._tracks.map(x => x.source._flush());
+		await Promise.all(promises);
+
+		await this._writer.close();
+
+		release();
+	}
+
 	async finalize() {
 		if (!this._started) {
 			throw new Error('Cannot finalize before starting.');
diff --git a/src/writer.ts b/src/writer.ts
index b1056ee..2fcaf62 100644
--- a/src/writer.ts
+++ b/src/writer.ts
@@ -17,6 +17,8 @@ export abstract class Writer {
 	abstract flush(): Promise<void>;
 	/** Called after muxing has finished. */
 	abstract finalize(): Promise<void>;
+	/** Closes the writer. */
+	abstract close(): Promise<void>;
 }
 
 /**
@@ -74,6 +76,8 @@ export class ArrayBufferTargetWriter extends Writer {
 		this.target.buffer = this.buffer.slice(0, Math.max(this.maxPos, this.pos));
 	}
 
+	async close() {}
+
 	getSlice(start: number, end: number) {
 		return this.bytes.slice(start, end);
 	}
@@ -185,6 +189,10 @@ export class StreamTargetWriter extends Writer {
 		assert(this.writer);
 		return this.writer.close();
 	}
+
+	async close() {
+		return this.writer?.close();
+	}
 }
 
 const DEFAULT_CHUNK_SIZE = 2 ** 24;
@@ -375,4 +383,8 @@ export class ChunkedStreamTargetWriter extends Writer {
 
 		return this.writer.close();
 	}
+
+	async close() {
+		return this.writer?.close();
+	}
 }
diff --git a/todo.txt b/todo.txt
index 0e57799..e59ea96 100644
--- a/todo.txt
+++ b/todo.txt
@@ -4,5 +4,4 @@
 - Audio codec string regex validation
 - Add the new audio codecs to muxers
 - Mov muxer!!! Only mov can hold PCM audio, MP4 cannot
-- Metadata methods for computing average fps and bitrate
-- close method on Output to cancel it (and dispose the encoders). Perhaps abort is a better title
\ No newline at end of file
+- Metadata methods for computing average fps and bitrate
\ No newline at end of file