Skip to content

Commit

Permalink
Log when waiting for too long
Browse files Browse the repository at this point in the history
  • Loading branch information
willcrichton committed Sep 4, 2024
1 parent b81bb45 commit d6b5e17
Showing 1 changed file with 19 additions and 8 deletions.
27 changes: 19 additions & 8 deletions crates/server/src/main.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
use std::{
path::PathBuf,
pin::pin,
sync::{Arc, LazyLock},
time::{Duration, Instant},
};

use miniserve::{http::StatusCode, Content, Request, Response};
Expand Down Expand Up @@ -50,14 +52,23 @@ fn chatbot_thread() -> (mpsc::Sender<Payload>, mpsc::Sender<()>) {
while let Some((messages, responder)) = req_rx.recv().await {
let doc_paths = chatbot.retrieval_documents(&messages);
let docs = load_docs(doc_paths).await;
let chat_fut = chatbot.query_chat(&messages, &docs);
let cancel_fut = cancel_rx.recv();
tokio::select! {
response = chat_fut => {
responder.send(Some(response)).unwrap();
}
_ = cancel_fut => {
responder.send(None).unwrap();
let mut chat_fut = pin!(chatbot.query_chat(&messages, &docs));
let mut cancel_fut = pin!(cancel_rx.recv());
let start = Instant::now();
loop {
let log_fut = tokio::time::sleep(Duration::from_secs(1));
tokio::select! {
response = &mut chat_fut => {
responder.send(Some(response)).unwrap();
break;
}
_ = &mut cancel_fut => {
responder.send(None).unwrap();
break;
}
_ = log_fut => {
println!("Waiting for {} seconds", start.elapsed().as_secs());
}
}
}
}
Expand Down

0 comments on commit d6b5e17

Please sign in to comment.