Rewrite mediaSamplesAtTimestamps to fix B-frame bug, fix incorrect timestamps when cloning samples

dev/demux.html

@@ -16,8 +16,21 @@
 			source
 		});
 
-		/*
+		let timestamps = [];
-		const timestamps = [
+		const videoTrack = await input.getPrimaryVideoTrack();
+
+		const packetSink = new Metamuxer.EncodedPacketSink(videoTrack);
+		for await (const packet of packetSink.packets()) {
+			if (packet.timestamp >= 1) {
+				break;
+			}
+
+			timestamps.push(packet.timestamp);
+		}
+		timestamps.push(11);
+
+		timestamps.sort((a, b) => a-b)
+		timestamps = [
 			4.16,
 			4.266666666666667,
 			4.373333333333333,
@@ -49,20 +62,6 @@
 			7.1466666666666665,
 			7.253333333333333
 		];
-		*/
-		const timestamps = [];
-		const videoTrack = await input.getPrimaryVideoTrack();
-
-		const packetSink = new Metamuxer.EncodedPacketSink(videoTrack);
-		for await (const packet of packetSink.packets()) {
-			if (packet.timestamp >= 1) {
-				break;
-			}
-
-			timestamps.push(packet.timestamp);
-		}
-
-		timestamps.sort((a, b) => a-b)
 
 		console.log(timestamps);
 		//return;
@@ -73,12 +72,38 @@
 		}
 		*/
 
+		const sink = new Metamuxer.VideoSampleSink(videoTrack);
+		for await (const sample of sink.samplesAtTimestamps(timestamps)) {
+			console.log("wee", sample.timestamp);
+			sample.close();
+		}
+		console.log("don")
 
+		/*
 		const sink = new Metamuxer.CanvasSink(videoTrack, { width: 320 });
-
 		for await (const wrappedCanvas of sink.canvasesAtTimestamps(timestamps)) {
 			console.log(wrappedCanvas);
 		}
+		console.log("don")
+		*/
+
+		
+		/*
+		
+
+		for (let i = 0; i < 1000; i++) {
+			const count = Math.floor(50 * Math.random());
+			const timestamps = Array.from({ length: count }, () => Math.random() * 20).sort((a, b) => a-b);
+
+			for await (const wrappedCanvas of sink.canvasesAtTimestamps(timestamps)) {
+				//console.log(wrappedCanvas);
+			}
+			console.log("don")
+		}
+
+		console.log("DONE!")
+		*/
+		
 			
 
 		/*

src/media-sink.ts

@@ -504,12 +504,21 @@ export abstract class BaseMediaSampleSink<
 			});
 
 			const packetSink = this._createPacketSink();
-			let lastKeyPacket: EncodedPacket | null = null;
 			let lastPacket: EncodedPacket | null = null;
+			let lastKeyPacket: EncodedPacket | null = null;
 
-			for await (const timestamp of timestampIterator) {
+			// The end sequence number (inclusive) in the next batch of packets that will be decoded. The batch starts
-				validateTimestamp(timestamp);
+			// at the last key frame and goes until this sequence number.
+			let maxSequenceNumber = -1;
 
+			const decodePackets = async () => {
+				assert(lastKeyPacket);
+
+				// Start at the current key packet
+				let currentPacket = lastKeyPacket;
+				decoder.decode(currentPacket);
+
+				while (currentPacket.sequenceNumber < maxSequenceNumber) {
 					const maxQueueSize = computeMaxQueueSize(sampleQueue.length);
 					while (sampleQueue.length + decoder.getDecodeQueueSize() > maxQueueSize && !terminated) {
 						({ promise: queueDequeue, resolve: onQueueDequeue } = promiseWithResolvers());
@@ -520,67 +529,78 @@ export abstract class BaseMediaSampleSink<
 						break;
 					}
 
-				const targetPacket = await packetSink.getPacket(timestamp);
+					const nextPacket = await packetSink.getNextPacket(currentPacket);
-				if (!targetPacket) {
-					pushToQueue(null);
-					continue;
-				}
-
-				// console.log('target', targetPacket);
-
-				const keyPacket = await packetSink.getKeyPacket(timestamp);
-				if (!keyPacket) {
-					pushToQueue(null);
-					continue;
-				}
-
-				if (lastPacket && targetPacket.sequenceNumber < lastPacket.sequenceNumber) {
-					// console.log('FLUSH?', lastPacket, targetPacket.sequenceNumber, lastPacket.sequenceNumber);
-					// We're going back in time with this one, let's flush and reset to a clean state
-					await decoder.flush();
-					timestampsOfInterest.length = 0;
-				}
-
-				if (
-					lastKeyPacket
-					&& keyPacket.sequenceNumber === lastKeyPacket.sequenceNumber // => they're the same packet
-					// && targetPacket.timestamp >= lastPacket!.timestamp
-					&& targetPacket.sequenceNumber >= lastPacket!.sequenceNumber
-				) {
-					assert(lastPacket);
-
-					if (
-						targetPacket.sequenceNumber === lastPacket.sequenceNumber
-						&& timestampsOfInterest.length === 0
-					) {
-						// Special case: We have a repeat packet, but the sample for that packet has already been
-						// decoded. Therefore, we need to push the sample here instead of in the decoder callback.
-						if (lastUsedSample) {
-							pushToQueue(lastUsedSample.clone() as MediaSample);
-						}
-					} else {
-						timestampsOfInterest.push(targetPacket.timestamp);
-					}
-				} else {
-					// console.log('==== RESET');
-					// The key packet has changed
-					lastKeyPacket = keyPacket;
-					lastPacket = keyPacket;
-					decoder.decode(keyPacket);
-					timestampsOfInterest.push(targetPacket.timestamp);
-				}
-
-				while (lastPacket.sequenceNumber < targetPacket.sequenceNumber) {
-					const nextPacket = await packetSink.getNextPacket(lastPacket);
 					assert(nextPacket);
 
-					lastPacket = nextPacket;
+					currentPacket = nextPacket;
 					decoder.decode(nextPacket);
 				}
+
+				maxSequenceNumber = -1;
+			};
+
+			const flushDecoder = async () => {
+				await decoder.flush();
+
+				// We don't expect this list to have any elements in it anymore, but in case it does, let's emit
+				// nulls for every remaining element, then clear it.
+				for (let i = 0; i < timestampsOfInterest.length; i++) {
+					pushToQueue(null);
+				}
+				timestampsOfInterest.length = 0;
+			};
+
+			for await (const timestamp of timestampIterator) {
+				validateTimestamp(timestamp);
+
+				if (terminated) {
+					break;
+				}
+
+				const targetPacket = await packetSink.getPacket(timestamp);
+				const keyPacket = targetPacket && await packetSink.getKeyPacket(timestamp);
+
+				if (!keyPacket) {
+					if (maxSequenceNumber !== -1) {
+						await decodePackets();
+						await flushDecoder();
+					}
+
+					pushToQueue(null);
+					lastPacket = null;
+					continue;
+				}
+
+				// Check if the key packet has changed or if we're going back in time
+				if (
+					lastPacket
+					&& (
+						keyPacket.sequenceNumber !== lastKeyPacket!.sequenceNumber
+						|| targetPacket.timestamp < lastPacket.timestamp
+					)
+				) {
+					await decodePackets();
+
+					if (targetPacket.timestamp < lastPacket.timestamp) {
+						// We're going back in time with this one, let's flush and reset to a clean state
+						await flushDecoder();
+					}
+				}
+
+				timestampsOfInterest.push(targetPacket.timestamp);
+				maxSequenceNumber = Math.max(targetPacket.sequenceNumber, maxSequenceNumber);
+
+				lastPacket = targetPacket;
+				lastKeyPacket = keyPacket;
 			}
 
 			if (!terminated) {
-				await decoder.flush();
+				if (maxSequenceNumber !== -1) {
+					// We still need to decode packets
+					await decodePackets();
+				}
+
+				await flushDecoder();
 				lastUsedSample?.close();
 			}
 			decoder.close();

src/packet.ts

@@ -28,11 +28,10 @@ export class EncodedPacket {
 		/** The duration of this packet in seconds. */
 		public readonly duration: number,
 		/**
-		 * The sequence number of this packet. The sequence number indicates the decode order of the packets. Packet A
+		 * The sequence number indicates the decode order of the packets. Packet A  must be decoded before packet B if A
-		 * must be decoded before packet B if A has a lower sequence number than B. If two packets have the same
+		 * has a lower sequence number than B. If two packets have the same sequence number, they are the same packet.
-		 * sequence number, they are the same packet. Otherwise, sequence numbers are arbitrary and are not guaranteed
+		 * Otherwise, sequence numbers are arbitrary and are not guaranteed to have any meaning besides their relative
-		 * to have any meaning besides their relative ordering. Negative sequence numbers mean the sequence number
+		 * ordering. Negative sequence numbers mean the sequence number is undefined.
-		 * is undefined.
 		 */
 		public readonly sequenceNumber = -1,
 		/**

src/sample.ts

@@ -225,7 +225,10 @@ export class VideoSample {
 		assert(this._data !== null);
 
 		if (isVideoFrame(this._data)) {
-			return new VideoSample(this._data.clone());
+			return new VideoSample(this._data.clone(), {
+				timestamp: this.timestamp,
+				duration: this.duration,
+			});
 		} else if (this._data instanceof Uint8Array) {
 			return new VideoSample(this._data.slice(), {
 				format: this.format!,
@@ -812,10 +815,12 @@ export class AudioSample {
 		}
 
 		if (isAudioData(this._data)) {
-			return new AudioSample(this._data.clone());
+			const sample = new AudioSample(this._data.clone());
+			sample.setTimestamp(this.timestamp); // Make sure the timestamp is precise (beyond microsecond accuracy)
+
+			return sample;
 		} else {
 			return new AudioSample({
-
 				format: this.format,
 				sampleRate: this.sampleRate,
 				numberOfFrames: this.numberOfFrames,