mediabunny/packages/server/src/video-encoder.ts

/*!
 * Copyright (c) 2026-present, Vanilagy and contributors
 *
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/.
 */

import {
	CustomVideoEncoder,
	type MaybePromise,
	QUALITY_MEDIUM,
	VideoCodec,
	VideoSample,
	EncodedPacket,
	EncodedPacketSideData,
} from 'mediabunny';
import * as NodeAv from 'node-av';
import {
	CODEC_TO_CODEC_ID,
	getHardwareEncoderCodec,
	unmapColorPrimaries,
	unmapMatrixCoefficients,
	unmapTransferCharacteristics,
} from './misc';
import { copyVideoSampleToAvFrame, AvFrameVideoSampleResource } from './video-sample';
import {
	AvcNalUnitType,
	extractAv1CodecInfoFromPacket,
	extractAvcDecoderConfigurationRecord,
	extractHevcDecoderConfigurationRecord,
	extractNalUnitTypeForAvc,
	extractNalUnitTypeForHevc,
	extractVp9CodecInfoFromPacket,
	HevcNalUnitType,
	iterateNalUnitsInAnnexB,
	NalUnitLocation,
	serializeAvcDecoderConfigurationRecord,
	serializeHevcDecoderConfigurationRecord,
} from '../../../src/codec-data';
import { extractVideoCodecString } from '../../../src/codec';
import { assert, binarySearchLessOrEqual, simplifyRational, toUint8Array } from '../../../src/misc';

type NodeAvState = {
	frame: NodeAv.Frame;
	packet: NodeAv.Packet;
	codecContext: NodeAv.CodecContext | null;
	scaler: NodeAv.SoftwareScaleContext | null;
	dstFrame: NodeAv.Frame | null;
};

const freeState = (state: NodeAvState) => {
	state.codecContext?.freeContext();
	state.frame.free();
	state.packet.free();
	state.scaler?.freeContext();
	state.dstFrame?.free();
};

// Needed for proper freeing if close isn't called
let finalizationRegistry: FinalizationRegistry<NodeAvState> | null = null;
if (typeof FinalizationRegistry !== 'undefined') {
	finalizationRegistry = new FinalizationRegistry<NodeAvState>((state) => {
		freeState(state);
	});
}

export class NodeAvVideoEncoder extends CustomVideoEncoder {
	state!: NodeAvState;
	avCodec!: NodeAv.Codec;
	lastBuffer: Buffer | null = null;
	packetEmitted = false;
	lastScalerKey: string | null = null;

	// Bookkeeping to restore the original timing information
	preciseTimings: {
		microsecondTimestamp: number;
		timestamp: number;
		duration: number;
		timestampIsValid: boolean;
		durationIsValid: boolean;
	}[] = [];

	static override supports(codec: VideoCodec, config: VideoEncoderConfig): boolean {
		return (codec === 'avc' || codec === 'hevc' || codec === 'vp8' || codec === 'vp9' || codec === 'av1')
			&& config.bitrateMode !== 'quantizer';
	}

	async init(): Promise<void> {
		const frame = new NodeAv.Frame();
		frame.alloc();
		frame.timeBase = new NodeAv.Rational(1, 1e6);

		const packet = new NodeAv.Packet();
		packet.alloc();

		this.state = { frame, packet, codecContext: null, scaler: null, dstFrame: null };

		finalizationRegistry?.register(this, this.state, this);

		const codecId = CODEC_TO_CODEC_ID[this.codec];
		assert(codecId !== undefined);

		let codec: NodeAv.Codec | null = null;
		if (this.codec === 'vp9' && this.config.alpha === 'keep') {
			codec = NodeAv.Codec.findEncoderByName(NodeAv.FF_ENCODER_LIBVPX_VP9) ?? NodeAv.Codec.findEncoder(codecId);
		} else if (this.config.hardwareAcceleration === 'prefer-software') {
			codec = NodeAv.Codec.findEncoder(codecId);
		} else {
			codec = getHardwareEncoderCodec(codecId) ?? NodeAv.Codec.findEncoder(codecId);
		}

		if (!codec) {
			throw new Error(`Unable to obtain libav codec for '${this.codec}'.`);
		}

		this.avCodec = codec;

		await this.createCodecContext();
	}

	async createCodecContext() {
		assert(this.state.codecContext === null);

		const codecContext = new NodeAv.CodecContext();
		codecContext.allocContext3(this.avCodec);

		let pixelFormat = NodeAv.AV_PIX_FMT_YUV420P;

		if (this.avCodec.pixelFormats) {
			if (!this.avCodec.pixelFormats.includes(NodeAv.AV_PIX_FMT_YUV420P)) {
				pixelFormat = this.avCodec.pixelFormats[0]!;
			}

			if (this.config.alpha === 'keep' && this.avCodec.pixelFormats.includes(NodeAv.AV_PIX_FMT_YUVA420P)) {
				pixelFormat = NodeAv.AV_PIX_FMT_YUVA420P;
			}
		}

		const pixelAspectRatio = simplifyRational({
			num: (this.config.displayWidth ?? this.config.width) * this.config.height,
			den: (this.config.displayHeight ?? this.config.height) * this.config.width,
		});

		codecContext.width = this.config.width;
		codecContext.height = this.config.height;
		codecContext.pixelFormat = pixelFormat;
		codecContext.timeBase = new NodeAv.Rational(1, 1e6);
		codecContext.gopSize = 60;
		codecContext.framerate = new NodeAv.Rational(Math.round(this.config.framerate ?? 0) || 30, 1);
		codecContext.bitRate = BigInt(
			this.config.bitrate ?? QUALITY_MEDIUM._toVideoBitrate(this.codec, this.config.width, this.config.height),
		);
		codecContext.sampleAspectRatio = new NodeAv.Rational(pixelAspectRatio.num, pixelAspectRatio.den);

		if (this.config.bitrateMode === 'constant') {
			codecContext.rcMinRate = codecContext.bitRate;
			codecContext.rcMaxRate = codecContext.bitRate;
		}

		const isRealtime = this.config.latencyMode === 'realtime';

		if (this.avCodec.name === 'libx264') {
			if (isRealtime) {
				codecContext.setOption('tune', 'zerolatency');
				codecContext.setOption('preset', 'ultrafast');
			}
		} else if (this.avCodec.name === 'libx265') {
			codecContext.setOption('x265-params', 'log-level=error');

			if (isRealtime) {
				codecContext.setOption('tune', 'zerolatency');
				codecContext.setOption('preset', 'ultrafast');
			}
		} else if (this.avCodec.name === 'libvpx') {
			if (isRealtime) {
				codecContext.setOption('deadline', 'realtime');
				codecContext.setOption('cpu-used', '8');
			} else {
				codecContext.setOption('cpu-used', '8');
			}
		} else if (this.avCodec.name === 'libvpx-vp9') {
			codecContext.setOption('deadline', 'realtime');

			if (isRealtime) {
				codecContext.setOption('cpu-used', '8');
			} else {
				codecContext.setOption('cpu-used', '5');
			}
		} else if (this.avCodec.name === 'libsvtav1') {
			// SVTAV1 can be silenced by setting an environment variable:
			// https://superuser.com/questions/1775236/how-to-remove-svt-av1-information-from-ffmpeg-output
			process.env['SVT_LOG'] = '1';

			if (isRealtime) {
				codecContext.setOption('preset', '12');
			}
		}

		const ret = await codecContext.open2();
		NodeAv.FFmpegError.throwIfError(ret, 'Open codec context');

		this.state.codecContext = codecContext;
	}

	async encode(videoSample: VideoSample, options: VideoEncoderEncodeOptions): Promise<void> {
		if (this.state.codecContext === null) {
			await this.createCodecContext();
		}
		assert(this.state.codecContext);

		if (videoSample._data instanceof AvFrameVideoSampleResource) {
			// Release any buffers still referenced from the previous encode before reffing the new frame, otherwise
			// av_frame_ref leaks them
			// https://github.com/Vanilagy/mediabunny/issues/392
			this.state.frame.unref();
			this.state.frame.ref(videoSample._data.frame);
		} else {
			if (videoSample.format === null) {
				throw new Error('Cannot encode foreign VideoSample with unknown (null) format.');
			}

			this.lastBuffer = await copyVideoSampleToAvFrame(videoSample, this.state.frame, this.lastBuffer);
		}

		let frameToEncode = this.state.frame;

		const requiresScaler
			= this.state.codecContext.pixelFormat !== this.state.frame.format
				|| this.state.codecContext.width !== this.state.frame.width
				|| this.state.codecContext.height !== this.state.frame.height;

		if (requiresScaler) {
			if (!this.state.scaler) {
				this.state.scaler = new NodeAv.SoftwareScaleContext();
			}

			const key = `${this.state.frame.width}x${this.state.frame.height}:${this.state.frame.format}`;
			const needsConfigure = key !== this.lastScalerKey;

			if (needsConfigure) {
				this.state.scaler.getContext(
					this.state.frame.width, this.state.frame.height, this.state.frame.format as NodeAv.AVPixelFormat,
					this.state.codecContext.width, this.state.codecContext.height, this.state.codecContext.pixelFormat,
					NodeAv.SWS_FAST_BILINEAR,
				);

				this.lastScalerKey = key;

				const ret = this.state.scaler.initContext();
				NodeAv.FFmpegError.throwIfError(ret, 'initContext');
			}

			if (!this.state.dstFrame) {
				this.state.dstFrame = new NodeAv.Frame();
				this.state.dstFrame.alloc();
				this.state.dstFrame.width = this.state.codecContext.width;
				this.state.dstFrame.height = this.state.codecContext.height;
				this.state.dstFrame.format = this.state.codecContext.pixelFormat;
				this.state.dstFrame.allocBuffer();
			}

			await this.state.scaler.scaleFrame(this.state.dstFrame, this.state.frame);
			this.state.dstFrame.copyProps(this.state.frame);
			frameToEncode = this.state.dstFrame;
		}

		frameToEncode.pts = BigInt(videoSample.microsecondTimestamp);
		frameToEncode.duration = BigInt(videoSample.microsecondDuration);
		frameToEncode.timeBase = new NodeAv.Rational(1, 1e6);

		// Let's just set both for good measure
		frameToEncode.pictType = options?.keyFrame
			? NodeAv.AV_PICTURE_TYPE_I
			: NodeAv.AV_PICTURE_TYPE_NONE;
		frameToEncode.keyFrame = options?.keyFrame
			? 1
			: 0;

		const preciseTimingIndex = binarySearchLessOrEqual(
			this.preciseTimings,
			videoSample.microsecondTimestamp,
			x => x.microsecondTimestamp,
		);
		const existingEntry = preciseTimingIndex !== -1
			? this.preciseTimings[preciseTimingIndex]
			: null;
		if (existingEntry && existingEntry.microsecondTimestamp === videoSample.microsecondTimestamp) {
			if (existingEntry.timestamp !== videoSample.timestamp) {
				// Mapping isn't unique, can't use the timestamp
				existingEntry.timestampIsValid = false;
			}
			if (existingEntry.duration !== videoSample.duration) {
				// Mapping isn't unique, can't use the duration
				existingEntry.durationIsValid = false;
			}
		} else {
			this.preciseTimings.splice(preciseTimingIndex + 1, 0, {
				microsecondTimestamp: videoSample.microsecondTimestamp,
				timestamp: videoSample.timestamp,
				duration: videoSample.duration,
				timestampIsValid: true,
				durationIsValid: true,
			});

			// Make sure it doesn't grow indefinitely
			if (this.preciseTimings.length > 128) {
				this.preciseTimings.shift();
			}
		}

		const ret = await this.state.codecContext.sendFrame(frameToEncode);
		NodeAv.FFmpegError.throwIfError(ret, 'Send frame');

		// Keep receiving packets until no more are available for this frame
		while (true) {
			const receiveRet = await this.state.codecContext.receivePacket(this.state.packet);
			if (receiveRet === NodeAv.AVERROR_EAGAIN || receiveRet === NodeAv.AVERROR_EOF) {
				break;
			}

			this.receivePacket(receiveRet);
		}
	}

	receivePacket(ret: number) {
		assert(this.state.codecContext);
		NodeAv.FFmpegError.throwIfError(ret, 'Receive packet');

		if (!this.state.packet.data) {
			return;
		}
		let packetData = toUint8Array(this.state.packet.data);

		let timestamp = Number(this.state.packet.pts) / 1e6;
		let duration = Number(this.state.packet.duration) / 1e6;

		const preciseTimingIndex = binarySearchLessOrEqual(
			this.preciseTimings,
			Number(this.state.packet.pts),
			x => x.microsecondTimestamp,
		);
		const entry = preciseTimingIndex !== -1
			? this.preciseTimings[preciseTimingIndex]
			: null;

		// If there's a relevant timing entry, refine the packet's timing data to get better accuracy than
		// microseconds
		if (entry && entry.microsecondTimestamp === Number(this.state.packet.pts)) {
			if (entry.timestampIsValid) {
				timestamp = entry.timestamp;
			}
			if (entry.durationIsValid) {
				duration = entry.duration;
			}
		}

		const metadata: EncodedVideoChunkMetadata = {};
		let decoderConfigCodecString: string | null = null;
		let decoderConfigDescription: Uint8Array | null = null;

		if (this.codec === 'avc' || this.codec === 'hevc') {
			let expectsAnnexB = false;
			if (this.codec === 'avc') {
				expectsAnnexB = this.config.avc?.format === 'annexb';
			} else {
				// eslint-disable-next-line @stylistic/max-len
				// eslint-disable-next-line @typescript-eslint/no-explicit-any, @typescript-eslint/no-unsafe-member-access
				expectsAnnexB = (this.config as any).hevc?.format === 'annexb';
			}

			if (!this.packetEmitted) {
				let serializedRecord: Uint8Array;

				if (this.codec === 'avc') {
					const record = extractAvcDecoderConfigurationRecord(this.state.packet.data);
					if (!record) {
						throw new Error('Invalid AVC data, could not extract decoder configuration record.');
					}

					serializedRecord = serializeAvcDecoderConfigurationRecord(record);
				} else {
					const record = extractHevcDecoderConfigurationRecord(this.state.packet.data);
					if (!record) {
						throw new Error('Invalid HEVC data, could not extract decoder configuration record.');
					}

					serializedRecord = serializeHevcDecoderConfigurationRecord(record);
				}

				decoderConfigCodecString = extractVideoCodecString({
					width: this.config.width,
					height: this.config.height,
					codec: this.codec,
					codecDescription: serializedRecord,
					colorSpace: null,
					avcType: 1,
					avcCodecInfo: null,
					hevcCodecInfo: null,
					vp9CodecInfo: null,
					av1CodecInfo: null,
				});

				if (!expectsAnnexB) {
					decoderConfigDescription = serializedRecord;
				}
			}

			if (!expectsAnnexB) {
				const NAL_UNIT_LENGTH_SIZE = 4;

				const nalUnits: NalUnitLocation[] = [];
				for (const loc of iterateNalUnitsInAnnexB(packetData)) {
					if (this.codec === 'avc') {
						const naluType = extractNalUnitTypeForAvc(packetData[loc.offset]!);

						// Certain NALUs get stripped
						if (
							naluType !== AvcNalUnitType.SPS
							&& naluType !== AvcNalUnitType.PPS
							&& naluType !== AvcNalUnitType.SPS_EXT
						) {
							nalUnits.push(loc);
						}
					} else {
						const naluType = extractNalUnitTypeForHevc(packetData[loc.offset]!);

						// Certain NALUs get stripped
						if (
							naluType !== HevcNalUnitType.SPS_NUT
							&& naluType !== HevcNalUnitType.PPS_NUT
							&& naluType !== HevcNalUnitType.VPS_NUT
						) {
							nalUnits.push(loc);
						}
					}
				}

				let totalSize = 0;
				for (const nalUnit of nalUnits) {
					totalSize += NAL_UNIT_LENGTH_SIZE + nalUnit.length;
				}

				const lengthPrefixedData = new Uint8Array(totalSize);
				const dataView = new DataView(lengthPrefixedData.buffer);
				let offset = 0;

				// Write each NAL unit with its length prefix
				for (const nalUnit of nalUnits) {
					const length = nalUnit.length;

					dataView.setUint32(offset, length, false);
					offset += 4;

					lengthPrefixedData.set(
						packetData.subarray(nalUnit.offset, nalUnit.offset + nalUnit.length),
						offset,
					);
					offset += nalUnit.length;
				}

				packetData = lengthPrefixedData;
			}
		} else if (this.codec === 'vp8') {
			if (!this.packetEmitted) {
				decoderConfigCodecString = extractVideoCodecString({
					width: this.config.width,
					height: this.config.height,
					codec: 'vp8',
					codecDescription: null,
					colorSpace: null,
					avcType: null,
					avcCodecInfo: null,
					hevcCodecInfo: null,
					vp9CodecInfo: null,
					av1CodecInfo: null,
				});
			}
		} else if (this.codec === 'vp9') {
			if (!this.packetEmitted) {
				const vp9CodecInfo = extractVp9CodecInfoFromPacket(packetData);

				decoderConfigCodecString = extractVideoCodecString({
					width: this.config.width,
					height: this.config.height,
					codec: 'vp9',
					codecDescription: null,
					colorSpace: null,
					avcType: null,
					avcCodecInfo: null,
					hevcCodecInfo: null,
					vp9CodecInfo,
					av1CodecInfo: null,
				});
			}
		} else if (this.codec === 'av1') {
			if (!this.packetEmitted) {
				const av1CodecInfo = extractAv1CodecInfoFromPacket(packetData);

				decoderConfigCodecString = extractVideoCodecString({
					width: this.config.width,
					height: this.config.height,
					codec: 'av1',
					codecDescription: null,
					colorSpace: null,
					avcType: null,
					avcCodecInfo: null,
					hevcCodecInfo: null,
					vp9CodecInfo: null,
					av1CodecInfo,
				});
			}
		} else {
			throw new Error('Unreachable.');
		}

		const sideData: EncodedPacketSideData = {};
		const matroskaBlockAdditional = this.state.packet.getSideData(NodeAv.AV_PKT_DATA_MATROSKA_BLOCKADDITIONAL);
		if (matroskaBlockAdditional) {
			sideData.alpha = toUint8Array(matroskaBlockAdditional).subarray(8); // Skip the BlockAddId
		}

		const packet = new EncodedPacket(
			packetData,
			this.state.packet.isKeyframe ? 'key' : 'delta',
			timestamp,
			duration,
			undefined,
			undefined,
			sideData,
		);

		if (decoderConfigCodecString !== null) {
			// Create the decoder config
			metadata.decoderConfig = {
				codec: decoderConfigCodecString,
				codedWidth: this.state.codecContext.width,
				codedHeight: this.state.codecContext.height,
				displayAspectWidth: this.config.displayWidth ?? this.state.codecContext.width,
				displayAspectHeight: this.config.displayHeight ?? this.state.codecContext.height,
				description: decoderConfigDescription ?? undefined,
				colorSpace: {
					primaries: unmapColorPrimaries(this.state.codecContext.colorPrimaries) as VideoColorPrimaries,
					matrix: unmapMatrixCoefficients(this.state.codecContext.colorSpace) as VideoMatrixCoefficients,
					transfer:
						unmapTransferCharacteristics(this.state.codecContext.colorTrc) as VideoTransferCharacteristics,
					fullRange: this.state.codecContext.colorRange === NodeAv.AVCOL_RANGE_JPEG
						? true
						: this.state.codecContext.colorRange === NodeAv.AVCOL_RANGE_MPEG
							? false
							: undefined,
				},
			};
		}

		this.packetEmitted = true;
		this.onPacket(packet, metadata);
	}

	async flush(): Promise<void> {
		if (this.state.codecContext) {
			// Send null frame to signal flush
			const ret = await this.state.codecContext.sendFrame(null);
			NodeAv.FFmpegError.throwIfError(ret, 'Send frame');

			// Keep receiving packets until no more are available
			while (true) {
				const receiveRet = await this.state.codecContext.receivePacket(this.state.packet);
				if (receiveRet === NodeAv.AVERROR_EAGAIN || receiveRet === NodeAv.AVERROR_EOF) {
					break;
				}

				this.receivePacket(receiveRet);
			}

			this.state.codecContext.freeContext();
			this.state.codecContext = null;
			// The codec is done now and can't be reused. Any subsequent encode call will first need to recreate a
			// codec context.
		}

		this.packetEmitted = false;
	}

	close(): MaybePromise<void> {
		finalizationRegistry?.unregister(this);
		freeState(this.state);
	}
}