Move audio decode logic to ClientStream

This way is possible to deduplicate the opus::Decoder used by audio output.

The audio effects and client network streams are unified in only one place,
allowing the Audio Output device to consume the Samples with only one
call.

1 files changed, 82 insertions, 26 deletions

diff --git a/mumd/src/audio/output.rs b/mumd/src/audio/output.rs
index 421d395..797cf84 100644
--- a/mumd/src/audio/output.rs
+++ b/mumd/src/audio/output.rs
@@ -1,22 +1,29 @@
 use crate::network::VoiceStreamType;
+use log::*;
 
 use cpal::{OutputCallbackInfo, Sample};
 use mumble_protocol::voice::VoicePacketPayload;
 use opus::Channels;
-use std::collections::{HashMap, VecDeque};
+use std::collections::{HashMap, VecDeque, hash_map::Entry};
 use std::ops::AddAssign;
 use std::sync::{Arc, Mutex};
 use tokio::sync::watch;
 
+type ClientStreamKey = Option<(VoiceStreamType, u32)>;
+type ConsumerInput<'a> = std::collections::hash_map::IterMut<'a, ClientStreamKey, VecDeque<f32>>;
+
 pub struct ClientStream {
-    buffer: VecDeque<f32>, //TODO ring buffer?
+    buffer: HashMap<ClientStreamKey, VecDeque<f32>>, //TODO ring buffer?
     opus_decoder: opus::Decoder,
 }
 
 impl ClientStream {
     pub fn new(sample_rate: u32, channels: u16) -> Self {
+        let mut buffer = HashMap::new();
+        //None is the system audio effects
+        buffer.insert(None, VecDeque::new());
         Self {
-            buffer: VecDeque::new(),
+            buffer,
             opus_decoder: opus::Decoder::new(
                 sample_rate,
                 match channels {
@@ -29,7 +36,7 @@ impl ClientStream {
         }
     }
 
-    pub fn decode_packet(&mut self, payload: VoicePacketPayload, channels: usize) {
+    pub fn decode_packet(&mut self, client: ClientStreamKey, payload: VoicePacketPayload, channels: usize) {
         match payload {
             VoicePacketPayload::Opus(bytes, _eot) => {
                 let mut out: Vec<f32> = vec![0.0; 720 * channels * 4]; //720 is because that is the max size of packet we can get that we want to decode
@@ -38,13 +45,67 @@ impl ClientStream {
                     .decode_float(&bytes, &mut out, false)
                     .expect("Error decoding");
                 out.truncate(parsed);
-                self.buffer.extend(out);
+                self.extend(client, &out);
             }
             _ => {
                 unimplemented!("Payload type not supported");
             }
         }
     }
+
+    pub fn consume<'a, F>(&mut self, consumer: F)
+        where
+            F: FnOnce(ConsumerInput) {
+        let iter = self.buffer.iter_mut();
+        consumer(iter);
+    }
+
+    pub fn extend(&mut self, key: ClientStreamKey, values: &[f32]) {
+        match self.buffer.entry(key) {
+            Entry::Occupied(mut entry) => {
+                entry.get_mut().extend(values.iter().copied());
+            }
+            Entry::Vacant(_) => {
+                match key {
+                    None => warn!("Can't find session None"),
+                    Some(key) => warn!("Can't find session id {}", key.1),
+                }
+            }
+        }
+    }
+
+    pub fn add_client(&mut self, session_id: u32) {
+        for stream_type in [VoiceStreamType::TCP, VoiceStreamType::UDP].iter() {
+            match self.buffer.entry(Some((*stream_type, session_id))) {
+                Entry::Occupied(_) => {
+                    warn!("Session id {} already exists", session_id);
+                }
+                Entry::Vacant(entry) => {
+                    entry.insert(VecDeque::new());
+                }
+            }
+        }
+    }
+
+    pub fn remove_client(&mut self, session_id: u32) {
+        for stream_type in [VoiceStreamType::TCP, VoiceStreamType::UDP].iter() {
+            match self.buffer.entry(Some((*stream_type, session_id))) {
+                Entry::Occupied(entry) => {
+                    entry.remove();
+                }
+                Entry::Vacant(_) => {
+                    warn!(
+                        "Tried to remove session id {} that doesn't exist",
+                        session_id
+                    );
+                }
+            }
+        }
+    }
+
+    pub fn clear_clients(&mut self) {
+        self.buffer.retain(|k , _| k.is_none());
+    }
 }
 
 pub trait SaturatingAdd {
@@ -74,8 +135,7 @@ impl SaturatingAdd for u16 {
 }
 
 pub fn curry_callback<T: Sample + AddAssign + SaturatingAdd + std::fmt::Display>(
-    effect_sound: Arc<Mutex<VecDeque<f32>>>,
-    user_bufs: Arc<Mutex<HashMap<(VoiceStreamType, u32), ClientStream>>>,
+    user_bufs: Arc<Mutex<ClientStream>>,
     output_volume_receiver: watch::Receiver<f32>,
     user_volumes: Arc<Mutex<HashMap<u32, (f32, bool)>>>,
 ) -> impl FnMut(&mut [T], &OutputCallbackInfo) + Send + 'static {
@@ -86,27 +146,23 @@ pub fn curry_callback<T: Sample + AddAssign + SaturatingAdd + std::fmt::Display>
 
         let volume = *output_volume_receiver.borrow();
 
-        let mut effects_sound = effect_sound.lock().unwrap();
         let mut user_bufs = user_bufs.lock().unwrap();
-        for ((_, id), client_stream) in &mut *user_bufs {
-            let (user_volume, muted) = user_volumes
-                .lock()
-                .unwrap()
-                .get(id)
-                .cloned()
-                .unwrap_or((1.0, false));
-            for sample in data.iter_mut() {
-                let s = client_stream.buffer.pop_front().unwrap_or(0.0) * volume * user_volume;
-                if !muted {
-                    *sample = sample.saturating_add(Sample::from(&s));
+        let user_volumes = user_volumes.lock().unwrap();
+        let mut saturating_add = |input: ConsumerInput| {
+            for (k, v) in input {
+                let (user_volume, muted) = match k {
+                    Some((_, id)) => user_volumes.get(id).cloned().unwrap_or((1.0, false)),
+                    None => (1.0, false),
+                };
+                for sample in data.iter_mut() {
+                    if !muted {
+                        *sample = sample.saturating_add(Sample::from(
+                            &(v.pop_front().unwrap_or(0.0) * volume * user_volume),
+                        ));
+                    }
                 }
             }
-        }
-
-        for sample in data.iter_mut() {
-            *sample = sample.saturating_add(Sample::from(
-                &(effects_sound.pop_front().unwrap_or(0.0) * volume),
-            ));
-        }
+        };
+        user_bufs.consume(&mut saturating_add);
     }
 }