From 11dc4d104b21b75d63e10e42ea4ee767c7bd5fdf Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Tue, 7 Mar 2023 11:51:52 +0800
Subject: refactor: optimize ctrl+c/ctrl+d abort handling (#27)

---
 src/client.rs | 15 +++++++--------
 1 file changed, 7 insertions(+), 8 deletions(-)

(limited to 'src/client.rs')

diff --git a/src/client.rs b/src/client.rs
index 589ad1a..19e6c27 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -1,13 +1,12 @@
 use crate::config::SharedConfig;
-use crate::repl::ReplyStreamHandler;
+use crate::repl::{ReplyStreamHandler, SharedAbortSignal};
 
 use anyhow::{anyhow, Context, Result};
 use eventsource_stream::Eventsource;
 use futures_util::StreamExt;
 use reqwest::{Client, Proxy, RequestBuilder};
 use serde_json::{json, Value};
-use std::sync::atomic::{AtomicBool, Ordering};
-use std::{sync::Arc, time::Duration};
+use std::time::Duration;
 use tokio::runtime::Runtime;
 use tokio::time::sleep;
 
@@ -43,27 +42,27 @@ impl ChatGptClient {
         prompt: Option<String>,
         handler: &mut ReplyStreamHandler,
     ) -> Result<()> {
-        async fn watch_ctrlc(ctrlc: Arc<AtomicBool>) {
+        async fn watch_abort(abort: SharedAbortSignal) {
             loop {
-                if ctrlc.load(Ordering::SeqCst) {
+                if abort.aborted() {
                     break;
                 }
                 sleep(Duration::from_millis(100)).await;
             }
         }
-        let ctrlc = handler.get_ctrlc();
+        let abort = handler.get_abort();
         self.runtime.block_on(async {
             tokio::select! {
                 ret = self.send_message_streaming_inner(input, prompt, handler) => {
                     handler.done();
                     ret.with_context(|| "Failed to send message streaming")
                 }
-                _ = watch_ctrlc(ctrlc.clone()) => {
+                _ = watch_abort(abort.clone()) => {
                     handler.done();
                     Ok(())
                  },
                 _ =  tokio::signal::ctrl_c() => {
-                    ctrlc.store(true, Ordering::SeqCst);
+                    abort.set_ctrlc();
                     Ok(())
                 }
             }
-- 
cgit v1.2.3