From 5985551abaf9418c1c37ef9fd6db39a6c9c94d69 Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Wed, 26 Jun 2024 21:51:06 +0800
Subject: feat: rag load websites (#655)

---
 src/utils/spinner.rs | 119 ++++++++++++++++++++++++++++-----------------------
 1 file changed, 66 insertions(+), 53 deletions(-)

(limited to 'src/utils/spinner.rs')

diff --git a/src/utils/spinner.rs b/src/utils/spinner.rs
index 8df4f57..a04a18f 100644
--- a/src/utils/spinner.rs
+++ b/src/utils/spinner.rs
@@ -5,46 +5,34 @@ use std::{
     io::{stdout, Write},
     time::Duration,
 };
-use tokio::{
-    sync::{mpsc, oneshot},
-    time::interval,
-};
+use tokio::{sync::mpsc, time::interval};
 
-pub struct Spinner {
+pub struct SpinnerInner {
     index: usize,
     message: String,
-    stopped: bool,
+    is_not_terminal: bool,
 }
 
-impl Spinner {
+impl SpinnerInner {
     const DATA: [&'static str; 10] = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"];
 
-    pub fn new(message: &str) -> Self {
-        Spinner {
+    fn new(message: &str) -> Self {
+        SpinnerInner {
             index: 0,
             message: message.to_string(),
-            stopped: false,
+            is_not_terminal: !stdout().is_terminal(),
         }
     }
 
-    pub fn set_message(&mut self, message: &str) {
-        self.message = format!(" {message}");
-    }
-
-    pub fn step(&mut self) -> Result<()> {
-        if self.stopped {
+    fn step(&mut self) -> Result<()> {
+        if self.is_not_terminal || self.message.is_empty() {
             return Ok(());
         }
         let mut writer = stdout();
         let frame = Self::DATA[self.index % Self::DATA.len()];
         let dots = ".".repeat((self.index / 5) % 4);
         let line = format!("{frame}{}{:<3}", self.message, dots);
-        queue!(
-            writer,
-            cursor::MoveToColumn(0),
-            terminal::Clear(terminal::ClearType::FromCursorDown),
-            style::Print(line),
-        )?;
+        queue!(writer, cursor::MoveToColumn(0), style::Print(line),)?;
         if self.index == 0 {
             queue!(writer, cursor::Hide)?;
         }
@@ -53,12 +41,20 @@ impl Spinner {
         Ok(())
     }
 
-    pub fn stop(&mut self) -> Result<()> {
-        if self.stopped {
+    fn set_message(&mut self, message: String) -> Result<()> {
+        self.clear_message()?;
+        if !message.is_empty() {
+            self.message = format!(" {message}");
+        }
+        Ok(())
+    }
+
+    fn clear_message(&mut self) -> Result<()> {
+        if self.is_not_terminal || self.message.is_empty() {
             return Ok(());
         }
+        self.message.clear();
         let mut writer = stdout();
-        self.stopped = true;
         queue!(
             writer,
             cursor::MoveToColumn(0),
@@ -70,43 +66,60 @@ impl Spinner {
     }
 }
 
-pub async fn run_spinner(message: &str) -> (oneshot::Sender<()>, mpsc::UnboundedSender<String>) {
+#[derive(Clone)]
+pub struct Spinner(mpsc::UnboundedSender<SpinnerEvent>);
+
+impl Drop for Spinner {
+    fn drop(&mut self) {
+        self.stop();
+    }
+}
+
+impl Spinner {
+    pub fn set_message(&self, message: String) -> Result<()> {
+        self.0.send(SpinnerEvent::SetMessage(message))?;
+        Ok(())
+    }
+
+    pub fn stop(&self) {
+        let _ = self.0.send(SpinnerEvent::Stop);
+    }
+}
+
+enum SpinnerEvent {
+    SetMessage(String),
+    Stop,
+}
+
+pub async fn create_spinner(message: &str) -> Spinner {
     let message = format!(" {message}");
-    let (stop_tx, stop_rx) = oneshot::channel();
-    let (message_tx, message_rx) = mpsc::unbounded_channel();
-    tokio::spawn(run_spinner_inner(message, stop_rx, message_rx));
-    (stop_tx, message_tx)
+    let (tx, rx) = mpsc::unbounded_channel();
+    tokio::spawn(run_spinner(message, rx));
+    Spinner(tx)
 }
 
-async fn run_spinner_inner(
-    message: String,
-    stop_rx: oneshot::Receiver<()>,
-    mut message_rx: mpsc::UnboundedReceiver<String>,
-) -> Result<()> {
-    let is_stdout_terminal = stdout().is_terminal();
-    let mut spinner = Spinner::new(&message);
+async fn run_spinner(message: String, mut rx: mpsc::UnboundedReceiver<SpinnerEvent>) -> Result<()> {
+    let mut spinner = SpinnerInner::new(&message);
     let mut interval = interval(Duration::from_millis(50));
-    tokio::select! {
-        _ = async {
-            loop {
-                tokio::select! {
-                    _ = interval.tick() => {
-                        if is_stdout_terminal {
-                            let _ = spinner.step();
+    loop {
+        tokio::select! {
+            _ = interval.tick() => {
+                let _ = spinner.step();
+            }
+            evt = rx.recv() => {
+                if let Some(evt) = evt {
+                    match evt {
+                        SpinnerEvent::SetMessage(message) => {
+                            spinner.set_message(message)?;
                         }
-                    }
-                    message = message_rx.recv() => {
-                        if let Some(message) = message {
-                            spinner.set_message(&message);
+                        SpinnerEvent::Stop => {
+                            spinner.clear_message()?;
+                            break;
                         }
                     }
+
                 }
             }
-        } => {}
-        _ = stop_rx => {
-            if is_stdout_terminal {
-                spinner.stop()?;
-            }
         }
     }
     Ok(())
-- 
cgit v1.2.3