karyon_p2p/peer/
mod.rs

1mod connection;
2mod peer_id;
3
4use std::{
5    collections::HashSet,
6    fmt,
7    sync::{Arc, Weak},
8};
9
10use async_channel::{Receiver, Sender};
11use log::{error, trace};
12
13use karyon_core::{
14    async_runtime::Executor,
15    async_util::{TaskGroup, TaskResult},
16};
17
18use crate::{
19    conn_queue::QueuedConn,
20    endpoint::Endpoint,
21    peer_pool::PeerPool,
22    protocol::{PeerConn, ProtocolEvent, ProtocolID},
23    Config, Result,
24};
25
26pub use peer_id::PeerID;
27
28use connection::Wire;
29
30#[derive(Clone, Debug)]
31pub enum ConnDirection {
32    Inbound,
33    Outbound,
34}
35
36impl fmt::Display for ConnDirection {
37    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
38        match self {
39            ConnDirection::Inbound => write!(f, "Inbound"),
40            ConnDirection::Outbound => write!(f, "Outbound"),
41        }
42    }
43}
44
45/// A connected peer. Holds a `Wire` that hides the wire shape
46/// (single framed pipe vs. per-protocol streams).
47pub struct Peer {
48    id: PeerID,
49    peer_pool: Weak<PeerPool>,
50
51    direction: ConnDirection,
52    remote_endpoint: Endpoint,
53
54    connection: Arc<dyn Wire>,
55    disconnect_signal: Sender<Result<()>>,
56
57    negotiated_protocols: HashSet<ProtocolID>,
58    stop_chan: (Sender<Result<()>>, Receiver<Result<()>>),
59    config: Arc<Config>,
60    executor: Executor,
61    task_group: TaskGroup,
62}
63
64impl Peer {
65    pub async fn send(&self, proto_id: ProtocolID, msg: Vec<u8>) -> Result<()> {
66        self.connection.send(&proto_id, msg).await
67    }
68
69    pub async fn recv(&self, proto_id: &ProtocolID) -> Result<ProtocolEvent> {
70        self.connection.recv(proto_id).await
71    }
72
73    pub async fn broadcast(&self, proto_id: &ProtocolID, msg: Vec<u8>) {
74        self.peer_pool().broadcast(proto_id, msg).await;
75    }
76
77    pub fn id(&self) -> &PeerID {
78        &self.id
79    }
80
81    pub fn config(&self) -> Arc<Config> {
82        self.config.clone()
83    }
84
85    pub fn executor(&self) -> Executor {
86        self.executor.clone()
87    }
88
89    pub fn remote_endpoint(&self) -> &Endpoint {
90        &self.remote_endpoint
91    }
92
93    pub fn is_inbound(&self) -> bool {
94        matches!(self.direction, ConnDirection::Inbound)
95    }
96
97    pub fn direction(&self) -> &ConnDirection {
98        &self.direction
99    }
100
101    pub fn negotiated_protocols(&self) -> &HashSet<ProtocolID> {
102        &self.negotiated_protocols
103    }
104
105    pub(crate) async fn run(self: Arc<Self>) -> Result<()> {
106        self.run_connect_protocols().await;
107        let stop_signal = self.stop_chan.1.recv().await?;
108        stop_signal
109    }
110
111    pub(crate) async fn shutdown(self: &Arc<Self>) -> Result<()> {
112        trace!("peer {} shutting down", self.id);
113
114        let _ = self.connection.shutdown().await;
115        let _ = self.stop_chan.0.try_send(Ok(()));
116
117        let _ = self.disconnect_signal.send(Ok(())).await;
118        self.task_group.cancel().await;
119        Ok(())
120    }
121
122    async fn run_connect_protocols(self: &Arc<Self>) {
123        for (proto_id, constructor) in self.peer_pool().protocols.read().await.iter() {
124            if !self.negotiated_protocols.contains(proto_id) {
125                trace!("peer {} skip protocol {proto_id} (not negotiated)", self.id);
126                continue;
127            }
128            trace!("peer {} run protocol {proto_id}", self.id);
129
130            let peer_conn = PeerConn::new(self.clone(), proto_id.clone());
131            let protocol = match constructor(peer_conn) {
132                Ok(p) => p,
133                Err(err) => {
134                    error!("Failed to build protocol {proto_id}: {err}");
135                    continue;
136                }
137            };
138
139            let on_failure = {
140                let this = self.clone();
141                let proto_id = proto_id.clone();
142                |result: TaskResult<Result<()>>| async move {
143                    if let TaskResult::Completed(res) = result {
144                        if res.is_err() {
145                            error!("protocol {proto_id} stopped");
146                        }
147                        let _ = this.stop_chan.0.try_send(res);
148                    }
149                }
150            };
151
152            self.task_group.spawn_then(protocol.start(), on_failure);
153        }
154    }
155
156    fn peer_pool(&self) -> Arc<PeerPool> {
157        self.peer_pool.upgrade().unwrap()
158    }
159}
160
161impl Peer {
162    pub(crate) async fn new(
163        peer_pool: Arc<PeerPool>,
164        queued: QueuedConn,
165        id: PeerID,
166        negotiated_protocols: HashSet<ProtocolID>,
167        protocol_ids: impl IntoIterator<Item = ProtocolID> + Clone,
168    ) -> Result<Arc<Self>> {
169        let config = peer_pool.config.clone();
170        let executor = peer_pool.executor.clone();
171        let task_group = TaskGroup::with_executor(executor.clone());
172        let stop_chan = async_channel::bounded::<Result<()>>(1);
173
174        let remote_endpoint = queued.remote_endpoint.clone();
175        let direction = queued.direction.clone();
176        let disconnect_signal = queued.disconnect_signal.clone();
177
178        let connection = connection::from_queued(
179            queued,
180            &negotiated_protocols,
181            protocol_ids,
182            &task_group,
183            stop_chan.0.clone(),
184        )
185        .await?;
186
187        let peer_pool_weak = Arc::downgrade(&peer_pool);
188        Ok(Arc::new(Peer {
189            id,
190            peer_pool: peer_pool_weak,
191            direction,
192            remote_endpoint,
193            connection,
194            disconnect_signal,
195            negotiated_protocols,
196            stop_chan,
197            config,
198            executor,
199            task_group,
200        }))
201    }
202}
karyon_p2p/peer/mod.rs

karyon_p2p/peer/
mod.rs
