feat: load audio in chunks

Author: fractalbeauty

player/src/audio.rs

@@ -11,34 +11,38 @@ use symphonia::core::errors::Error as SymphoniaError;
 use symphonia::core::formats::FormatOptions;
 use symphonia::core::io::MediaSourceStream;
 use symphonia::core::meta::MetadataOptions;
-use symphonia::core::probe::Hint;
+use symphonia::core::probe::{Hint, ProbeResult};
 
 use cpal::traits::{DeviceTrait, HostTrait, StreamTrait};
-use cpal::{Data, PlayStreamError, Sample, SampleFormat, Stream};
+use cpal::{Sample, Stream};
 use rubato::{FftFixedIn, Resampler};
 
-pub struct Track {
-    pub samples: Arc<Vec<f32>>,
-    pub sample_rate: u32,
-    pub channel_count: usize,
-    pub position: usize,
+type Ringbuf<T> = LocalRb<T, Vec<MaybeUninit<T>>>;
+
+pub struct AudioReader {
+    probe_result: ProbeResult,
+    decoder: Box<dyn symphonia::core::codecs::Decoder>,
     encoder: Encoder,
+
+    buffer: Ringbuf<f32>,
+
+    position: usize, // we might not actually need this for anything now?
+    finished: bool,
+
+    // not read until the first packet
+    sample_rate: Option<u32>,
+    channel_count: Option<usize>,
 }
 
-impl Track {
+impl AudioReader {
     pub fn load(path: &str) -> Result<Self, Box<dyn Error>> {
         let src = std::fs::File::open(path)?;
 
         let mss = MediaSourceStream::new(Box::new(src), Default::default());
 
-        let mut hint = Hint::new();
-        //hint.with_extension("mp3");
-
-        let mut samples_not_interleaved: Option<Vec<Vec<f32>>> = None;
-        let mut sample_rate: u32 = 0;
-        let mut channel_count: usize = 0;
+        let hint = Hint::new();
 
-        let mut probe_result = symphonia::default::get_probe()
+        let probe_result = symphonia::default::get_probe()
             .format(
                 &hint,
                 mss,
@@ -50,7 +54,7 @@ impl Track {
             )
             .unwrap();
 
-        let mut decoder = symphonia::default::get_codecs()
+        let decoder = symphonia::default::get_codecs()
             .make(
                 &probe_result
                     .format
@@ -61,59 +65,100 @@ impl Track {
             )
             .unwrap();
 
-        println!("loading into memory...");
+        let encoder =
+            opus::Encoder::new(48000, opus::Channels::Stereo, opus::Application::Audio).unwrap();
+        //encoder.set_bitrate(opus::Bitrate::Bits(256)).unwrap();
+
+        Ok(Self {
+            probe_result,
+            decoder,
+            encoder,
+
+            buffer: Ringbuf::new(16384),
+
+            position: 0,
+            finished: false,
+
+            sample_rate: None,
+            channel_count: None,
+        })
+    }
+
+    // Ok(true) if there's more to read
+    fn read_packets(&mut self) -> Result<bool, Box<dyn Error>> {
+        println!("read_packet() called");
+
+        let mut has_more = true;
+        let mut channel_samples: Vec<Vec<f32>> = vec![Vec::new(); self.channel_count.unwrap_or(0)];
+        let mut sample_count = 0;
 
         loop {
-            match probe_result.format.next_packet() {
+            match self.probe_result.format.next_packet() {
                 Ok(packet) => {
-                    let decoded = decoder.decode(&packet).unwrap();
+                    let decoded = self.decoder.decode(&packet).unwrap();
                     let spec = *decoded.spec();
-                    let song_samples = match &mut samples_not_interleaved {
-                        Some(s) => Some(s),
-                        None => {
-                            samples_not_interleaved = Some(vec![Vec::new(); spec.channels.count()]);
-                            sample_rate = spec.rate;
-                            channel_count = spec.channels.count();
-                            samples_not_interleaved.as_mut()
-                        }
+
+                    if self.sample_rate == None {
+                        self.sample_rate = Some(spec.rate);
+                    }
+                    if self.channel_count == None {
+                        self.channel_count = Some(spec.channels.count());
+
+                        // we probably also need to allocate channel_samples now since this is the 1st packet we're processing
+                        channel_samples = vec![Vec::new(); spec.channels.count()];
                     }
-                    .unwrap();
 
                     if decoded.frames() > 0 {
-                        let mut samples: SampleBuffer<f32> =
+                        let mut sample_buffer: SampleBuffer<f32> =
                             SampleBuffer::new(decoded.frames() as u64, spec);
 
-                        samples.copy_interleaved_ref(decoded);
-                        for frame in samples.samples().chunks(spec.channels.count()) {
+                        sample_buffer.copy_interleaved_ref(decoded);
+
+                        let samples = sample_buffer.samples();
+
+                        println!("packet had {:?}", samples.len());
+                        sample_count += samples.len();
+
+                        for frame in samples.chunks(spec.channels.count()) {
                             for (chan, sample) in frame.iter().enumerate() {
-                                song_samples[chan].push(*sample)
+                                channel_samples[chan].push(*sample)
                             }
                         }
                     } else {
                         eprintln!("Empty packet encountered while loading song!");
                     }
                 }
-                Err(SymphoniaError::IoError(_)) => break,
+                Err(SymphoniaError::IoError(_)) => {
+                    // no more packets
+                    has_more = false;
+                    break;
+                }
                 Err(e) => return Err("failed to parse track".into()),
             }
+
+            // we should stop after a reasonable number of samples
+            // (a reasonable number --> nobody knows)
+            // we probably want more than 960 so we can resample enough at once (does that matter?)
+            if sample_count >= 2048 {
+                break;
+            }
         }
 
         println!("resampling....");
 
         // resample to standard 48000 if needed
-        let samples_correct_rate = if sample_rate != 48000 {
-            let l = samples_not_interleaved.as_ref().unwrap()[0].len().clone();
+        let samples_correct_rate = if self.sample_rate.unwrap() != 48000 {
+            let l = channel_samples[0].len();
             let mut resampler =
-                FftFixedIn::<f32>::new(sample_rate as usize, 48000, l, 100, 2).unwrap();
+                FftFixedIn::<f32>::new(self.sample_rate.unwrap() as usize, 48000, l, 100, 2)
+                    .unwrap();
 
-            resampler
-                .process(&samples_not_interleaved.unwrap(), None)
-                .unwrap()
+            resampler.process(&channel_samples, None).unwrap()
         } else {
-            samples_not_interleaved.unwrap()
+            channel_samples
         };
 
-        println!("interlacing....");
+        println!("interleaving....");
 
         // now we have to interleave it since we had to use the resampling thing
         let samples: Vec<f32> = samples_correct_rate[0]
@@ -123,38 +168,42 @@ impl Track {
             .copied()
             .collect();
 
-        // it would be really sick to do all that at once (while reading it)
-        // but since we read like.. 2 samples at a time idk how
-        // the way it is it's just too slow for anything realistically-sized
+        println!(
+            "copying {:?} samples to the buffer which has {:?} free",
+            samples.len(),
+            self.buffer.free_len()
+        );
 
-        let mut encoder =
-            opus::Encoder::new(48000, opus::Channels::Stereo, opus::Application::Audio).unwrap();
-        //encoder.set_bitrate(opus::Bitrate::Bits(256)).unwrap();
-        Ok(Self {
-            samples: Arc::new(samples),
-            sample_rate,
-            channel_count,
-            position: 0,
-            encoder,
-        })
+        for sample in samples.iter() {
+            self.buffer.push(*sample).unwrap();
+        }
+
+        Ok(has_more)
     }
 
     pub fn encode_frame(&mut self) -> Result<AudioFrame, ()> {
-        const pcm_length: usize = 960;
-
-        if pcm_length * (self.position + 1) >= self.samples.len() {
-            println!(
-                "encode_frame was about to run out of samples at pos {}",
-                self.position
-            );
-            return Err(());
+        const PCM_LENGTH: usize = 960;
+
+        // we don't have enough samples buffered to encode an opus frame, read more packets
+        if self.buffer.len() < PCM_LENGTH {
+            if let Ok(has_more) = self.read_packets() {
+                if !has_more {
+                    // no more packets, fill with 0s until we have a multiple of 960 frames
+                    while self.buffer.len() % 960 != 0 {
+                        self.buffer.push(0.0).unwrap();
+                    }
+
+                    self.finished = true;
+                }
+            } else {
+                // unrecoverable error?
+                return Err(());
+            }
         }
 
-        let mut pcm = [0.0; pcm_length];
+        let mut pcm = [0.0; PCM_LENGTH];
 
-        for i in 0..pcm_length {
-            pcm[i] = self.samples[i + pcm_length * self.position];
-        }
+        self.buffer.pop_slice(&mut pcm);
 
         let x = self.encoder.encode_vec_float(&pcm, 256).unwrap();
         self.position += 1;
@@ -164,20 +213,22 @@ impl Track {
             data: x,
         })
     }
-}
 
-type Ringbuf = LocalRb<f32, Vec<MaybeUninit<f32>>>;
+    pub fn finished(&self) -> bool {
+        self.finished
+    }
+}
 
 pub struct Player {
     decoder: Decoder,
-    buffer: Arc<Mutex<Ringbuf>>,
+    buffer: Arc<Mutex<Ringbuf<f32>>>,
     stream: Option<Stream>,
 }
 impl Player {
     pub fn new() -> Self {
         Self {
             decoder: Decoder::new(48000, opus::Channels::Stereo).unwrap(),
-            buffer: Arc::new(Mutex::new(Ringbuf::new(48000 * 10))), // 10s??
+            buffer: Arc::new(Mutex::new(Ringbuf::<f32>::new(48000 * 10))), // 10s??
             stream: None,
         }
     }
@@ -261,11 +312,15 @@ impl Player {
 fn write_audio<T: Sample>(
     out_data: &mut [f32],
     _: &cpal::OutputCallbackInfo,
-    buffer_mutex: &Arc<Mutex<Ringbuf>>,
+    buffer_mutex: &Arc<Mutex<Ringbuf<f32>>>,
 ) {
     let mut buffer = buffer_mutex.lock().unwrap();
     if out_data.len() < buffer.len() {
         buffer.pop_slice(out_data);
+
+        for sample in out_data.iter_mut() {
+            *sample = *sample * 0.1;
+        }
     } else {
         // uhhhh
         println!("write_audio: buffer underrun!! (but no panic)");

player/src/main.rs

@@ -7,10 +7,7 @@ use tokio::time::timeout;
 use tokio::{net::TcpStream, sync::mpsc};
 
 mod audio;
-use audio::Track;
-
-// use cpal::traits::{DeviceTrait, HostTrait};
-
+use audio::AudioReader;
 struct Connection {
     stream: FrameStream,
     my_id: String,
@@ -97,8 +94,11 @@ impl Connection {
                     None
                 }
             } else if line.starts_with("c ") {
-                let mut c = line.split_ascii_whitespace().into_iter().collect::<Vec<&str>>();
-                if c.len() < 2{
+                let mut c = line
+                    .split_ascii_whitespace()
+                    .into_iter()
+                    .collect::<Vec<&str>>();
+                if c.len() < 2 {
                     None
                 } else {
                     c.remove(0);
@@ -182,13 +182,11 @@ impl Connection {
 
         self.play_state = PlayState::Transmitting;
         tokio::spawn(async move {
-            let mut t: Track = match Track::load(&path) {
+            let mut t: AudioReader = match AudioReader::load(&path) {
                 Ok(t) => t,
-                Err(_) => return
+                Err(_) => return,
             };
 
-
-            dbg!(&t.samples.len());
             audio_tx_2.send(AudioData::Clear).unwrap();
             audio_tx_2.send(AudioData::Start).unwrap();
 
@@ -199,16 +197,21 @@ impl Connection {
                 .unwrap();
 
             let mut interval = tokio::time::interval(tokio::time::Duration::from_millis(200));
-            while (t.position + 1) * 960 < t.samples.len() {
+            loop {
                 // ok let's actually do the math for this
                 // each frame is 960 samples
                 // at 48k that means it's 20ms per frame
                 // SO we need to send a frame at least every 20ms.
                 // i think.................
                 // LOL OK it's two frames idk why maybe because it's stereo interleaved??????
                 interval.tick().await;
-                
+
                 for _ in 0..20 {
+                    // stop in the middle of this "tick" of 20 frames
+                    if t.finished() {
+                        break;
+                    };
+
                     let f = t.encode_frame();
                     if let Ok(frame) = f {
                         audio_tx_2.send(AudioData::Frame(frame.clone())).unwrap();
@@ -217,8 +220,14 @@ impl Connection {
                         break; // we're done i guess
                     }
                 }
+
+                // we encoded the entire file
+                if t.finished() {
+                    break;
+                };
             }
 
+            // TODO: should be "End"? and should finish playing without stopping immediately
             audio_tx_2.send(AudioData::Stop).unwrap();
             message_tx_2
                 .send(Message::PlaybackState(PlaybackState {