From ac1b009486cf2821378d0769076ef7e5d11ca590 Mon Sep 17 00:00:00 2001
From: Vanilagy <1696106+Vanilagy@users.noreply.github.com>
Date: Mon, 28 Apr 2025 18:31:53 +0200
Subject: [PATCH] Implement "cross-track offset" for MediaStreamTrack sources &
 small fixes

---
 dev/live.html       | 50 ++++++++++++++++++++++++++++++++++++++++++
 src/media-source.ts | 53 +++++++++++++++++++++++++++++++++++----------
 src/muxer.ts        | 19 +++++++++-------
 todo.txt            |  1 -
 4 files changed, 102 insertions(+), 21 deletions(-)
 create mode 100644 dev/live.html

diff --git a/dev/live.html b/dev/live.html
new file mode 100644
index 0000000..73e2a4b
--- /dev/null
+++ b/dev/live.html
@@ -0,0 +1,50 @@
+<button>Go</button>
+
+<script src="../dist/metamuxer.js"></script>
+
+<script type="module">
+    function download(blob, filename) {
+		const url = URL.createObjectURL(blob);
+		const a = document.createElement('a');
+		a.href = url;
+		a.download = filename;
+		a.click();
+		URL.revokeObjectURL(url);
+	}
+
+    const button = document.querySelector('button');
+    button.addEventListener('click', async () => {
+        const stream = await navigator.mediaDevices.getUserMedia({ video: true, audio: true });
+        const videoTrack = stream.getVideoTracks()[0];
+        const audioTrack = stream.getAudioTracks()[0];
+
+        const output = new Metamuxer.Output({
+            target: new Metamuxer.BufferTarget(),
+            format: new Metamuxer.Mp4OutputFormat(),
+        });
+        if (videoTrack) {
+            output.addVideoTrack(new Metamuxer.MediaStreamVideoTrackSource(videoTrack, {
+                codec: 'avc',
+                bitrate: Metamuxer.QUALITY_MEDIUM
+            }));
+        }
+        if (audioTrack) {
+            output.addAudioTrack(new Metamuxer.MediaStreamAudioTrackSource(audioTrack, {
+                codec: 'aac',
+                bitrate: Metamuxer.QUALITY_MEDIUM
+            }));
+        }
+
+        await output.start();
+
+        await new Promise(resolve => setTimeout(resolve, 3000));
+
+        await output.finalize();
+
+        console.log(output.target.buffer);
+        download(new Blob([output.target.buffer]), 'livetest' + output.format.fileExtension);
+
+        videoTrack?.stop();
+        audioTrack?.stop();
+    });
+</script>
\ No newline at end of file
diff --git a/src/media-source.ts b/src/media-source.ts
index 719ac57..58e83e5 100644
--- a/src/media-source.ts
+++ b/src/media-source.ts
@@ -40,8 +40,11 @@ export abstract class MediaSource {
 	_closingPromise: Promise<void> | null = null;
 	/** @internal */
 	_closed = false;
-	/** @internal */
-	_offsetTimestamps = false;
+	/**
+	 * @internal
+	 * A time offset in seconds that is added to all timestamps generated by this source.
+	 */
+	_timestampOffset = 0;
 
 	/** @internal */
 	_ensureValidAdd() {
@@ -431,8 +434,7 @@ class VideoEncoderWrapper {
 		if (this.customEncoder) {
 			return this.customEncoderQueueSize;
 		} else {
-			assert(this.encoder);
-			return this.encoder.encodeQueueSize;
+			return this.encoder?.encodeQueueSize ?? 0;
 		}
 	}
 }
@@ -540,9 +542,6 @@ export class MediaStreamVideoTrackSource extends VideoSource {
 	/** @internal */
 	private _track: MediaStreamVideoTrack;
 
-	/** @internal */
-	override _offsetTimestamps = true;
-
 	constructor(track: MediaStreamVideoTrack, encodingConfig: VideoEncodingConfig) {
 		if (!(track instanceof MediaStreamTrack) || track.kind !== 'video') {
 			throw new TypeError('track must be a video MediaStreamTrack.');
@@ -563,9 +562,26 @@ export class MediaStreamVideoTrackSource extends VideoSource {
 	override _start() {
 		this._abortController = new AbortController();
 
+		let frameReceived = false;
+
 		const processor = new MediaStreamTrackProcessor({ track: this._track });
 		const consumer = new WritableStream<VideoFrame>({
 			write: (videoFrame) => {
+				const timestampInSeconds = videoFrame.timestamp / 1e6;
+
+				assert(this._connectedTrack);
+				const muxer = this._connectedTrack.output._muxer;
+				if (muxer.firstMediaStreamTimestamp === null) {
+					// We're the first MediaStreamTrack of this output to receive data
+					muxer.firstMediaStreamTimestamp = timestampInSeconds;
+				}
+
+				if (!frameReceived) {
+					// Math.min to ensure the timestamps can't get negative
+					this._timestampOffset = -Math.min(muxer.firstMediaStreamTimestamp, timestampInSeconds);
+					frameReceived = true;
+				}
+
 				if (this._encoder.getQueueSize() >= 4) {
 					// Drop frames if the encoder is overloaded
 					videoFrame.close();
@@ -1053,8 +1069,7 @@ class AudioEncoderWrapper {
 		} else if (this.isPcmEncoder) {
 			return 0;
 		} else {
-			assert(this.encoder);
-			return this.encoder.encodeQueueSize;
+			return this.encoder?.encodeQueueSize ?? 0;
 		}
 	}
 }
@@ -1191,9 +1206,6 @@ export class MediaStreamAudioTrackSource extends AudioSource {
 	/** @internal */
 	private _track: MediaStreamAudioTrack;
 
-	/** @internal */
-	override _offsetTimestamps = true;
-
 	constructor(track: MediaStreamAudioTrack, encodingConfig: AudioEncodingConfig) {
 		if (!(track instanceof MediaStreamTrack) || track.kind !== 'audio') {
 			throw new TypeError('track must be an audio MediaStreamTrack.');
@@ -1209,9 +1221,26 @@ export class MediaStreamAudioTrackSource extends AudioSource {
 	override _start() {
 		this._abortController = new AbortController();
 
+		let dataReceived = false;
+
 		const processor = new MediaStreamTrackProcessor({ track: this._track });
 		const consumer = new WritableStream<AudioData>({
 			write: (audioData) => {
+				const timestampInSeconds = audioData.timestamp / 1e6;
+
+				assert(this._connectedTrack);
+				const muxer = this._connectedTrack.output._muxer;
+				if (muxer.firstMediaStreamTimestamp === null) {
+					// We're the first MediaStreamTrack of this output to receive data
+					muxer.firstMediaStreamTimestamp = timestampInSeconds;
+				}
+
+				if (!dataReceived) {
+					// Math.min to ensure the timestamps can't get negative
+					this._timestampOffset = -Math.min(muxer.firstMediaStreamTimestamp, timestampInSeconds);
+					dataReceived = true;
+				}
+
 				if (this._encoder.getQueueSize() >= 4) {
 					// Drop data if the encoder is overloaded
 					audioData.close();
diff --git a/src/muxer.ts b/src/muxer.ts
index 7add20a..77b88ba 100644
--- a/src/muxer.ts
+++ b/src/muxer.ts
@@ -7,6 +7,13 @@ export abstract class Muxer {
 	output: Output;
 	mutex = new AsyncMutex();
 
+	/**
+	 * This field is used to synchronize multiple MediaStreamTracks. They use the same time coordinate system across
+	 * tracks, and to ensure correct audio-video sync, we must use the same offset for all of them. The reason an offset
+	 * is needed at all is because the timestamps typically don't start at zero.
+	 */
+	firstMediaStreamTimestamp: number | null = null;
+
 	constructor(output: Output) {
 		this.output = output;
 	}
@@ -29,12 +36,13 @@ export abstract class Muxer {
 	onTrackClose(track: OutputTrack) {}
 
 	private trackTimestampInfo = new WeakMap<OutputTrack, {
-		timestampOffset: number;
 		maxTimestamp: number;
 		maxTimestampBeforeLastKeyFrame: number;
 	}>();
 
 	protected validateAndNormalizeTimestamp(track: OutputTrack, timestampInSeconds: number, isKeyFrame: boolean) {
+		timestampInSeconds += track.source._timestampOffset;
+
 		let timestampInfo = this.trackTimestampInfo.get(track);
 		if (!timestampInfo) {
 			if (!isKeyFrame) {
@@ -42,17 +50,12 @@ export abstract class Muxer {
 			}
 
 			timestampInfo = {
-				timestampOffset: timestampInSeconds,
-				maxTimestamp: track.source._offsetTimestamps ? 0 : timestampInSeconds,
-				maxTimestampBeforeLastKeyFrame: track.source._offsetTimestamps ? 0 : timestampInSeconds,
+				maxTimestamp: timestampInSeconds,
+				maxTimestampBeforeLastKeyFrame: timestampInSeconds,
 			};
 			this.trackTimestampInfo.set(track, timestampInfo);
 		}
 
-		if (track.source._offsetTimestamps) {
-			timestampInSeconds -= timestampInfo.timestampOffset;
-		}
-
 		if (timestampInSeconds < 0) {
 			throw new Error(`Timestamps must be non-negative (got ${timestampInSeconds}s).`);
 		}
diff --git a/todo.txt b/todo.txt
index 3f5dfc6..0bb2534 100644
--- a/todo.txt
+++ b/todo.txt
@@ -1,4 +1,3 @@
 - https://github.com/Vanilagy/mp4-muxer/issues/83 tell him it's possible now
-- cross-track offset for streaming sources
 - textsubtitlesource, chunked piping
 - More efficient MP3 loading when reading sequentially
\ No newline at end of file