Implement "cross-track offset" for MediaStreamTrack sources & small fixes

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();

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).`);
 		}