Skip to content

Commit

Permalink
[ampc] ability to send requests without timeouts (should be used spar…
Browse files Browse the repository at this point in the history
…ringly)
  • Loading branch information
mikkeldenker committed Dec 10, 2024
1 parent 852a10c commit bd9f1cb
Show file tree
Hide file tree
Showing 3 changed files with 21 additions and 2 deletions.
19 changes: 19 additions & 0 deletions crates/core/src/ampc/worker.rs
Original file line number Diff line number Diff line change
Expand Up @@ -98,6 +98,12 @@ where
Ok(res)
}

fn send_raw_without_timeout(&self, req: &JobReq<Self::Job>) -> Result<JobResp<Self::Job>> {
let mut conn = self.conn()?;
let res = block_on(conn.send_without_timeout(req))?;
Ok(res)
}

fn send<R>(&self, req: R) -> R::Response
where
R: RequestWrapper<<Self::Job as Job>::Worker>,
Expand All @@ -107,6 +113,19 @@ where
Resp::User(res) => R::unwrap_response(res).unwrap(),
}
}

fn send_without_timeout<R>(&self, req: R) -> R::Response
where
R: RequestWrapper<<Self::Job as Job>::Worker>,
{
match self
.send_raw_without_timeout(&Req::User(R::wrap(req)))
.unwrap()
{
Resp::Coordinator(_) => panic!("unexpected coordinator response"),
Resp::User(res) => R::unwrap_response(res).unwrap(),
}
}
}

pub trait RequestWrapper<W: Worker>: Message<W> {
Expand Down
2 changes: 1 addition & 1 deletion crates/core/src/distributed/sonic/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -119,7 +119,7 @@ where
}
}

async fn send_without_timeout(&mut self, request: &Req) -> Result<Res> {
pub async fn send_without_timeout(&mut self, request: &Req) -> Result<Res> {
self.awaiting_res = true;
let bytes = bincode::encode_to_vec(request, common::bincode_config()).unwrap();

Expand Down
2 changes: 1 addition & 1 deletion crates/core/src/entrypoint/ampc/shortest_path/worker.rs
Original file line number Diff line number Diff line change
Expand Up @@ -175,7 +175,7 @@ impl RemoteShortestPathWorker {
}

pub fn sample_nodes(&self, num_nodes: u64) -> Vec<webgraph::NodeID> {
self.send(SampleNodes(num_nodes))
self.send_without_timeout(SampleNodes(num_nodes))
}
}

Expand Down

0 comments on commit bd9f1cb

Please sign in to comment.