1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
|
use std::collections::HashMap;
use async_channel::{Receiver, Sender};
use karyon_core::async_runtime::lock::Mutex;
use crate::{message, Error, Result};
use super::RequestID;
/// Manages client requests
pub(super) struct MessageDispatcher {
chans: Mutex<HashMap<RequestID, Sender<message::Response>>>,
}
impl MessageDispatcher {
/// Creates a new MessageDispatcher
pub(super) fn new() -> Self {
Self {
chans: Mutex::new(HashMap::new()),
}
}
/// Registers a new request with a given ID and returns a Receiver channel
/// to wait for the response.
pub(super) async fn register(&self, id: RequestID) -> Receiver<message::Response> {
let (tx, rx) = async_channel::bounded(1);
self.chans.lock().await.insert(id, tx);
rx
}
/// Unregisters the request with the provided ID
pub(super) async fn unregister(&self, id: &RequestID) {
self.chans.lock().await.remove(id);
}
/// Clear the registered channels.
pub(super) async fn clear(&self) {
let mut chans = self.chans.lock().await;
for (_, tx) in chans.iter() {
tx.close();
}
chans.clear();
}
/// Dispatches a response to the channel associated with the response's ID.
///
/// If a channel is registered for the response's ID, the response is sent
/// through that channel. If no channel is found for the ID, returns an error.
pub(super) async fn dispatch(&self, res: message::Response) -> Result<()> {
let res_id = match res.id {
Some(ref rid) => rid.clone(),
None => {
return Err(Error::InvalidMsg("Response id is none"));
}
};
let id: RequestID = serde_json::from_value(res_id)?;
let val = self.chans.lock().await.remove(&id);
match val {
Some(tx) => tx.send(res).await.map_err(Error::from),
None => Err(Error::InvalidMsg("Receive unknown message")),
}
}
}
|
