index.ts

import * as Moq from "@kixelated/moq";
import { Effect, type Getter, Signal } from "@kixelated/signals";
import type * as Catalog from "../../catalog";
import { u53, u8 } from "../../catalog/integers";
import * as Container from "../../container";
import { Captions, type CaptionsProps } from "./captions";
import type * as Capture from "./capture";
import { Speaking, type SpeakingProps } from "./speaking";

export * from "./captions";

const GAIN_MIN = 0.001;
const FADE_TIME = 0.2;

export type AudioConstraints = Omit<
	MediaTrackConstraints,
	"aspectRatio" | "backgroundBlur" | "displaySurface" | "facingMode" | "frameRate" | "height" | "width"
>;

// Stronger typing for the MediaStreamTrack interface.
export interface AudioTrack extends MediaStreamTrack {
	kind: "audio";
	clone(): AudioTrack;
}

// MediaTrackSettings can represent both audio and video, which means a LOT of possibly undefined properties.
// This is a fork of the MediaTrackSettings interface with properties required for audio or vidfeo.
export interface AudioTrackSettings {
	deviceId: string;
	groupId: string;

	autoGainControl: boolean;
	channelCount: number;
	echoCancellation: boolean;
	noiseSuppression: boolean;
	sampleRate: number;
	sampleSize: number;
}

// The initial values for our signals.
export type AudioProps = {
	enabled?: boolean;
	media?: AudioTrack;
	constraints?: AudioConstraints;

	muted?: boolean;
	volume?: number;
	captions?: CaptionsProps;
	speaking?: SpeakingProps;

	// The size of each group. Larger groups mean fewer drops but the viewer can fall further behind.
	// NOTE: Each frame is always flushed to the network immediately.
	maxLatency?: DOMHighResTimeStamp;
};

export class Audio {
	broadcast: Moq.BroadcastProducer;
	enabled: Signal<boolean>;

	muted: Signal<boolean>;
	volume: Signal<number>;
	captions: Captions;
	speaking: Speaking;
	maxLatency: DOMHighResTimeStamp;

	media: Signal<AudioTrack | undefined>;
	constraints: Signal<AudioConstraints | undefined>;

	#catalog = new Signal<Catalog.Audio | undefined>(undefined);
	readonly catalog: Getter<Catalog.Audio | undefined> = this.#catalog;

	#config = new Signal<Catalog.AudioConfig | undefined>(undefined);
	readonly config: Getter<Catalog.AudioConfig | undefined> = this.#config;

	#worklet = new Signal<AudioWorkletNode | undefined>(undefined);

	#gain = new Signal<GainNode | undefined>(undefined);
	readonly root: Getter<AudioNode | undefined> = this.#gain;

	#track = new Moq.TrackProducer("audio", 1);
	#group?: Moq.GroupProducer;
	#groupTimestamp = 0;

	#signals = new Effect();

	constructor(broadcast: Moq.BroadcastProducer, props?: AudioProps) {
		this.broadcast = broadcast;
		this.media = new Signal(props?.media);
		this.enabled = new Signal(props?.enabled ?? false);
		this.speaking = new Speaking(this, props?.speaking);
		this.captions = new Captions(this, props?.captions);
		this.constraints = new Signal(props?.constraints);
		this.muted = new Signal(props?.muted ?? false);
		this.volume = new Signal(props?.volume ?? 1);
		this.maxLatency = props?.maxLatency ?? 100; // Default is a group every 100ms

		this.#signals.effect(this.#runSource.bind(this));
		this.#signals.effect(this.#runGain.bind(this));
		this.#signals.effect(this.#runEncoder.bind(this));
		this.#signals.effect(this.#runCatalog.bind(this));
	}

	#runSource(effect: Effect): void {
		const enabled = effect.get(this.enabled);
		const media = effect.get(this.media);
		if (!enabled || !media) return;

		const settings = media.getSettings();
		if (!settings) {
			throw new Error("track has no settings");
		}

		const context = new AudioContext({
			latencyHint: "interactive",
			sampleRate: settings.sampleRate,
		});
		effect.cleanup(() => context.close());

		const root = new MediaStreamAudioSourceNode(context, {
			mediaStream: new MediaStream([media]),
		});
		effect.cleanup(() => root.disconnect());

		const gain = new GainNode(context, {
			gain: this.volume.peek(),
		});
		root.connect(gain);
		effect.cleanup(() => gain.disconnect());

		// Async because we need to wait for the worklet to be registered.
		effect.spawn(async () => {
			// Hacky workaround to support Webpack and Vite:
			// https://github.com/webpack/webpack/issues/11543#issuecomment-2045809214

			const { register } = navigator.serviceWorker;
			// @ts-ignore hack to make webpack believe that it is registering a worker
			navigator.serviceWorker.register = (url: URL) => context.audioWorklet.addModule(url);
			await navigator.serviceWorker.register(new URL("./capture-worklet", import.meta.url));
			navigator.serviceWorker.register = register;

			const worklet = new AudioWorkletNode(context, "capture", {
				numberOfInputs: 1,
				numberOfOutputs: 0,
				channelCount: settings.channelCount,
			});

			effect.set(this.#worklet, worklet);

			gain.connect(worklet);
			effect.cleanup(() => worklet.disconnect());

			// Only set the gain after the worklet is registered.
			effect.set(this.#gain, gain);
		});
	}

	#runGain(effect: Effect): void {
		const gain = effect.get(this.#gain);
		if (!gain) return;

		effect.cleanup(() => gain.gain.cancelScheduledValues(gain.context.currentTime));

		const volume = effect.get(this.muted) ? 0 : effect.get(this.volume);
		if (volume < GAIN_MIN) {
			gain.gain.exponentialRampToValueAtTime(GAIN_MIN, gain.context.currentTime + FADE_TIME);
			gain.gain.setValueAtTime(0, gain.context.currentTime + FADE_TIME + 0.01);
		} else {
			gain.gain.exponentialRampToValueAtTime(volume, gain.context.currentTime + FADE_TIME);
		}
	}

	#runEncoder(effect: Effect): void {
		if (!effect.get(this.enabled)) return;

		const media = effect.get(this.media);
		if (!media) return;

		const worklet = effect.get(this.#worklet);
		if (!worklet) return;

		this.broadcast.insertTrack(this.#track.consume());
		effect.cleanup(() => {
			this.broadcast.removeTrack(this.#track.name);
		});

		const settings = media.getSettings() as AudioTrackSettings;

		const config = {
			// TODO get codec and description from decoderConfig
			codec: "opus",
			// Firefox doesn't provide the sampleRate in the settings.
			sampleRate: u53(settings.sampleRate ?? worklet?.context.sampleRate),
			numberOfChannels: u53(settings.channelCount),
			// TODO configurable
			bitrate: u53(settings.channelCount * 32_000),
		};

		const encoder = new AudioEncoder({
			output: (frame) => {
				if (frame.type !== "key") {
					throw new Error("only key frames are supported");
				}

				if (!this.#group || frame.timestamp - this.#groupTimestamp >= 1000 * this.maxLatency) {
					this.#group?.close();
					this.#group = this.#track.appendGroup();
					this.#groupTimestamp = frame.timestamp;
				}

				const buffer = Container.encodeFrame(frame, frame.timestamp);
				this.#group.writeFrame(buffer);
			},
			error: (err) => {
				this.#group?.abort(err);
				this.#group = undefined;

				this.#track.abort(err);
			},
		});
		effect.cleanup(() => encoder.close());

		effect.cleanup(() => {
			this.#group?.close();
			this.#group = undefined;
			this.#groupTimestamp = 0;
		});

		encoder.configure({
			codec: config.codec,
			numberOfChannels: config.numberOfChannels,
			sampleRate: config.sampleRate,
			bitrate: config.bitrate,
		});

		effect.set(this.#config, config);

		worklet.port.onmessage = ({ data }: { data: Capture.AudioFrame }) => {
			const channels = data.channels.slice(0, settings.channelCount);
			const joinedLength = channels.reduce((a, b) => a + b.length, 0);
			const joined = new Float32Array(joinedLength);

			channels.reduce((offset: number, channel: Float32Array): number => {
				joined.set(channel, offset);
				return offset + channel.length;
			}, 0);

			const frame = new AudioData({
				format: "f32-planar",
				sampleRate: worklet.context.sampleRate,
				numberOfFrames: channels[0].length,
				numberOfChannels: channels.length,
				timestamp: (1_000_000 * data.timestamp) / worklet.context.sampleRate,
				data: joined,
				transfer: [joined.buffer],
			});

			encoder.encode(frame);
			frame.close();
		};
	}

	#runCatalog(effect: Effect): void {
		const config = effect.get(this.#config);
		if (!config) return;

		const captions = effect.get(this.captions.catalog);
		const speaking = effect.get(this.speaking.catalog);

		const catalog: Catalog.Audio = {
			track: {
				name: this.#track.name,
				priority: u8(this.#track.priority),
			},
			config,
			captions,
			speaking,
		};

		effect.set(this.#catalog, catalog);
	}

	close() {
		this.#signals.close();
		this.captions.close();
		this.#track.close();
	}
}