fix: address thread shutdown

- replace unwrap() calls with checks
- add stop flag to allow thread shutdown
- implement drop trait to automatically set the stop flag

src/main.rs

@@ -1,6 +1,10 @@
 use std::{
   io::{self, stdout},
-  sync::{Arc, Mutex, mpsc},
+  sync::{
+    Arc, Mutex,
+    atomic::{AtomicBool, Ordering},
+    mpsc,
+  },
   thread,
   time::Duration,
 };
@@ -143,8 +147,9 @@ impl Drop for TerminalGuard {
   }
 }
 
-pub struct AudioVisualizer {
+struct AudioVisualizer {
-  pub bar_data: Arc<Mutex<Vec<u64>>>,
+  bar_data: Arc<Mutex<Vec<u64>>>,
+  stop_flag: Arc<AtomicBool>,
 }
 
 // This is a implementation largely lifted from these
@@ -152,18 +157,22 @@ pub struct AudioVisualizer {
 // 1. CAVA (C): https://github.com/karlstav/cava
 // 2. cli-visualizer (C++): https://github.com/dpayne/cli-visualizer
 impl AudioVisualizer {
-  pub fn new(num_bars: usize) -> Self {
+  fn new(num_bars: usize) -> Self {
-    // 
     let bar_data = Arc::new(Mutex::new(vec![0; num_bars]));
     let bar_data_clone = Arc::clone(&bar_data);
+    let stop_flag = Arc::new(AtomicBool::new(false));
+    let stop_clone = Arc::new(AtomicBool::new(false));
 
     thread::spawn(move || {
-      if let Err(e) = Self::run_audio_loop(bar_data_clone, num_bars) {
+      if let Err(e) = Self::run_audio_loop(bar_data_clone, num_bars, stop_clone) {
         eprintln!("Audio capture error: {:?}", e);
       }
     });
 
-    Self { bar_data }
+    Self {
+      bar_data,
+      stop_flag,
+    }
   }
 
   // Doing best effort to support cross-platform functionality
@@ -304,7 +313,7 @@ impl AudioVisualizer {
     smoothed
   }
 
-  // songs can be quiet and load so we try to adjust
+  // songs can be quiet and loud so we try to adjust
   // sensitivity to avoid bars becoming flattened or
   // clipped
   fn adjust_autosens(autosens: &mut f32, bars: &[f32]) {
@@ -321,6 +330,7 @@ impl AudioVisualizer {
   fn run_audio_loop(
     bar_data: Arc<Mutex<Vec<u64>>>,
     num_bars: usize,
+    stop: Arc<AtomicBool>,
   ) -> Result<(), Box<dyn std::error::Error>> {
     let host = cpal::default_host();
     let (device, config) = Self::get_audio_device(&host)?;
@@ -353,8 +363,7 @@ impl AudioVisualizer {
     let stream = device.build_input_stream(
       config,
       move |data: &[f32], _| {
-        let mut buf = buffer_clone.lock().unwrap();
+        if let Ok(mut buf) = buffer_clone.lock() {
-
           for chunk in data.chunks(channels) {
             let mono: f32 = chunk.iter().sum::<f32>() / channels as f32;
             buf.push(mono);
@@ -364,6 +373,7 @@ impl AudioVisualizer {
             let drain_amt = buf.len() - chunk_size;
             buf.drain(0..drain_amt);
           }
+        }
       },
       |err| eprintln!("Stream error: {}", err),
       None,
@@ -371,11 +381,14 @@ impl AudioVisualizer {
 
     stream.play()?;
 
-    loop {
+    while !stop.load(Ordering::Relaxed) {
       thread::sleep(std::time::Duration::from_millis(16));
 
       let samples = {
-        let buf = audio_buffer.lock().unwrap();
+        let buf = match audio_buffer.lock() {
+          Ok(b) => b,
+          Err(_) => continue,
+        };
         if buf.len() < chunk_size {
           continue;
         }
@@ -384,7 +397,10 @@ impl AudioVisualizer {
 
       Self::apply_hann_window(&samples, &mut windowed_buffer);
 
-      fft.process(&mut windowed_buffer, &mut spectrum).unwrap();
+      if let Err(e) = fft.process(&mut windowed_buffer, &mut spectrum) {
+        eprintln!("FFT error: {:?}", e);
+        continue;
+      }
 
       let raw_bars = Self::process_fft_magnitudes(
         &spectrum,
@@ -407,6 +423,14 @@ impl AudioVisualizer {
         }
       }
     }
+
+    Ok(())
+  }
+}
+
+impl Drop for AudioVisualizer {
+  fn drop(&mut self) {
+    self.stop_flag.store(true, Ordering::Relaxed)
   }
 }