zeromq/
router.rs

1use crate::backend::GenericSocketBackend;
2use crate::codec::*;
3use crate::endpoint::Endpoint;
4use crate::error::{ZmqError, ZmqResult};
5use crate::fair_queue::FairQueue;
6use crate::message::*;
7use crate::transport::AcceptStopHandle;
8use crate::util::PeerIdentity;
9use crate::{MultiPeerBackend, SocketEvent, SocketOptions, SocketRecv, SocketSend, SocketType};
10use crate::{Socket, SocketBackend};
11
12use async_trait::async_trait;
13use futures_channel::mpsc;
14use futures_util::{SinkExt, StreamExt};
15
16use std::collections::HashMap;
17use std::convert::TryInto;
18use std::sync::Arc;
19
20pub struct RouterSocket {
21    backend: Arc<GenericSocketBackend>,
22    binds: HashMap<Endpoint, AcceptStopHandle>,
23    fair_queue: FairQueue<ZmqFramedRead, PeerIdentity>,
24}
25
26impl Drop for RouterSocket {
27    fn drop(&mut self) {
28        self.backend.shutdown();
29    }
30}
31
32#[async_trait]
33impl Socket for RouterSocket {
34    fn with_options(options: SocketOptions) -> Self {
35        let fair_queue = FairQueue::new(true);
36        Self {
37            backend: Arc::new(GenericSocketBackend::with_options(
38                Some(fair_queue.inner()),
39                SocketType::ROUTER,
40                options,
41            )),
42            binds: HashMap::new(),
43            fair_queue,
44        }
45    }
46
47    fn backend(&self) -> Arc<dyn MultiPeerBackend> {
48        self.backend.clone()
49    }
50
51    fn binds(&mut self) -> &mut HashMap<Endpoint, AcceptStopHandle> {
52        &mut self.binds
53    }
54
55    fn monitor(&mut self) -> mpsc::Receiver<SocketEvent> {
56        let (sender, receiver) = mpsc::channel(1024);
57        self.backend.socket_monitor.lock().replace(sender);
58        receiver
59    }
60}
61
62#[async_trait]
63impl SocketRecv for RouterSocket {
64    async fn recv(&mut self) -> ZmqResult<ZmqMessage> {
65        loop {
66            match self.fair_queue.next().await {
67                Some((peer_id, Ok(Message::Message(mut message)))) => {
68                    message.push_front(peer_id.into());
69                    return Ok(message);
70                }
71                Some((_peer_id, Ok(msg))) => todo!("Unimplemented message: {:?}", msg),
72                Some((peer_id, Err(_))) => {
73                    self.backend.peer_disconnected(&peer_id);
74                }
75                None => todo!(),
76            };
77        }
78    }
79}
80
81#[async_trait]
82impl SocketSend for RouterSocket {
83    async fn send(&mut self, mut message: ZmqMessage) -> ZmqResult<()> {
84        assert!(message.len() > 1);
85        let peer_id: PeerIdentity = message.pop_front().unwrap().try_into()?;
86        match self.backend.peers.get_mut(&peer_id) {
87            Some(mut peer) => {
88                peer.send_queue.send(Message::Message(message)).await?;
89                Ok(())
90            }
91            None => Err(ZmqError::Other("Destination client not found by identity")),
92        }
93    }
94}
zeromq/router.rs

zeromq/
router.rs
