chanora/crates/chanora_audio/src/opus_voice.rs

use audiopus::coder::Encoder as OpusEncoder;
use audiopus::{
    Application as OpusApp, Bitrate as OpusBitrate, Channels as OpusChannels,
    SampleRate as OpusSampleRate,
};
use crossbeam::queue::ArrayQueue;
use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
use std::sync::Arc;
use tokio::sync::mpsc;
use tracing::{debug, info, warn};

use chanora_protocol::{AudioData, CodecType, OutAudio, OutPacket};

use crate::AudioError;

pub(crate) const MAX_OPUS_FRAME: usize = 1275;
const VOICE_FRAME_QUEUE_CAPACITY: usize = 64;

const VOIP_BITRATE_BPS: i32 = 32_000;
const VOIP_COMPLEXITY: u8 = 10;
const VOIP_PACKET_LOSS_PERC: u8 = 5;

pub(crate) fn new_voip_encoder(context: &str) -> Result<OpusEncoder, AudioError> {
    let mut encoder = OpusEncoder::new(OpusSampleRate::Hz48000, OpusChannels::Mono, OpusApp::Voip)
        .map_err(|e| AudioError::Opus(format!("encoder new ({context}): {e}")))?;
    tune_voip_encoder(&mut encoder, context);
    Ok(encoder)
}

pub(crate) fn tune_voip_encoder(encoder: &mut OpusEncoder, context: &str) {
    if let Err(e) = encoder.set_bitrate(OpusBitrate::BitsPerSecond(VOIP_BITRATE_BPS)) {
        warn!(target: "chanora_audio", context = %context, error = %e, "opus set_bitrate failed");
    }
    if let Err(e) = encoder.set_complexity(VOIP_COMPLEXITY) {
        warn!(target: "chanora_audio", context = %context, error = %e, "opus set_complexity failed");
    }
    if let Err(e) = encoder.set_inband_fec(true) {
        warn!(target: "chanora_audio", context = %context, error = %e, "opus set_inband_fec failed");
    }
    if let Err(e) = encoder.set_packet_loss_perc(VOIP_PACKET_LOSS_PERC) {
        warn!(
            target: "chanora_audio",
            context = %context,
            error = %e,
            "opus set_packet_loss_perc failed"
        );
    }

    info!(
        target: "chanora_audio",
        context = %context,
        bitrate_bps = VOIP_BITRATE_BPS,
        complexity = VOIP_COMPLEXITY,
        inband_fec = true,
        packet_loss_perc = VOIP_PACKET_LOSS_PERC,
        "opus encoder tuned for VoIP"
    );
}

pub(crate) struct EncodedVoiceFrame {
    data: [u8; MAX_OPUS_FRAME],
    len: usize,
}

pub(crate) struct EncodedVoiceFrameSender {
    queue: Arc<ArrayQueue<EncodedVoiceFrame>>,
    open: Arc<AtomicBool>,
}

enum EncodedVoiceFrameSendError {
    Full,
    Closed,
}

impl EncodedVoiceFrameSender {
    fn new(capacity: usize) -> Self {
        Self {
            queue: Arc::new(ArrayQueue::new(capacity)),
            open: Arc::new(AtomicBool::new(true)),
        }
    }

    fn worker_queue(&self) -> Arc<ArrayQueue<EncodedVoiceFrame>> {
        Arc::clone(&self.queue)
    }

    fn worker_open_flag(&self) -> Arc<AtomicBool> {
        Arc::clone(&self.open)
    }

    fn push(&self, frame: EncodedVoiceFrame) -> Result<(), EncodedVoiceFrameSendError> {
        if !self.open.load(Ordering::Relaxed) {
            return Err(EncodedVoiceFrameSendError::Closed);
        }
        self.queue
            .push(frame)
            .map_err(|_| EncodedVoiceFrameSendError::Full)
    }
}

pub(crate) fn start_out_packet_worker(
    voice_out_tx: mpsc::Sender<OutPacket>,
    frames_sent: Arc<AtomicU32>,
    context: &'static str,
) -> Result<EncodedVoiceFrameSender, AudioError> {
    start_out_packet_worker_with_spawner(voice_out_tx, frames_sent, context, |name, worker| {
        std::thread::Builder::new()
            .name(name)
            .spawn(worker)
            .map(|_| ())
    })
}

fn start_out_packet_worker_with_spawner<S>(
    voice_out_tx: mpsc::Sender<OutPacket>,
    frames_sent: Arc<AtomicU32>,
    context: &'static str,
    spawn: S,
) -> Result<EncodedVoiceFrameSender, AudioError>
where
    S: FnOnce(String, Box<dyn FnOnce() + Send + 'static>) -> std::io::Result<()>,
{
    let tx = EncodedVoiceFrameSender::new(VOICE_FRAME_QUEUE_CAPACITY);
    let rx = tx.worker_queue();
    let worker_open = tx.worker_open_flag();
    spawn(
        format!("chanora-{context}-voice-packets"),
        Box::new(move || {
            loop {
                let Some(frame) = rx.pop() else {
                    if Arc::strong_count(&rx) == 1 {
                        break;
                    }
                    std::thread::sleep(std::time::Duration::from_millis(1));
                    continue;
                };
                let packet = OutAudio::new(&AudioData::C2S {
                    id: 0,
                    codec: CodecType::OpusVoice,
                    data: frame.as_slice(),
                });
                match voice_out_tx.try_send(packet) {
                    Ok(()) => {
                        frames_sent.fetch_add(1, Ordering::Relaxed);
                    }
                    Err(mpsc::error::TrySendError::Full(_)) => {
                        warn!(target: "chanora_audio", context = %context, "voice_out queue full; dropping frame");
                    }
                    Err(mpsc::error::TrySendError::Closed(_)) => {
                        debug!(target: "chanora_audio", context = %context, "voice_out closed; voice packet worker stopping");
                        worker_open.store(false, Ordering::Relaxed);
                        break;
                    }
                }
            }
        }),
    )
    .map_err(|e| {
        tx.open.store(false, Ordering::Relaxed);
        AudioError::Backend(format!("voice packet worker spawn ({context}): {e}"))
    })?;
    Ok(tx)
}

impl EncodedVoiceFrame {
    fn try_from_opus(opus_out: &[u8], len: usize) -> Option<Self> {
        if len > opus_out.len() || len > MAX_OPUS_FRAME {
            return None;
        }
        let mut data = [0u8; MAX_OPUS_FRAME];
        data[..len].copy_from_slice(&opus_out[..len]);
        Some(Self { data, len })
    }

    fn as_slice(&self) -> &[u8] {
        &self.data[..self.len]
    }
}

/// Encode-scope send helper for a freshly encoded Opus voice frame.
pub(crate) fn send_voip_frame<F, G>(
    voice_out_tx: &EncodedVoiceFrameSender,
    opus_out: &[u8],
    len: usize,
    on_full: F,
    on_closed: G,
) where
    F: FnOnce(),
    G: FnOnce(),
{
    let Some(frame) = EncodedVoiceFrame::try_from_opus(opus_out, len) else {
        on_full();
        return;
    };
    match voice_out_tx.push(frame) {
        Ok(()) => {}
        Err(EncodedVoiceFrameSendError::Full) => on_full(),
        Err(EncodedVoiceFrameSendError::Closed) => on_closed(),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn encoded_voice_frame_copies_into_fixed_storage() {
        let source = [7u8; MAX_OPUS_FRAME];

        let frame = EncodedVoiceFrame::try_from_opus(&source, MAX_OPUS_FRAME).unwrap();

        assert_eq!(frame.as_slice().len(), MAX_OPUS_FRAME);
        assert!(frame.as_slice().iter().all(|byte| *byte == 7));
    }

    #[test]
    fn encoded_voice_frame_rejects_lengths_beyond_fixed_storage() {
        let source = [0u8; MAX_OPUS_FRAME];

        assert!(EncodedVoiceFrame::try_from_opus(&source, MAX_OPUS_FRAME + 1).is_none());
    }

    #[test]
    fn encoded_voice_frame_sender_reports_full_without_blocking() {
        let sender = EncodedVoiceFrameSender::new(1);
        let source = [3u8; MAX_OPUS_FRAME];
        let first = EncodedVoiceFrame::try_from_opus(&source, 4).unwrap();
        let second = EncodedVoiceFrame::try_from_opus(&source, 4).unwrap();

        assert!(sender.push(first).is_ok());
        assert!(sender.push(second).is_err());
    }

    #[test]
    fn encoded_voice_frame_sender_reports_closed_without_queueing() {
        let sender = EncodedVoiceFrameSender::new(1);
        sender.open.store(false, Ordering::Relaxed);
        let source = [3u8; MAX_OPUS_FRAME];
        let frame = EncodedVoiceFrame::try_from_opus(&source, 4).unwrap();

        assert!(matches!(
            sender.push(frame),
            Err(EncodedVoiceFrameSendError::Closed)
        ));
        assert_eq!(sender.queue.len(), 0);
    }

    #[test]
    fn encoded_voice_frame_sender_reports_spawn_failure() {
        let (voice_out_tx, _voice_out_rx) = mpsc::channel(1);
        let frames_sent = Arc::new(AtomicU32::new(0));

        let result = start_out_packet_worker_with_spawner(
            voice_out_tx,
            frames_sent,
            "test",
            |_name, _worker| Err(std::io::Error::other("spawn failed")),
        );

        assert!(
            matches!(result, Err(AudioError::Backend(message)) if message.contains("spawn failed"))
        );
    }
}