testnet/brain
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: add timeout for new VM request

Browse Source
This commit is contained in:
Noor 2025-05-05 16:46:32 +05:30
parent a1664d61e4
commit 8fa8a64862
Signed by: noormohammedb
GPG Key ID: D83EFB8B3B967146
2 changed files with 29 additions and 24 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
src/db.rs
View File

@ -1,4 +1,4 @@
use std::str::FromStr; use std::{str::FromStr, time::Duration};
pub use crate::constants::{ pub use crate::constants::{
ACCOUNT, ACTIVE_APP, ACTIVE_VM, DB_SCHEMA_FILE, DELETED_VM, ID_ALPHABET, NEW_VM_REQ, 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<DaemonNotification>), DaemonConnection(#[from] tokio::sync::mpsc::error::SendError<DaemonNotification>),
#[error(transparent)] #[error(transparent)]
StdIo(#[from] std::io::Error), StdIo(#[from] std::io::Error),
#[error(transparent)]
TimeOut(#[from] tokio::time::error::Elapsed),
} }
pub async fn db_connection( pub async fn db_connection(
@ -328,6 +330,7 @@ impl NewVmResp {
let mut args_stream = let mut args_stream =
resp.stream::<Notification<detee_shared::snp::pb::vm_proto::MeasurementArgs>>(1)?; resp.stream::<Notification<detee_shared::snp::pb::vm_proto::MeasurementArgs>>(1)?;
tokio::time::timeout(Duration::from_secs(10), async {
loop { loop {
tokio::select! { tokio::select! {
new_vm_req_notif = new_vm_stream.next() => { new_vm_req_notif = new_vm_stream.next() => {
@ -336,7 +339,7 @@ impl NewVmResp {
match new_vm_req_notif { match new_vm_req_notif {
Ok(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() { 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)); return Ok::<NewVmResp, Error>(Self::Error(vm_id.to_string(), new_vm_req_notif.data.error));
}; };
}, },
Err(e) => return Err(e.into()), Err(e) => return Err(e.into()),
@ -357,6 +360,7 @@ impl NewVmResp {
} }
} }
} }
}).await?
} }
} }

1
src/grpc.rs
View File

@ -625,6 +625,7 @@ impl BrainVmCli for BrainVmCliForReal {
new_vm_req.submit(&self.db).await?; new_vm_req.submit(&self.db).await?;
match oneshot_rx.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())), Ok(new_vm_resp) => Ok(Response::new(new_vm_resp?.into())),
Err(e) => { Err(e) => {
log::error!("Something weird happened. Reached error {e:?}"); log::error!("Something weird happened. Reached error {e:?}");