moq/

audio.rs

//! Raw-audio import/export via [`moq_audio`].
//!
//! Sibling to `moq_publish_media_*` / `moq_consume_audio`
//! (those handle already-encoded frames). These functions accept and
//! return raw PCM, with Opus encode/decode happening inside the FFI
//! boundary.
//!
//! Format / sample rate / channel count are fixed at producer or
//! consumer construction via [`moq_audio_encoder_input`] /
//! [`moq_audio_encoder_output`] / [`moq_audio_decoder_output`], so
//! each [`moq_audio_frame`] carries only payload bytes and a
//! timestamp.

use std::ffi::{c_char, c_void};
use std::sync::Arc;
use std::time::Duration;

use bytes::Bytes;
use tokio::sync::oneshot;

use crate::ffi::OnStatus;
use crate::{Error, Id, NonZeroSlab, Shared, State, ffi};

// ---- C-visible types ----

/// Raw PCM sample layout, mirroring WebCodecs `AudioData.format`.
///
/// The enum is exposed in the C header for readability, but ABI
/// fields/parameters that carry it are typed `u32`. A C caller
/// passing an unknown discriminant gets `Error::InvalidCode` instead
/// of UB.
///
/// <https://developer.mozilla.org/en-US/docs/Web/API/AudioData/format>
#[repr(C)]
#[allow(non_camel_case_types)]
#[derive(Clone, Copy, Debug)]
pub enum moq_audio_sample_format {
	MOQ_AUDIO_SAMPLE_FORMAT_U8 = 0,
	MOQ_AUDIO_SAMPLE_FORMAT_S16 = 1,
	MOQ_AUDIO_SAMPLE_FORMAT_S32 = 2,
	MOQ_AUDIO_SAMPLE_FORMAT_F32 = 3,
	MOQ_AUDIO_SAMPLE_FORMAT_U8_PLANAR = 4,
	MOQ_AUDIO_SAMPLE_FORMAT_S16_PLANAR = 5,
	MOQ_AUDIO_SAMPLE_FORMAT_S32_PLANAR = 6,
	MOQ_AUDIO_SAMPLE_FORMAT_F32_PLANAR = 7,
}

fn audio_format_from_u32(value: u32) -> Result<moq_audio::Format, Error> {
	use moq_audio::Format;
	Ok(match value {
		v if v == moq_audio_sample_format::MOQ_AUDIO_SAMPLE_FORMAT_U8 as u32 => Format::U8,
		v if v == moq_audio_sample_format::MOQ_AUDIO_SAMPLE_FORMAT_S16 as u32 => Format::S16,
		v if v == moq_audio_sample_format::MOQ_AUDIO_SAMPLE_FORMAT_S32 as u32 => Format::S32,
		v if v == moq_audio_sample_format::MOQ_AUDIO_SAMPLE_FORMAT_F32 as u32 => Format::F32,
		v if v == moq_audio_sample_format::MOQ_AUDIO_SAMPLE_FORMAT_U8_PLANAR as u32 => Format::U8Planar,
		v if v == moq_audio_sample_format::MOQ_AUDIO_SAMPLE_FORMAT_S16_PLANAR as u32 => Format::S16Planar,
		v if v == moq_audio_sample_format::MOQ_AUDIO_SAMPLE_FORMAT_S32_PLANAR as u32 => Format::S32Planar,
		v if v == moq_audio_sample_format::MOQ_AUDIO_SAMPLE_FORMAT_F32_PLANAR as u32 => Format::F32Planar,
		_ => return Err(Error::InvalidCode),
	})
}

/// PCM layout the caller hands to [`moq_encode_audio_frame`].
#[repr(C)]
#[allow(non_camel_case_types)]
pub struct moq_audio_encoder_input {
	/// `moq_audio_sample_format` discriminant.
	pub format: u32,
	pub sample_rate: u32,
	/// Interleaved channel count, which also names the speaker layout by the
	/// WAVE convention: 1 mono, 2 stereo, 3 2.1, 4 quad, 5 5.0, 6 5.1, 7 6.1,
	/// 8 7.1, in front left, front right, center, LFE, back, side order.
	pub channels: u32,
}

/// Codec-side configuration. `sample_rate` / `channels` = 0 means
/// "match the input (snapping the rate up to a libopus-supported
/// value if necessary)".
#[repr(C)]
#[allow(non_camel_case_types)]
pub struct moq_audio_encoder_output {
	/// Codec id, UTF-8: "opus", "pcm", or "aac". AAC encodes through the
	/// platform's encoder, so a host without one refuses it.
	pub codec: *const c_char,
	pub codec_len: usize,
	/// 0 = derive from input.
	pub sample_rate: u32,
	/// 0 = derive from input.
	pub channels: u32,
	/// 0 = libopus default.
	pub bitrate: u32,
	/// Encoded frame duration in microseconds. Opus accepts exactly
	/// 2500/5000/10000/20000/40000/60000 us. 0 = the codec's default: 20 ms for
	/// Opus, which matches the JS publish path, and 1024 samples for AAC.
	pub frame_duration_us: u32,
}

/// PCM layout the caller wants out of [`moq_decode_audio`].
#[repr(C)]
#[allow(non_camel_case_types)]
pub struct moq_audio_decoder_output {
	pub format: u32,
	/// 0 = deliver at the codec's native sample rate.
	pub sample_rate: u32,
	/// 0 = deliver at the codec's native channel count. A count names its
	/// layout as `moq_audio_encoder_input.channels` describes, and the decoder
	/// remixes to it.
	pub channels: u32,
	/// Upper bound on buffering before skipping a stalled group, in
	/// microseconds. Same congestion-control knob as
	/// `moq_consume_audio`'s `max_age_us`. 0 = skip
	/// aggressively (the moq-mux default); set to your playout
	/// buffer (tens to a few hundred ms) for a softer skip. Named
	/// `_max` to leave room for a future `min_buffer_us`, a
	/// jitter-buffer floor rather than a staleness bound.
	pub max_age_us: u64,
}

/// One audio frame: payload bytes plus a presentation timestamp.
///
/// `data` is owned by the consume slab (see
/// [`moq_decode_audio_frame_free`]) or borrowed by the publish call
/// (the publisher copies before returning).
#[repr(C)]
#[allow(non_camel_case_types)]
pub struct moq_audio_frame {
	pub timestamp_us: u64,
	pub data: *const u8,
	pub data_size: usize,
}

// ---- State extensions (used internally by lib.rs) ----

/// An audio producer, shared so the Opus encode in `write` runs with the global
/// lock released. See [`Shared`].
pub(crate) struct AudioEncoder {
	producer: moq_audio::encode::Producer<moq_mux::catalog::hang::Extra>,
	reservation: Option<Arc<moq_net::bandwidth::Reservation>>,
}

type AudioProducer = Shared<AudioEncoder>;

#[derive(Default)]
pub struct Audio {
	producers: NonZeroSlab<AudioProducer>,
	consumer_tasks: NonZeroSlab<Option<AudioTaskEntry>>,
	frames: NonZeroSlab<moq_audio::Frame>,
}

/// A spawned task entry: `close` signals shutdown, `callback` delivers status.
///
/// `close` is an `Option` so `consume_close` can drop just the sender without
/// removing the entry. The task delivers one final terminal callback and then
/// removes itself, so `user_data` stays valid until that callback fires.
struct AudioTaskEntry {
	close: Option<oneshot::Sender<()>>,
	callback: OnStatus,
}

impl Audio {
	pub fn publish(
		&mut self,
		broadcast: &mut moq_net::broadcast::Producer,
		catalog: moq_mux::catalog::Producer<moq_mux::catalog::hang::Extra>,
		input: moq_audio::encode::Input,
		options: moq_audio::encode::Options,
		reserve: bool,
	) -> Result<Id, Error> {
		let producer = moq_audio::encode::Producer::new(broadcast, catalog, input, &options)?;
		let reservation = reserve.then(|| Arc::new(options.bandwidth.reserve(&producer.demand(), producer.bitrate())));
		self.producers
			.insert(Shared::new(AudioEncoder { producer, reservation }))
	}

	pub(crate) fn reservation(&self, id: Id) -> Result<Option<Arc<moq_net::bandwidth::Reservation>>, Error> {
		Ok(self
			.producer(id)?
			.lock()
			.as_ref()
			.ok_or(Error::MediaNotFound)?
			.reservation
			.clone())
	}

	/// A watch-only handle to the encoded track's subscriber demand.
	pub(crate) fn demand(&self, id: Id) -> Result<moq_net::track::Demand, Error> {
		Ok(self
			.producer(id)?
			.lock()
			.as_ref()
			.ok_or(Error::MediaNotFound)?
			.producer
			.demand())
	}

	/// Resolve a producer handle, so the caller can encode with the global lock
	/// released.
	///
	/// Bind the result before locking it: a temporary [`State`] guard lives to the
	/// end of the statement that created it, so resolving and locking in one
	/// expression would put the encode back under the global lock.
	pub(crate) fn producer(&self, id: Id) -> Result<AudioProducer, Error> {
		self.producers.get(id).cloned().ok_or(Error::MediaNotFound)
	}

	/// Resolve a producer and drop its id, so nothing can be published to it after.
	pub(crate) fn remove(&mut self, id: Id) -> Result<AudioProducer, Error> {
		self.producers.remove(id).ok_or(Error::MediaNotFound)
	}

	pub fn consume(
		&mut self,
		broadcast: &moq_net::broadcast::Consumer,
		catalog: &hang::catalog::AudioConfig,
		name: &str,
		config: moq_audio::decode::Options,
		on_frame: OnStatus,
	) -> Result<Id, Error> {
		let broadcast = broadcast.clone();
		let catalog = catalog.clone();
		let name = name.to_string();

		let channel = oneshot::channel();
		let entry = AudioTaskEntry {
			close: Some(channel.0),
			callback: on_frame,
		};
		let id = self.consumer_tasks.insert(Some(entry))?;

		// `decode::Consumer::new` subscribes (blocking on SUBSCRIBE_OK), so run it
		// inside the task to keep this entrypoint non-blocking.
		tokio::spawn(async move {
			let res = async move {
				let consumer = moq_audio::decode::Consumer::new(&broadcast, &catalog, name, config).await?;
				Self::run(on_frame, consumer, channel.1).await
			}
			.await;

			// Deliver one final terminal callback (code <= 0), then drop the entry.
			// Pull it out from under the lock so the callback never runs while held.
			let entry = State::lock().audio.consumer_tasks.remove(id).flatten();
			if let Some(entry) = entry {
				entry.callback.call(res);
			}
		});

		Ok(id)
	}

	async fn run(
		callback: OnStatus,
		mut consumer: moq_audio::decode::Consumer,
		mut close: oneshot::Receiver<()>,
	) -> Result<(), Error> {
		loop {
			// `biased` so a pending close always wins over a ready frame.
			let frame = tokio::select! {
				biased;
				_ = &mut close => return Ok(()),
				frame = consumer.read() => match frame {
					Ok(Some(frame)) => frame,
					Ok(None) => return Ok(()),
					// One packet the codec rejected is that packet's problem: the
					// decoder stays usable, so drop it and keep the subscription rather
					// than ending the caller's stream over a single bad frame.
					Err(moq_audio::Error::Decode(err)) => {
						tracing::warn!(%err, "dropping an audio frame");
						continue;
					}
					Err(err) => return Err(err.into()),
				},
			};

			// Hold the lock only to buffer the frame; release it before the callback.
			let frame_id = State::lock().audio.frames.insert(frame)?;
			callback.call(Ok(frame_id));
		}
	}

	pub fn consume_close(&mut self, id: Id) -> Result<(), Error> {
		// Signal shutdown; the task delivers a final callback and removes itself.
		self.consumer_tasks
			.get_mut(id)
			.and_then(|entry| entry.as_mut())
			.ok_or(Error::TrackNotFound)?
			.close
			.take()
			.ok_or(Error::TrackNotFound)?;
		Ok(())
	}

	pub fn frame_info(&self, id: Id, dst: &mut moq_audio_frame) -> Result<(), Error> {
		let frame = self.frames.get(id).ok_or(Error::FrameNotFound)?;
		*dst = moq_audio_frame {
			// The C ABI carries plain microseconds, so flatten the scaled
			// `Timestamp` here at the boundary. Saturating rather than erroring:
			// this is a getter on a frame we already decoded, and a u64 overflow
			// needs a timestamp ~580,000 years out.
			timestamp_us: u64::try_from(frame.timestamp.as_micros()).unwrap_or(u64::MAX),
			data: frame.data.as_ptr(),
			data_size: frame.data.len(),
		};
		Ok(())
	}

	pub fn frame_free(&mut self, id: Id) -> Result<(), Error> {
		self.frames.remove(id).ok_or(Error::FrameNotFound)?;
		Ok(())
	}
}

// ---- C entry points ----

/// Open an audio track on a broadcast.
///
/// The encoder configuration is fixed at construction; subsequent
/// frame writes pass only payload + timestamp via
/// [`moq_encode_audio_frame`].
///
/// Returns a non-zero handle on success or a negative error code.
///
/// # Safety
/// - `name` must point to `name_len` bytes of UTF-8.
/// - `input` / `output` must point to fully populated structs.
/// - `output->codec` must point to `output->codec_len` bytes of UTF-8.
/// - `bandwidth` is a handle from [`crate::moq_session_bandwidth`], or 0 to leave the
///   configured bitrate unclaimed.
#[unsafe(no_mangle)]
pub unsafe extern "C" fn moq_encode_audio(
	broadcast: u32,
	name: *const c_char,
	name_len: usize,
	input: *const moq_audio_encoder_input,
	output: *const moq_audio_encoder_output,
	bandwidth: u32,
) -> i32 {
	ffi::enter(move || {
		let broadcast = ffi::parse_id(broadcast)?;
		let name = unsafe { ffi::parse_str(name, name_len)? }.to_string();
		let raw_input = unsafe { input.as_ref() }.ok_or(Error::InvalidPointer)?;
		let raw_output = unsafe { output.as_ref() }.ok_or(Error::InvalidPointer)?;
		let codec_str = unsafe { ffi::parse_str(raw_output.codec, raw_output.codec_len)? };

		let layout = moq_audio::Layout::from_channels(raw_input.channels)?;
		let mut encoder_input = moq_audio::encode::Input::new(raw_input.sample_rate, layout);
		encoder_input.format = audio_format_from_u32(raw_input.format)?;

		// The C ABI takes an explicit track name and spells "unset" as 0, so map
		// both onto the Rust options here rather than leaking either convention.
		let mut options = moq_audio::encode::Options::default();
		options.track = Some(name);
		let codec = codec_str
			.parse()
			.map_err(|_| Error::UnknownFormat(codec_str.to_string()))?;
		options.settings = moq_audio::encode::Settings::from_input(codec, &encoder_input);
		if let Some(sample_rate) = zeroable(raw_output.sample_rate) {
			options.settings.sample_rate = sample_rate;
		}
		if let Some(channels) = zeroable(raw_output.channels) {
			options.settings.layout = moq_audio::Layout::from_channels(channels)?;
		}
		options.settings.bitrate =
			zeroable(raw_output.bitrate).map(|bps| moq_net::bandwidth::Rate::from_bps(bps.into()));
		if let Some(micros) = zeroable(raw_output.frame_duration_us) {
			options.settings.frame_duration = Duration::from_micros(micros.into());
		}

		let bandwidth = ffi::parse_id_optional(bandwidth)?;
		let mut state = State::lock();
		if let Some(id) = bandwidth {
			options.bandwidth = state.bandwidth.allocator(id)?;
		}
		let State { publish, audio, .. } = &mut *state;
		let (broadcast_producer, catalog) = publish.pair_mut(broadcast)?;

		audio.publish(
			broadcast_producer,
			catalog.clone(),
			encoder_input,
			options,
			bandwidth.is_some(),
		)
	})
}

/// This encoder's bandwidth reservation, or 0 if it was published without one.
///
/// Closing the returned handle does not release the encoder's claim; that lasts
/// until [moq_encode_audio_finish].
#[unsafe(no_mangle)]
pub extern "C" fn moq_encode_audio_reservation(producer: u32) -> i32 {
	ffi::enter(move || {
		let producer = ffi::parse_id(producer)?;
		let mut state = State::lock();
		match state.audio.reservation(producer)? {
			Some(reservation) => Ok(i32::from(state.bandwidth.hold(reservation)?)),
			None => Ok(0),
		}
	})
}

/// Watch whether the encoded audio track has subscribers, so the microphone and
/// encoder run only while someone listens. See [`crate::moq_publish_media_demand`]
/// for the callback contract.
///
/// Returns a non-zero watcher handle on success, or a negative code on failure.
///
/// # Safety
/// - `on_demand` must be non-NULL.
/// - The caller must keep `user_data` valid until the terminal (`<= 0`) `on_demand` callback.
#[unsafe(no_mangle)]
pub unsafe extern "C" fn moq_encode_audio_demand(
	producer: u32,
	on_demand: crate::moq_status_callback,
	user_data: *mut c_void,
) -> i32 {
	ffi::enter(move || {
		let producer = ffi::parse_id(producer)?;
		let on_demand = unsafe { OnStatus::new(user_data, on_demand)? };
		let mut state = State::lock();
		let demand = state.audio.demand(producer)?;
		state.publish.demand(demand, on_demand)
	})
}

/// The C ABI spells an unset `u32` knob as 0, which no field here accepts as a
/// real value.
fn zeroable(value: u32) -> Option<u32> {
	(value != 0).then_some(value)
}

/// Push one audio frame.
///
/// `frame->data` is borrowed for the duration of the call; the
/// producer copies before returning.
///
/// # Safety
/// - `frame` must point to a valid [`moq_audio_frame`].
/// - `frame->data` must point to `frame->data_size` bytes.
#[unsafe(no_mangle)]
pub unsafe extern "C" fn moq_encode_audio_frame(producer: u32, frame: *const moq_audio_frame) -> i32 {
	ffi::enter(move || {
		let producer = ffi::parse_id(producer)?;
		let frame = unsafe { frame.as_ref() }.ok_or(Error::InvalidPointer)?;
		let data = unsafe { ffi::parse_slice(frame.data, frame.data_size)? };

		// The C ABI carries plain microseconds; scale them at the boundary.
		let timestamp = moq_net::Timestamp::from_micros(frame.timestamp_us).map_err(moq_audio::Error::from)?;
		let owned = moq_audio::Frame::new(Bytes::copy_from_slice(data), timestamp);

		let producer = State::lock().audio.producer(producer)?;
		producer
			.lock()
			.as_mut()
			.ok_or(Error::MediaNotFound)?
			.producer
			.write(&owned)?;
		Ok(())
	})
}

/// Flush any pending samples and finalize an audio producer.
#[unsafe(no_mangle)]
pub extern "C" fn moq_encode_audio_finish(producer: u32) -> i32 {
	ffi::enter(move || {
		let producer = ffi::parse_id(producer)?;
		// The id is dropped first, so nothing new queues behind the flush; whatever
		// is mid-encode still finishes before this takes the producer.
		let producer = State::lock().audio.remove(producer)?;
		let mut producer = producer.take().ok_or(Error::MediaNotFound)?.producer;
		producer.finish()?;
		Ok(())
	})
}

/// Subscribe to an audio track and decode it into PCM.
///
/// The catalog `index` identifies which audio rendition to subscribe
/// to, matching the existing `moq_consume_audio` selection
/// model. TODO: a future API will pick the right rendition
/// automatically (ABR).
///
/// Returns a non-zero handle on success or a negative error code.
///
/// `on_frame` is called with a positive frame ID per frame, then exactly once
/// more with a terminal code: `0` (closed cleanly) or a negative error. After
/// the terminal (`<= 0`) callback, `on_frame` is never called again and
/// `user_data` is never touched again, so release `user_data` there. The
/// terminal callback fires even after [`moq_decode_audio_cancel`].
///
/// Starts at the newest cached group so reopening live playback skips the backlog.
///
/// A packet the codec cannot decode is logged and skipped rather than ending
/// the subscription, so a single bad frame costs that frame and not the stream.
///
/// # Safety
/// - `output` must point to a valid [`moq_audio_decoder_output`].
/// - `user_data` must stay valid until the terminal (`<= 0`) `on_frame` callback.
#[unsafe(no_mangle)]
pub unsafe extern "C" fn moq_decode_audio(
	catalog: u32,
	index: u32,
	output: *const moq_audio_decoder_output,
	on_frame: crate::moq_status_callback,
	user_data: *mut c_void,
) -> i32 {
	ffi::enter(move || {
		let catalog = ffi::parse_id(catalog)?;
		let raw = unsafe { output.as_ref() }.ok_or(Error::InvalidPointer)?;

		let mut config = moq_audio::decode::Options::default();
		config.start = moq_audio::decode::Start::Latest;
		config.output.format = audio_format_from_u32(raw.format)?;
		config.output.sample_rate = zeroable(raw.sample_rate);
		config.output.layout = zeroable(raw.channels)
			.map(moq_audio::Layout::from_channels)
			.transpose()?;
		config.max_age = Duration::from_micros(raw.max_age_us);

		let on_frame = unsafe { OnStatus::new(user_data, on_frame)? };

		let mut state = State::lock();
		let (broadcast, audio_cfg, name) = state.consume.audio_rendition(catalog, index as usize)?;

		let State { audio, .. } = &mut *state;
		audio.consume(&broadcast, &audio_cfg, &name, config, on_frame)
	})
}

/// Stop an audio (raw PCM) consumer's background task.
///
/// Returns immediately: zero on success, or a negative code if already closed.
/// Does NOT free `user_data`; the on-frame callback still fires once more with a
/// terminal `0` (or a negative error), which is where `user_data` should be
/// released. Frame IDs already delivered to the callback are likewise not freed;
/// release each with [`moq_decode_audio_frame_free`].
#[unsafe(no_mangle)]
pub extern "C" fn moq_decode_audio_cancel(consumer: u32) -> i32 {
	ffi::enter(move || {
		let consumer = ffi::parse_id(consumer)?;
		State::lock().audio.consume_close(consumer)
	})
}

/// Copy a delivered frame's metadata into `dst`.
///
/// The written `dst->data` pointer remains valid until the same `id`
/// is released with [`moq_decode_audio_frame_free`].
///
/// # Safety
/// - `dst` must point to a writable [`moq_audio_frame`].
#[unsafe(no_mangle)]
pub unsafe extern "C" fn moq_decode_audio_frame(id: u32, dst: *mut moq_audio_frame) -> i32 {
	ffi::enter(move || {
		let id = ffi::parse_id(id)?;
		let dst = unsafe { dst.as_mut() }.ok_or(Error::InvalidPointer)?;
		State::lock().audio.frame_info(id, dst)
	})
}

/// Free a frame previously delivered through the consume callback.
/// Required for every delivered frame ID; closing the parent consumer
/// is not enough.
#[unsafe(no_mangle)]
pub extern "C" fn moq_decode_audio_frame_free(id: u32) -> i32 {
	ffi::enter(move || {
		let id = ffi::parse_id(id)?;
		State::lock().audio.frame_free(id)
	})
}