diff --git a/src/db.rs b/src/db.rs index 995ded9..a20d86e 100644 --- a/src/db.rs +++ b/src/db.rs @@ -1,4 +1,4 @@ -use std::str::FromStr; +use std::{str::FromStr, time::Duration}; pub use crate::constants::{ ACCOUNT, ACTIVE_APP, ACTIVE_VM, DB_SCHEMA_FILE, DELETED_VM, ID_ALPHABET, NEW_VM_REQ, @@ -24,6 +24,8 @@ pub enum Error { DaemonConnection(#[from] tokio::sync::mpsc::error::SendError), #[error(transparent)] StdIo(#[from] std::io::Error), + #[error(transparent)] + TimeOut(#[from] tokio::time::error::Elapsed), } pub async fn db_connection( @@ -328,35 +330,37 @@ impl NewVmResp { let mut args_stream = resp.stream::>(1)?; - loop { - tokio::select! { - new_vm_req_notif = new_vm_stream.next() => { - log::debug!("Got stream 1..."); - if let Some(new_vm_req_notif) = new_vm_req_notif { - match new_vm_req_notif { - Ok(new_vm_req_notif) => { - if new_vm_req_notif.action == surrealdb::Action::Update && !new_vm_req_notif.data.error.is_empty() { - return Ok(Self::Error(vm_id.to_string(), new_vm_req_notif.data.error)); - }; - }, - Err(e) => return Err(e.into()), + tokio::time::timeout(Duration::from_secs(10), async { + loop { + tokio::select! { + new_vm_req_notif = new_vm_stream.next() => { + log::debug!("Got stream 1..."); + if let Some(new_vm_req_notif) = new_vm_req_notif { + match new_vm_req_notif { + Ok(new_vm_req_notif) => { + if new_vm_req_notif.action == surrealdb::Action::Update && !new_vm_req_notif.data.error.is_empty() { + return Ok::(Self::Error(vm_id.to_string(), new_vm_req_notif.data.error)); + }; + }, + Err(e) => return Err(e.into()), + } } } - } - args_notif = args_stream.next() => { - if let Some(args_notif) = args_notif { - match args_notif { - Ok(args_notif) => { - if args_notif.action == surrealdb::Action::Create { - return Ok(Self::Args(vm_id.to_string(), args_notif.data)); - }; - }, - Err(e) => return Err(e.into()), + args_notif = args_stream.next() => { + if let Some(args_notif) = args_notif { + match args_notif { + Ok(args_notif) => { + if args_notif.action == surrealdb::Action::Create { + return Ok(Self::Args(vm_id.to_string(), args_notif.data)); + }; + }, + Err(e) => return Err(e.into()), + } } } } } - } + }).await? } } diff --git a/src/grpc.rs b/src/grpc.rs index 0877ea8..2114a9b 100644 --- a/src/grpc.rs +++ b/src/grpc.rs @@ -625,6 +625,7 @@ impl BrainVmCli for BrainVmCliForReal { new_vm_req.submit(&self.db).await?; match oneshot_rx.await { + Ok(Err(db::Error::TimeOut(_))) => Err(Status::deadline_exceeded("Request failed due to timeout. Please try again later or contact the DeTEE devs team.")), Ok(new_vm_resp) => Ok(Response::new(new_vm_resp?.into())), Err(e) => { log::error!("Something weird happened. Reached error {e:?}");