From 1cc89eff514669d6459f62cc4eed8e0b34d8ef0c Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Tue, 23 Apr 2024 14:32:06 +0800
Subject: refactor: more async code (#427)

---
 src/render/stream.rs | 72 +++++++++++++++++++++++++---------------------------
 1 file changed, 35 insertions(+), 37 deletions(-)

(limited to 'src/render/stream.rs')

diff --git a/src/render/stream.rs b/src/render/stream.rs
index 9fdad5a..6007690 100644
--- a/src/render/stream.rs
+++ b/src/render/stream.rs
@@ -1,9 +1,8 @@
 use super::{MarkdownRender, ReplyEvent};
 
-use crate::utils::{AbortSignal, Spinner};
+use crate::utils::{run_spinner, AbortSignal};
 
 use anyhow::Result;
-use crossbeam::channel::Receiver;
 use crossterm::{
     cursor,
     event::{self, Event, KeyCode, KeyModifiers},
@@ -12,32 +11,32 @@ use crossterm::{
 };
 use std::{
     io::{self, stdout, Stdout, Write},
-    ops::Div,
-    time::{Duration, Instant},
+    time::Duration,
 };
 use textwrap::core::display_width;
+use tokio::sync::{mpsc::UnboundedReceiver, oneshot};
 
-pub fn markdown_stream(
-    rx: &Receiver<ReplyEvent>,
+pub async fn markdown_stream(
+    rx: UnboundedReceiver<ReplyEvent>,
     render: &mut MarkdownRender,
     abort: &AbortSignal,
 ) -> Result<()> {
     enable_raw_mode()?;
     let mut stdout = io::stdout();
 
-    let ret = markdown_stream_inner(rx, render, abort, &mut stdout);
+    let ret = markdown_stream_inner(rx, render, abort, &mut stdout).await;
 
     disable_raw_mode()?;
 
     ret
 }
 
-pub fn raw_stream(rx: &Receiver<ReplyEvent>, abort: &AbortSignal) -> Result<()> {
+pub async fn raw_stream(mut rx: UnboundedReceiver<ReplyEvent>, abort: &AbortSignal) -> Result<()> {
     loop {
         if abort.aborted() {
             return Ok(());
         }
-        if let Ok(evt) = rx.try_recv() {
+        if let Some(evt) = rx.recv().await {
             match evt {
                 ReplyEvent::Text(text) => {
                     print!("{}", text);
@@ -52,30 +51,29 @@ pub fn raw_stream(rx: &Receiver<ReplyEvent>, abort: &AbortSignal) -> Result<()>
     Ok(())
 }
 
-fn markdown_stream_inner(
-    rx: &Receiver<ReplyEvent>,
+async fn markdown_stream_inner(
+    mut rx: UnboundedReceiver<ReplyEvent>,
     render: &mut MarkdownRender,
     abort: &AbortSignal,
     writer: &mut Stdout,
 ) -> Result<()> {
-    let mut last_tick = Instant::now();
-    let tick_rate = Duration::from_millis(50);
-
     let mut buffer = String::new();
     let mut buffer_rows = 1;
 
     let columns = terminal::size()?.0;
 
-    let mut spinner = Spinner::new(" Generating");
+    let (spinner_tx, spinner_rx) = oneshot::channel();
+    let mut spinner_tx = Some(spinner_tx);
+    tokio::spawn(run_spinner(" Generating", spinner_rx));
 
     'outer: loop {
         if abort.aborted() {
             return Ok(());
         }
-        spinner.step(writer)?;
-
-        for reply_event in gather_events(rx) {
-            spinner.stop(writer)?;
+        for reply_event in gather_events(&mut rx).await {
+            if let Some(spinner_tx) = spinner_tx.take() {
+                let _ = spinner_tx.send(());
+            }
 
             match reply_event {
                 ReplyEvent::Text(mut text) => {
@@ -135,10 +133,7 @@ fn markdown_stream_inner(
             }
         }
 
-        let timeout = tick_rate
-            .checked_sub(last_tick.elapsed())
-            .unwrap_or_else(|| tick_rate.div(2));
-        if crossterm::event::poll(timeout)? {
+        if crossterm::event::poll(Duration::from_millis(25))? {
             if let Event::Key(key) = event::read()? {
                 match key.code {
                     KeyCode::Char('c') if key.modifiers == KeyModifiers::CONTROL => {
@@ -153,28 +148,31 @@ fn markdown_stream_inner(
                 }
             }
         }
-
-        if last_tick.elapsed() >= tick_rate {
-            last_tick = Instant::now();
-        }
     }
 
-    spinner.stop(writer)?;
-
+    if let Some(spinner_tx) = spinner_tx.take() {
+        let _ = spinner_tx.send(());
+    }
     Ok(())
 }
 
-fn gather_events(rx: &Receiver<ReplyEvent>) -> Vec<ReplyEvent> {
+async fn gather_events(rx: &mut UnboundedReceiver<ReplyEvent>) -> Vec<ReplyEvent> {
     let mut texts = vec![];
     let mut done = false;
-    for reply_event in rx.try_iter() {
-        match reply_event {
-            ReplyEvent::Text(v) => texts.push(v),
-            ReplyEvent::Done => {
-                done = true;
+    tokio::select! {
+        _ = async {
+            while let Some(reply_event) = rx.recv().await {
+                match reply_event {
+                    ReplyEvent::Text(v) => texts.push(v),
+                    ReplyEvent::Done => {
+                        done = true;
+                        break;
+                    }
+                }
             }
-        }
-    }
+        } => {}
+        _ = tokio::time::sleep(Duration::from_millis(50)) => {}
+    };
     let mut events = vec![];
     if !texts.is_empty() {
         events.push(ReplyEvent::Text(texts.join("")))
-- 
cgit v1.2.3