fix bad audio quality & save_to_disk example (#64)

- save_to_disk example: write raw pcm frames into a wav file

examples/Cargo.lock

@@ -2377,6 +2377,17 @@ dependencies = [
  "winapi-util",
 ]
 
+[[package]]
+name = "save_to_disk"
+version = "0.1.0"
+dependencies = [
+ "bytes",
+ "futures",
+ "livekit",
+ "tokio",
+ "tokio-util",
+]
+
 [[package]]
 name = "schannel"
 version = "0.1.21"

examples/save_to_disk/Cargo.toml

@@ -0,0 +1,13 @@
+[package]
+name = "save_to_disk"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+tokio = { version = "1", features = ["full"] }
+livekit = { path = "../../livekit", version = "0.1.1" }
+bytes = "1.4.0"
+tokio-util = "0.7.8"
+futures = "0.3.28"

examples/save_to_disk/src/main.rs

@@ -0,0 +1,143 @@
+use bytes::{BufMut, BytesMut};
+use futures::StreamExt;
+use livekit::prelude::*;
+use livekit::webrtc::audio_stream::native::NativeAudioStream;
+use std::env;
+use tokio::fs::File;
+use tokio::io::{AsyncWriteExt, BufWriter};
+
+const WAV_HEADER_SIZE: usize = 44;
+const FILE_PATH: &str = "record.wav";
+
+#[derive(Debug, Clone, Copy)]
+pub struct WavHeader {
+    pub sample_rate: u32,
+    pub bit_depth: u16,
+    pub num_channels: u16,
+}
+
+pub struct WavWriter {
+    header: WavHeader,
+    data: BytesMut,
+    writer: BufWriter<File>,
+}
+
+impl WavWriter {
+    pub async fn create<P: AsRef<std::path::Path>>(
+        path: P,
+        header: WavHeader,
+    ) -> Result<WavWriter, std::io::Error> {
+        let file = File::create(path).await?;
+        let writer = BufWriter::new(file);
+
+        let mut wav_writer = WavWriter {
+            header,
+            data: BytesMut::new(),
+            writer,
+        };
+
+        wav_writer.write_header()?;
+        Ok(wav_writer)
+    }
+
+    fn write_header(&mut self) -> Result<(), std::io::Error> {
+        let byte_rate = (self.header.sample_rate
+            * self.header.bit_depth as u32
+            * self.header.num_channels as u32);
+
+        let block_align = byte_rate / self.header.sample_rate as u32;
+
+        self.data.put_slice(b"RIFF");
+        self.data.put_u32_le(0); // Placeholder for file size
+        self.data.put_slice(b"WAVE");
+        self.data.put_slice(b"fmt ");
+        self.data.put_u32_le(16); // Subchunk1Size (16 for PCM)
+        self.data.put_u16_le(1); // AudioFormat (1 for PCM)
+        self.data.put_u16_le(self.header.num_channels);
+        self.data.put_u32_le(self.header.sample_rate);
+        self.data.put_u32_le(byte_rate);
+        self.data.put_u16_le(32);
+        self.data.put_u16_le(self.header.bit_depth);
+        self.data.put_slice(b"data");
+        self.data.put_u32_le(0); // Placeholder for data size
+
+        assert_eq!(self.data.len(), WAV_HEADER_SIZE);
+
+        Ok(())
+    }
+
+    pub async fn write_sample(&mut self, sample: i16) -> Result<(), std::io::Error> {
+        self.data.put_i16_le(sample);
+        Ok(())
+    }
+
+    pub async fn finalize(mut self) -> Result<(), std::io::Error> {
+        let data_size = self.data.len() as u32 - WAV_HEADER_SIZE as u32;
+        let file_size = data_size + WAV_HEADER_SIZE as u32 - 8;
+        self.data.as_mut()[4..8].copy_from_slice(&file_size.to_le_bytes());
+        self.data.as_mut()[40..44].copy_from_slice(&data_size.to_le_bytes());
+
+        self.writer.write_all(&self.data).await?;
+        self.writer.flush().await?;
+        Ok(())
+    }
+}
+
+#[tokio::main]
+async fn main() {
+    let url = env::var("LIVEKIT_URL").expect("LIVEKIT_URL is not set");
+    let token = env::var("LIVEKIT_TOKEN").expect("LIVEKIT_TOKEN is not set");
+
+    let (room, mut rx) = Room::connect(&url, &token).await.unwrap();
+    let session = room.session();
+    println!("Connected to room: {} - {}", session.name(), session.sid());
+
+    while let Some(msg) = rx.recv().await {
+        match msg {
+            RoomEvent::TrackSubscribed {
+                track,
+                publication: _,
+                participant: _,
+            } => {
+                if let RemoteTrack::Audio(audio_track) = track {
+                    record_track(audio_track).await.unwrap();
+                    break;
+                }
+            }
+            _ => {}
+        }
+    }
+
+    println!("Done");
+}
+
+async fn record_track(audio_track: RemoteAudioTrack) -> Result<(), std::io::Error> {
+    println!("Recording track {:?}", audio_track.sid());
+    let rtc_track = audio_track.rtc_track();
+
+    // TODO(theomonnom): Remove hardcoded values
+    let header = WavHeader {
+        sample_rate: 48000,
+        bit_depth: 16,
+        num_channels: 1,
+    };
+
+    let mut wav_writer = WavWriter::create(FILE_PATH, header).await?;
+    let mut audio_stream = NativeAudioStream::new(rtc_track);
+
+    let max_record = 5 * header.sample_rate * header.num_channels as u32;
+    let mut sample_count = 0;
+    'recv_loop: while let Some(frame) = audio_stream.next().await {
+        for sample in frame.data {
+            wav_writer.write_sample(sample).await.unwrap();
+            sample_count += 1;
+
+            if sample_count >= max_record {
+                break 'recv_loop;
+            }
+        }
+    }
+
+    wav_writer.finalize().await?;
+    Ok(())
+}