1use crate::{
2 api::{ConnectAddr, ParticipantError, ParticipantEvent, Stream},
3 channel::{Protocols, ProtocolsError, RecvProtocols, SendProtocols},
4 metrics::NetworkMetrics,
5 util::DeferredTracer,
6};
7use bytes::Bytes;
8use futures_util::{FutureExt, StreamExt};
9use hashbrown::HashMap;
10use network_protocol::{
11 _internal::SortedVec, Bandwidth, Cid, Pid, Prio, Promises, ProtocolEvent, RecvProtocol,
12 SendProtocol, Sid,
13};
14use std::{
15 sync::{
16 Arc,
17 atomic::{AtomicBool, AtomicI32, Ordering},
18 },
19 time::{Duration, Instant},
20};
21use tokio::{
22 select,
23 sync::{Mutex, RwLock, mpsc, oneshot, watch},
24 task::JoinHandle,
25};
26use tokio_stream::wrappers::UnboundedReceiverStream;
27use tracing::*;
28
29pub(crate) type A2bStreamOpen = (Prio, Promises, Bandwidth, oneshot::Sender<Stream>);
30pub(crate) type S2bCreateChannel = (Cid, Sid, Protocols, ConnectAddr, oneshot::Sender<()>);
31pub(crate) type S2bShutdownBparticipant = (Duration, oneshot::Sender<Result<(), ParticipantError>>);
32pub(crate) type B2sPrioStatistic = (Pid, u64, u64);
33
34#[derive(Debug)]
35#[expect(dead_code)]
36struct ChannelInfo {
37 cid: Cid,
38 cid_string: String, remote_con_addr: ConnectAddr,
40}
41
42#[derive(Debug)]
43struct StreamInfo {
44 #[expect(dead_code)]
45 prio: Prio,
46 #[expect(dead_code)]
47 promises: Promises,
48 send_closed: Arc<AtomicBool>,
49 b2a_msg_recv_s: Mutex<async_channel::Sender<Bytes>>,
50}
51
52#[derive(Debug)]
53struct ControlChannels {
54 a2b_open_stream_r: mpsc::UnboundedReceiver<A2bStreamOpen>,
55 b2a_stream_opened_s: mpsc::UnboundedSender<Stream>,
56 b2a_event_s: mpsc::UnboundedSender<ParticipantEvent>,
57 s2b_create_channel_r: mpsc::UnboundedReceiver<S2bCreateChannel>,
58 b2a_bandwidth_stats_s: watch::Sender<f32>,
59 s2b_shutdown_bparticipant_r: oneshot::Receiver<S2bShutdownBparticipant>, }
61
62#[derive(Debug)]
63struct OpenStreamInfo {
64 a2b_msg_s: crossbeam_channel::Sender<(Sid, Bytes)>,
65 a2b_close_stream_s: mpsc::UnboundedSender<Sid>,
66}
67
68#[derive(Debug)]
69pub struct BParticipant {
70 local_pid: Pid, remote_pid: Pid,
72 remote_pid_string: String, offset_sid: Sid,
74 channels: Arc<RwLock<HashMap<Cid, Mutex<ChannelInfo>>>>,
75 streams: RwLock<HashMap<Sid, StreamInfo>>,
76 run_channels: Option<ControlChannels>,
77 shutdown_barrier: AtomicI32,
78 metrics: Arc<NetworkMetrics>,
79 open_stream_channels: Arc<Mutex<Option<OpenStreamInfo>>>,
80 output_limit: usize,
81}
82
83impl BParticipant {
84 const BARR_CHANNEL: i32 = 1;
86 const BARR_RECV: i32 = 4;
87 const BARR_SEND: i32 = 2;
88 const TICK_TIME: Duration = Duration::from_millis(Self::TICK_TIME_MS);
89 const TICK_TIME_MS: u64 = 5;
90
91 pub(crate) fn new(
92 local_pid: Pid,
93 remote_pid: Pid,
94 offset_sid: Sid,
95 metrics: Arc<NetworkMetrics>,
96 output_limit: usize,
97 ) -> (
98 Self,
99 mpsc::UnboundedSender<A2bStreamOpen>,
100 mpsc::UnboundedReceiver<Stream>,
101 mpsc::UnboundedReceiver<ParticipantEvent>,
102 mpsc::UnboundedSender<S2bCreateChannel>,
103 oneshot::Sender<S2bShutdownBparticipant>,
104 watch::Receiver<f32>,
105 ) {
106 let (a2b_open_stream_s, a2b_open_stream_r) = mpsc::unbounded_channel::<A2bStreamOpen>();
107 let (b2a_stream_opened_s, b2a_stream_opened_r) = mpsc::unbounded_channel::<Stream>();
108 let (b2a_event_s, b2a_event_r) = mpsc::unbounded_channel::<ParticipantEvent>();
109 let (s2b_shutdown_bparticipant_s, s2b_shutdown_bparticipant_r) = oneshot::channel();
110 let (s2b_create_channel_s, s2b_create_channel_r) = mpsc::unbounded_channel();
111 let (b2a_bandwidth_stats_s, b2a_bandwidth_stats_r) = watch::channel::<f32>(0.0);
112
113 let run_channels = Some(ControlChannels {
114 a2b_open_stream_r,
115 b2a_stream_opened_s,
116 b2a_event_s,
117 s2b_create_channel_r,
118 b2a_bandwidth_stats_s,
119 s2b_shutdown_bparticipant_r,
120 });
121
122 (
123 Self {
124 local_pid,
125 remote_pid,
126 remote_pid_string: remote_pid.to_string(),
127 offset_sid,
128 channels: Arc::new(RwLock::new(HashMap::new())),
129 streams: RwLock::new(HashMap::new()),
130 shutdown_barrier: AtomicI32::new(
131 Self::BARR_CHANNEL + Self::BARR_SEND + Self::BARR_RECV,
132 ),
133 run_channels,
134 metrics,
135 open_stream_channels: Arc::new(Mutex::new(None)),
136 output_limit,
137 },
138 a2b_open_stream_s,
139 b2a_stream_opened_r,
140 b2a_event_r,
141 s2b_create_channel_s,
142 s2b_shutdown_bparticipant_s,
143 b2a_bandwidth_stats_r,
144 )
145 }
146
147 pub async fn run(mut self, b2s_prio_statistic_s: mpsc::UnboundedSender<B2sPrioStatistic>) {
148 let (b2b_add_send_protocol_s, b2b_add_send_protocol_r) =
149 mpsc::unbounded_channel::<(Cid, SendProtocols)>();
150 let (b2b_add_recv_protocol_s, b2b_add_recv_protocol_r) =
151 mpsc::unbounded_channel::<(Cid, RecvProtocols)>();
152 let (b2b_close_send_protocol_s, b2b_close_send_protocol_r) =
153 async_channel::unbounded::<Cid>();
154 let (b2b_force_close_recv_protocol_s, b2b_force_close_recv_protocol_r) =
155 async_channel::unbounded::<Cid>();
156 let (b2b_notify_send_of_recv_open_s, b2b_notify_send_of_recv_open_r) =
157 crossbeam_channel::unbounded::<(Cid, Sid, Prio, Promises, u64)>();
158 let (b2b_notify_send_of_recv_close_s, b2b_notify_send_of_recv_close_r) =
159 crossbeam_channel::unbounded::<(Cid, Sid)>();
160
161 let (a2b_close_stream_s, a2b_close_stream_r) = mpsc::unbounded_channel::<Sid>();
162 let (a2b_msg_s, a2b_msg_r) = crossbeam_channel::unbounded::<(Sid, Bytes)>();
163
164 *self.open_stream_channels.lock().await = Some(OpenStreamInfo {
165 a2b_msg_s,
166 a2b_close_stream_s,
167 });
168 let run_channels = self.run_channels.take().unwrap();
169 trace!("start all managers");
170 tokio::join!(
171 self.send_mgr(
172 run_channels.a2b_open_stream_r,
173 a2b_close_stream_r,
174 a2b_msg_r,
175 b2b_add_send_protocol_r,
176 b2b_close_send_protocol_r,
177 b2b_notify_send_of_recv_open_r,
178 b2b_notify_send_of_recv_close_r,
179 run_channels.b2a_event_s.clone(),
180 b2s_prio_statistic_s,
181 run_channels.b2a_bandwidth_stats_s,
182 )
183 .instrument(tracing::info_span!("send")),
184 self.recv_mgr(
185 run_channels.b2a_stream_opened_s,
186 b2b_add_recv_protocol_r,
187 b2b_force_close_recv_protocol_r,
188 b2b_close_send_protocol_s.clone(),
189 b2b_notify_send_of_recv_open_s,
190 b2b_notify_send_of_recv_close_s,
191 )
192 .instrument(tracing::info_span!("recv")),
193 self.create_channel_mgr(
194 run_channels.s2b_create_channel_r,
195 b2b_add_send_protocol_s,
196 b2b_add_recv_protocol_s,
197 run_channels.b2a_event_s,
198 ),
199 self.participant_shutdown_mgr(
200 run_channels.s2b_shutdown_bparticipant_r,
201 b2b_close_send_protocol_s.clone(),
202 b2b_force_close_recv_protocol_s,
203 ),
204 );
205 }
206
207 fn best_protocol(all: &SortedVec<Cid, SendProtocols>, promises: Promises) -> Option<Cid> {
208 all.data.iter().find(|(_, p)| matches!(p, SendProtocols::Mpsc(_))).map(|(c, _)| *c).or_else(
210 || if network_protocol::TcpSendProtocol::<crate::channel::TcpDrain>::supported_promises()
211 .contains(promises)
212 {
213 all.data.iter().find(|(_, p)| matches!(p, SendProtocols::Tcp(_))).map(|(c, _)| *c)
215 } else {
216 None
217 }
218 ).or_else(
219 || if network_protocol::QuicSendProtocol::<crate::channel::QuicDrain>::supported_promises()
221 .contains(promises)
222 {
223 all.data.iter().find(|(_, p)| matches!(p, SendProtocols::Quic(_))).map(|(c, _)| *c)
224 } else {
225 None
226 }
227 ).or_else(
228 || {
229 warn!("couldn't satisfy promises");
230 all.data.first().map(|(c, _)| *c)
231 }
232 )
233 }
234
235 async fn send_mgr(
237 &self,
238 mut a2b_open_stream_r: mpsc::UnboundedReceiver<A2bStreamOpen>,
239 mut a2b_close_stream_r: mpsc::UnboundedReceiver<Sid>,
240 a2b_msg_r: crossbeam_channel::Receiver<(Sid, Bytes)>,
241 mut b2b_add_protocol_r: mpsc::UnboundedReceiver<(Cid, SendProtocols)>,
242 b2b_close_send_protocol_r: async_channel::Receiver<Cid>,
243 b2b_notify_send_of_recv_open_r: crossbeam_channel::Receiver<(
244 Cid,
245 Sid,
246 Prio,
247 Promises,
248 Bandwidth,
249 )>,
250 b2b_notify_send_of_recv_close_r: crossbeam_channel::Receiver<(Cid, Sid)>,
251 b2a_event_s: mpsc::UnboundedSender<ParticipantEvent>,
252 _b2s_prio_statistic_s: mpsc::UnboundedSender<B2sPrioStatistic>,
253 b2a_bandwidth_stats_s: watch::Sender<f32>,
254 ) {
255 let mut sorted_send_protocols = SortedVec::<Cid, SendProtocols>::default();
256 let mut sorted_stream_protocols = SortedVec::<Sid, Cid>::default();
257 let mut interval = tokio::time::interval(Self::TICK_TIME);
258 let mut last_instant = Instant::now();
259 let mut stream_ids = self.offset_sid;
260 let mut part_bandwidth = 0.0f32;
261 trace!("workaround, actively wait for first protocol");
262 if let Some((c, p)) = b2b_add_protocol_r.recv().await {
263 sorted_send_protocols.insert(c, p)
264 }
265 loop {
266 let (open, close, _, addp, remp) = select!(
267 Some(n) = a2b_open_stream_r.recv().fuse() => (Some(n), None, None, None, None),
268 Some(n) = a2b_close_stream_r.recv().fuse() => (None, Some(n), None, None, None),
269 _ = interval.tick() => (None, None, Some(()), None, None),
270 Some(n) = b2b_add_protocol_r.recv().fuse() => (None, None, None, Some(n), None),
271 Ok(n) = b2b_close_send_protocol_r.recv().fuse() => (None, None, None, None, Some(n)),
272 );
273
274 if let Some((cid, p)) = addp {
275 debug!(?cid, "add protocol");
276 sorted_send_protocols.insert(cid, p);
277 }
278
279 if sorted_send_protocols.data.is_empty() {
281 warn!("no channel");
282 tokio::time::sleep(Self::TICK_TIME * 1000).await; continue;
284 }
285
286 let mut cid = u64::MAX;
289
290 let active_err = async {
291 if let Some((prio, promises, guaranteed_bandwidth, return_s)) = open {
292 let sid = stream_ids;
293 stream_ids += Sid::from(1);
294 cid = Self::best_protocol(&sorted_send_protocols, promises).unwrap();
295 trace!(?sid, ?cid, "open stream");
296
297 let stream = self
298 .create_stream(sid, prio, promises, guaranteed_bandwidth, self.output_limit)
299 .await;
300
301 let event = ProtocolEvent::OpenStream {
302 sid,
303 prio,
304 promises,
305 guaranteed_bandwidth,
306 };
307
308 sorted_stream_protocols.insert(sid, cid);
309 return_s.send(stream).unwrap();
310 sorted_send_protocols
311 .get_mut(&cid)
312 .unwrap()
313 .send(event)
314 .await?;
315 }
316
317 for (cid, sid, prio, promises, guaranteed_bandwidth) in
319 b2b_notify_send_of_recv_open_r.try_iter()
320 {
321 match sorted_send_protocols.get_mut(&cid) {
322 Some(p) => {
323 sorted_stream_protocols.insert(sid, cid);
324 p.notify_from_recv(ProtocolEvent::OpenStream {
325 sid,
326 prio,
327 promises,
328 guaranteed_bandwidth,
329 });
330 },
331 None => warn!(?cid, "couldn't notify create protocol, doesn't exist"),
332 };
333 }
334
335 for (sid, buffer) in a2b_msg_r.try_iter() {
337 cid = *sorted_stream_protocols.get(&sid).unwrap();
338 let event = ProtocolEvent::Message { data: buffer, sid };
339 sorted_send_protocols
340 .get_mut(&cid)
341 .unwrap()
342 .send(event)
343 .await?;
344 }
345
346 for (cid, sid) in b2b_notify_send_of_recv_close_r.try_iter() {
348 match sorted_send_protocols.get_mut(&cid) {
349 Some(p) => {
350 let _ = sorted_stream_protocols.delete(&sid);
351 p.notify_from_recv(ProtocolEvent::CloseStream { sid });
352 },
353 None => warn!(?cid, "couldn't notify close protocol, doesn't exist"),
354 };
355 }
356
357 if let Some(sid) = close {
358 trace!(?stream_ids, "delete stream");
359 self.delete_stream(sid).await;
360 if let Some(c) = sorted_stream_protocols.delete(&sid) {
363 cid = c;
364 let event = ProtocolEvent::CloseStream { sid };
365 sorted_send_protocols
366 .get_mut(&c)
367 .unwrap()
368 .send(event)
369 .await?;
370 }
371 }
372
373 let send_time = Instant::now();
374 let diff = send_time.duration_since(last_instant);
375 last_instant = send_time;
376 let mut cnt = 0;
377 for (c, p) in sorted_send_protocols.data.iter_mut() {
378 cid = *c;
379 cnt += p.flush(1_000_000_000, diff).await?; }
381 let flush_time = send_time.elapsed().as_secs_f32();
382 part_bandwidth = 0.99 * part_bandwidth + 0.01 * (cnt as f32 / flush_time);
383 self.metrics
384 .participant_bandwidth(&self.remote_pid_string, part_bandwidth);
385 let _ = b2a_bandwidth_stats_s.send(part_bandwidth);
386 let r: Result<(), network_protocol::ProtocolError<ProtocolsError>> = Ok(());
387 r
388 }
389 .await;
390 if let Err(e) = active_err {
391 info!(?cid, ?e, "protocol failed, shutting down channel");
392 trace!("TODO: for now decide to FAIL this participant and not wait for a failover");
395 sorted_send_protocols.delete(&cid).unwrap();
396 if let Some(info) = self.channels.write().await.get(&cid)
397 && let Err(e) = b2a_event_s.send(ParticipantEvent::ChannelDeleted(
398 info.lock().await.remote_con_addr.clone(),
399 ))
400 {
401 debug!(?e, "Participant was dropped during channel disconnect");
402 };
403 self.metrics.channels_disconnected(&self.remote_pid_string);
404 if sorted_send_protocols.data.is_empty() {
405 break;
406 }
407 }
408
409 if let Some(cid) = remp {
410 debug!(?cid, "remove protocol");
411 match sorted_send_protocols.delete(&cid) {
412 Some(mut prot) => {
413 if let Some(info) = self.channels.write().await.get(&cid)
414 && let Err(e) = b2a_event_s.send(ParticipantEvent::ChannelDeleted(
415 info.lock().await.remote_con_addr.clone(),
416 ))
417 {
418 debug!(?e, "Participant was dropped during channel disconnect");
419 };
420 self.metrics.channels_disconnected(&self.remote_pid_string);
421 trace!("blocking flush");
422 let _ = prot.flush(u64::MAX, Duration::from_secs(1)).await;
423 trace!("shutdown prot");
424 let _ = prot.send(ProtocolEvent::Shutdown).await;
425 },
426 None => trace!("tried to remove protocol twice"),
427 };
428 if sorted_send_protocols.data.is_empty() {
429 break;
430 }
431 }
432 }
433 trace!("stop sending in api!");
434 self.open_stream_channels.lock().await.take();
435 trace!("Stop send_mgr");
436 self.shutdown_barrier
437 .fetch_sub(Self::BARR_SEND, Ordering::SeqCst);
438 }
439
440 async fn recv_mgr(
441 &self,
442 b2a_stream_opened_s: mpsc::UnboundedSender<Stream>,
443 mut b2b_add_protocol_r: mpsc::UnboundedReceiver<(Cid, RecvProtocols)>,
444 b2b_force_close_recv_protocol_r: async_channel::Receiver<Cid>,
445 b2b_close_send_protocol_s: async_channel::Sender<Cid>,
446 b2b_notify_send_of_recv_open_r: crossbeam_channel::Sender<(
447 Cid,
448 Sid,
449 Prio,
450 Promises,
451 Bandwidth,
452 )>,
453 b2b_notify_send_of_recv_close_s: crossbeam_channel::Sender<(Cid, Sid)>,
454 ) {
455 let mut recv_protocols: HashMap<Cid, JoinHandle<()>> = HashMap::new();
456 let (hacky_recv_s, mut hacky_recv_r) = mpsc::unbounded_channel();
458
459 let retrigger = |cid: Cid, mut p: RecvProtocols, map: &mut HashMap<_, _>| {
460 let hacky_recv_s = hacky_recv_s.clone();
461 let handle = tokio::spawn(async move {
462 let r = p.recv().await;
463 let _ = hacky_recv_s.send((cid, r, p)); });
465 map.insert(cid, handle);
466 };
467
468 let remove_c = |recv_protocols: &mut HashMap<Cid, JoinHandle<()>>, cid: &Cid| {
469 match recv_protocols.remove(cid) {
470 Some(h) => {
471 h.abort();
472 debug!(?cid, "remove protocol");
473 },
474 None => trace!("tried to remove protocol twice"),
475 };
476 recv_protocols.is_empty()
477 };
478
479 let mut defered_orphan = DeferredTracer::new(Level::WARN);
480
481 loop {
482 let (event, addp, remp) = select!(
483 Some(n) = hacky_recv_r.recv().fuse() => (Some(n), None, None),
484 Some(n) = b2b_add_protocol_r.recv().fuse() => (None, Some(n), None),
485 Ok(n) = b2b_force_close_recv_protocol_r.recv().fuse() => (None, None, Some(n)),
486 else => {
487 error!("recv_mgr -> something is seriously wrong!, end recv_mgr");
488 break;
489 }
490 );
491
492 if let Some((cid, p)) = addp {
493 debug!(?cid, "add protocol");
494 retrigger(cid, p, &mut recv_protocols);
495 };
496 if let Some(cid) = remp {
497 if remove_c(&mut recv_protocols, &cid) {
499 break;
500 }
501 };
502
503 if let Some((cid, r, p)) = event {
504 match r {
505 Ok(ProtocolEvent::OpenStream {
506 sid,
507 prio,
508 promises,
509 guaranteed_bandwidth,
510 }) => {
511 trace!(?sid, "open stream");
512 let _ = b2b_notify_send_of_recv_open_r.send((
513 cid,
514 sid,
515 prio,
516 promises,
517 guaranteed_bandwidth,
518 ));
519 let stream = self
522 .create_stream(
523 sid,
524 prio,
525 promises,
526 guaranteed_bandwidth,
527 self.output_limit,
528 )
529 .await;
530 b2a_stream_opened_s.send(stream).unwrap();
531 retrigger(cid, p, &mut recv_protocols);
532 },
533 Ok(ProtocolEvent::CloseStream { sid }) => {
534 trace!(?sid, "close stream");
535 let _ = b2b_notify_send_of_recv_close_s.send((cid, sid));
536 self.delete_stream(sid).await;
537 retrigger(cid, p, &mut recv_protocols);
538 },
539 Ok(ProtocolEvent::Message { data, sid }) => {
540 let lock = self.streams.read().await;
541 match lock.get(&sid) {
542 Some(stream) => {
543 let _ = stream.b2a_msg_recv_s.lock().await.send(data).await;
544 },
545 None => defered_orphan.log(sid),
546 };
547 retrigger(cid, p, &mut recv_protocols);
548 },
549 Ok(ProtocolEvent::Shutdown) => {
550 info!(?cid, "shutdown protocol");
551 if let Err(e) = b2b_close_send_protocol_s.send(cid).await {
552 debug!(?e, ?cid, "send_mgr was already closed simultaneously");
553 }
554 if remove_c(&mut recv_protocols, &cid) {
555 break;
556 }
557 },
558 Err(e) => {
559 info!(?e, ?cid, "protocol failed, shutting down channel");
560 if let Err(e) = b2b_close_send_protocol_s.send(cid).await {
561 debug!(?e, ?cid, "send_mgr was already closed simultaneously");
562 }
563 if remove_c(&mut recv_protocols, &cid) {
564 break;
565 }
566 },
567 }
568 }
569
570 if let Some(table) = defered_orphan.print() {
571 for (sid, cnt) in table.iter() {
572 warn!(?sid, ?cnt, "recv messages with orphan stream");
573 }
574 }
575 }
576 trace!("receiving no longer possible, closing all streams");
577 for (_, si) in self.streams.write().await.drain() {
578 si.send_closed.store(true, Ordering::SeqCst);
579 self.metrics.streams_closed(&self.remote_pid_string);
580 }
581 trace!("Stop recv_mgr");
582 self.shutdown_barrier
583 .fetch_sub(Self::BARR_RECV, Ordering::SeqCst);
584 }
585
586 async fn create_channel_mgr(
587 &self,
588 s2b_create_channel_r: mpsc::UnboundedReceiver<S2bCreateChannel>,
589 b2b_add_send_protocol_s: mpsc::UnboundedSender<(Cid, SendProtocols)>,
590 b2b_add_recv_protocol_s: mpsc::UnboundedSender<(Cid, RecvProtocols)>,
591 b2a_event_s: mpsc::UnboundedSender<ParticipantEvent>,
592 ) {
593 let s2b_create_channel_r = UnboundedReceiverStream::new(s2b_create_channel_r);
594 s2b_create_channel_r
595 .for_each_concurrent(
596 None,
597 |(cid, _, protocol, remote_con_addr, b2s_create_channel_done_s)| {
598 let channels = Arc::clone(&self.channels);
601 let b2b_add_send_protocol_s = b2b_add_send_protocol_s.clone();
602 let b2b_add_recv_protocol_s = b2b_add_recv_protocol_s.clone();
603 let b2a_event_s = b2a_event_s.clone();
604 async move {
605 let mut lock = channels.write().await;
606 let mut channel_no = lock.len();
607 lock.insert(
608 cid,
609 Mutex::new(ChannelInfo {
610 cid,
611 cid_string: cid.to_string(),
612 remote_con_addr: remote_con_addr.clone(),
613 }),
614 );
615 drop(lock);
616 let (send, recv) = protocol.split();
617 b2b_add_send_protocol_s.send((cid, send)).unwrap();
618 b2b_add_recv_protocol_s.send((cid, recv)).unwrap();
619 if let Err(e) =
620 b2a_event_s.send(ParticipantEvent::ChannelCreated(remote_con_addr))
621 {
622 debug!(?e, "Participant was dropped during channel connect");
623 };
624 b2s_create_channel_done_s.send(()).unwrap();
625 if channel_no > 5 {
626 debug!(?channel_no, "metrics will overwrite channel #5");
627 channel_no = 5;
628 }
629 self.metrics
630 .channels_connected(&self.remote_pid_string, channel_no, cid);
631 }
632 },
633 )
634 .await;
635 trace!("Stop create_channel_mgr");
636 self.shutdown_barrier
637 .fetch_sub(Self::BARR_CHANNEL, Ordering::SeqCst);
638 }
639
640 async fn participant_shutdown_mgr(
661 &self,
662 s2b_shutdown_bparticipant_r: oneshot::Receiver<S2bShutdownBparticipant>,
663 b2b_close_send_protocol_s: async_channel::Sender<Cid>,
664 b2b_force_close_recv_protocol_s: async_channel::Sender<Cid>,
665 ) {
666 let wait_for_manager = || async {
667 let mut sleep = 0.01f64;
668 loop {
669 let bytes = self.shutdown_barrier.load(Ordering::SeqCst);
670 if bytes == 0 {
671 break;
672 }
673 sleep *= 1.4;
674 tokio::time::sleep(Duration::from_secs_f64(sleep)).await;
675 if sleep > 0.2 {
676 trace!(?bytes, "wait for mgr to close");
677 }
678 }
679 };
680
681 let awaited = s2b_shutdown_bparticipant_r.await.ok();
682 debug!("participant_shutdown_mgr triggered. Closing all streams for send");
683 {
684 let lock = self.streams.read().await;
685 for si in lock.values() {
686 si.send_closed.store(true, Ordering::SeqCst);
687 }
688 }
689
690 let lock = self.channels.read().await;
691 assert!(
692 !lock.is_empty(),
693 "no channel existed remote_pid={}",
694 self.remote_pid
695 );
696 for cid in lock.keys() {
697 if let Err(e) = b2b_close_send_protocol_s.send(*cid).await {
698 debug!(
699 ?e,
700 ?cid,
701 "closing send_mgr may fail if we got a recv error simultaneously"
702 );
703 }
704 }
705 drop(lock);
706
707 trace!("wait for other managers");
708 let timeout = tokio::time::sleep(
709 awaited
710 .as_ref()
711 .map(|(timeout_time, _)| *timeout_time)
712 .unwrap_or_default(),
713 );
714 let timeout = select! {
715 _ = wait_for_manager() => false,
716 _ = timeout => true,
717 };
718 if timeout {
719 warn!("timeout triggered: for killing recv");
720 let lock = self.channels.read().await;
721 for cid in lock.keys() {
722 if let Err(e) = b2b_force_close_recv_protocol_s.send(*cid).await {
723 debug!(
724 ?e,
725 ?cid,
726 "closing recv_mgr may fail if we got a recv error simultaneously"
727 );
728 }
729 }
730 }
731
732 trace!("wait again");
733 wait_for_manager().await;
734
735 if let Some((_, sender)) = awaited {
736 let _ = sender.send(Ok(()));
739 }
740
741 #[cfg(feature = "metrics")]
742 self.metrics.participants_disconnected_total.inc();
743 self.metrics.cleanup_participant(&self.remote_pid_string);
744 trace!("Stop participant_shutdown_mgr");
745 }
746
747 async fn delete_stream(&self, sid: Sid) {
750 let stream = { self.streams.write().await.remove(&sid) };
751 match stream {
752 Some(si) => {
753 si.send_closed.store(true, Ordering::SeqCst);
754 si.b2a_msg_recv_s.lock().await.close();
755 },
756 None => {
757 trace!("Couldn't find the stream, might be simultaneous close from local/remote")
758 },
759 }
760 self.metrics.streams_closed(&self.remote_pid_string);
761 }
762
763 async fn create_stream(
764 &self,
765 sid: Sid,
766 prio: Prio,
767 promises: Promises,
768 guaranteed_bandwidth: Bandwidth,
769 output_limit: usize,
770 ) -> Stream {
771 let (b2a_msg_recv_s, b2a_msg_recv_r) = async_channel::unbounded::<Bytes>();
772 let send_closed = Arc::new(AtomicBool::new(false));
773 self.streams.write().await.insert(sid, StreamInfo {
774 prio,
775 promises,
776 send_closed: Arc::clone(&send_closed),
777 b2a_msg_recv_s: Mutex::new(b2a_msg_recv_s),
778 });
779 self.metrics.streams_opened(&self.remote_pid_string);
780
781 let (a2b_msg_s, a2b_close_stream_s) = {
782 let lock = self.open_stream_channels.lock().await;
783 match &*lock {
784 Some(osi) => (osi.a2b_msg_s.clone(), osi.a2b_close_stream_s.clone()),
785 None => {
786 debug!(
788 "It seems that a stream was requested to open, while the send_mgr is \
789 already closed"
790 );
791 let (a2b_msg_s, _) = crossbeam_channel::unbounded();
792 let (a2b_close_stream_s, _) = mpsc::unbounded_channel();
793 (a2b_msg_s, a2b_close_stream_s)
794 },
795 }
796 };
797
798 Stream::new(
799 self.local_pid,
800 self.remote_pid,
801 sid,
802 prio,
803 promises,
804 guaranteed_bandwidth,
805 send_closed,
806 a2b_msg_s,
807 b2a_msg_recv_r,
808 a2b_close_stream_s,
809 output_limit,
810 )
811 }
812}
813
814#[cfg(test)]
815mod tests {
816 use super::*;
817 use core::assert_matches;
818 use network_protocol::{ProtocolMetricCache, ProtocolMetrics};
819 use tokio::{
820 runtime::Runtime,
821 sync::{mpsc, oneshot},
822 task::JoinHandle,
823 };
824
825 fn mock_bparticipant() -> (
826 Arc<Runtime>,
827 mpsc::UnboundedSender<A2bStreamOpen>,
828 mpsc::UnboundedReceiver<Stream>,
829 mpsc::UnboundedReceiver<ParticipantEvent>,
830 mpsc::UnboundedSender<S2bCreateChannel>,
831 oneshot::Sender<S2bShutdownBparticipant>,
832 mpsc::UnboundedReceiver<B2sPrioStatistic>,
833 watch::Receiver<f32>,
834 JoinHandle<()>,
835 ) {
836 let runtime = Arc::new(Runtime::new().unwrap());
837 let runtime_clone = Arc::clone(&runtime);
838
839 let (b2s_prio_statistic_s, b2s_prio_statistic_r) =
840 mpsc::unbounded_channel::<B2sPrioStatistic>();
841
842 let (
843 bparticipant,
844 a2b_open_stream_s,
845 b2a_stream_opened_r,
846 b2a_event_r,
847 s2b_create_channel_s,
848 s2b_shutdown_bparticipant_s,
849 b2a_bandwidth_stats_r,
850 ) = runtime_clone.block_on(async move {
851 let local_pid = Pid::fake(0);
852 let remote_pid = Pid::fake(1);
853 let sid = Sid::new(1000);
854 let metrics = Arc::new(NetworkMetrics::new(&local_pid).unwrap());
855
856 BParticipant::new(local_pid, remote_pid, sid, Arc::clone(&metrics), usize::MAX)
857 });
858
859 let handle = runtime_clone.spawn(bparticipant.run(b2s_prio_statistic_s));
860 (
861 runtime_clone,
862 a2b_open_stream_s,
863 b2a_stream_opened_r,
864 b2a_event_r,
865 s2b_create_channel_s,
866 s2b_shutdown_bparticipant_s,
867 b2s_prio_statistic_r,
868 b2a_bandwidth_stats_r,
869 handle,
870 )
871 }
872
873 #[expect(clippy::needless_pass_by_ref_mut)]
874 async fn mock_mpsc(
875 cid: Cid,
876 _runtime: &Arc<Runtime>,
877 create_channel: &mut mpsc::UnboundedSender<S2bCreateChannel>,
878 ) -> Protocols {
879 let (s1, r1) = mpsc::channel(100);
880 let (s2, r2) = mpsc::channel(100);
881 let met = Arc::new(ProtocolMetrics::new().unwrap());
882 let metrics = ProtocolMetricCache::new(&cid.to_string(), Arc::clone(&met));
883 let p1 = Protocols::new_mpsc(s1, r2, metrics);
884 let (complete_s, complete_r) = oneshot::channel();
885 create_channel
886 .send((cid, Sid::new(0), p1, ConnectAddr::Mpsc(42), complete_s))
887 .unwrap();
888 complete_r.await.unwrap();
889 let metrics = ProtocolMetricCache::new(&cid.to_string(), met);
890 Protocols::new_mpsc(s2, r1, metrics)
891 }
892
893 #[test]
894 fn close_bparticipant_by_timeout_during_close() {
895 let (
896 runtime,
897 a2b_open_stream_s,
898 b2a_stream_opened_r,
899 mut b2a_event_r,
900 mut s2b_create_channel_s,
901 s2b_shutdown_bparticipant_s,
902 b2s_prio_statistic_r,
903 _b2a_bandwidth_stats_r,
904 handle,
905 ) = mock_bparticipant();
906
907 let _remote = runtime.block_on(mock_mpsc(0, &runtime, &mut s2b_create_channel_s));
908 std::thread::sleep(Duration::from_millis(50));
909
910 let (s, r) = oneshot::channel();
911 let before = Instant::now();
912 runtime.block_on(async {
913 drop(s2b_create_channel_s);
914 s2b_shutdown_bparticipant_s
915 .send((Duration::from_secs(1), s))
916 .unwrap();
917 r.await.unwrap().unwrap();
918 });
919 assert!(
920 before.elapsed() > Duration::from_millis(900),
921 "timeout wasn't triggered"
922 );
923 assert_matches!(
924 b2a_event_r.try_recv().unwrap(),
925 ParticipantEvent::ChannelCreated(_)
926 );
927 assert_matches!(
928 b2a_event_r.try_recv().unwrap(),
929 ParticipantEvent::ChannelDeleted(_)
930 );
931 assert_matches!(b2a_event_r.try_recv(), Err(_));
932
933 runtime.block_on(handle).unwrap();
934
935 drop((a2b_open_stream_s, b2a_stream_opened_r, b2s_prio_statistic_r));
936 drop(runtime);
937 }
938
939 #[test]
940 fn close_bparticipant_cleanly() {
941 let (
942 runtime,
943 a2b_open_stream_s,
944 b2a_stream_opened_r,
945 mut b2a_event_r,
946 mut s2b_create_channel_s,
947 s2b_shutdown_bparticipant_s,
948 b2s_prio_statistic_r,
949 _b2a_bandwidth_stats_r,
950 handle,
951 ) = mock_bparticipant();
952
953 let remote = runtime.block_on(mock_mpsc(0, &runtime, &mut s2b_create_channel_s));
954 std::thread::sleep(Duration::from_millis(50));
955
956 let (s, r) = oneshot::channel();
957 let before = Instant::now();
958 runtime.block_on(async {
959 drop(s2b_create_channel_s);
960 s2b_shutdown_bparticipant_s
961 .send((Duration::from_secs(2), s))
962 .unwrap();
963 drop(remote); r.await.unwrap().unwrap();
965 });
966 assert!(
967 before.elapsed() < Duration::from_millis(1900),
968 "timeout was triggered"
969 );
970 assert_matches!(
971 b2a_event_r.try_recv().unwrap(),
972 ParticipantEvent::ChannelCreated(_)
973 );
974 assert_matches!(
975 b2a_event_r.try_recv().unwrap(),
976 ParticipantEvent::ChannelDeleted(_)
977 );
978 assert_matches!(b2a_event_r.try_recv(), Err(_));
979
980 runtime.block_on(handle).unwrap();
981
982 drop((a2b_open_stream_s, b2a_stream_opened_r, b2s_prio_statistic_r));
983 drop(runtime);
984 }
985
986 #[test]
987 fn create_stream() {
988 let (
989 runtime,
990 a2b_open_stream_s,
991 b2a_stream_opened_r,
992 _b2a_event_r,
993 mut s2b_create_channel_s,
994 s2b_shutdown_bparticipant_s,
995 b2s_prio_statistic_r,
996 _b2a_bandwidth_stats_r,
997 handle,
998 ) = mock_bparticipant();
999
1000 let remote = runtime.block_on(mock_mpsc(0, &runtime, &mut s2b_create_channel_s));
1001 std::thread::sleep(Duration::from_millis(50));
1002
1003 let (rs, mut rr) = remote.split();
1005 let (stream_sender, _stream_receiver) = oneshot::channel();
1006 a2b_open_stream_s
1007 .send((7u8, Promises::ENCRYPTED, 1_000_000, stream_sender))
1008 .unwrap();
1009
1010 let stream_event = runtime.block_on(rr.recv()).unwrap();
1011 match stream_event {
1012 ProtocolEvent::OpenStream {
1013 sid,
1014 prio,
1015 promises,
1016 guaranteed_bandwidth,
1017 } => {
1018 assert_eq!(sid, Sid::new(1000));
1019 assert_eq!(prio, 7u8);
1020 assert_eq!(promises, Promises::ENCRYPTED);
1021 assert_eq!(guaranteed_bandwidth, 1_000_000);
1022 },
1023 _ => panic!("wrong event"),
1024 };
1025
1026 let (s, r) = oneshot::channel();
1027 runtime.block_on(async {
1028 drop(s2b_create_channel_s);
1029 s2b_shutdown_bparticipant_s
1030 .send((Duration::from_secs(1), s))
1031 .unwrap();
1032 drop((rs, rr));
1033 r.await.unwrap().unwrap();
1034 });
1035
1036 runtime.block_on(handle).unwrap();
1037
1038 drop((a2b_open_stream_s, b2a_stream_opened_r, b2s_prio_statistic_r));
1039 drop(runtime);
1040 }
1041
1042 #[test]
1043 fn created_stream() {
1044 let (
1045 runtime,
1046 a2b_open_stream_s,
1047 mut b2a_stream_opened_r,
1048 _b2a_event_r,
1049 mut s2b_create_channel_s,
1050 s2b_shutdown_bparticipant_s,
1051 b2s_prio_statistic_r,
1052 _b2a_bandwidth_stats_r,
1053 handle,
1054 ) = mock_bparticipant();
1055
1056 let remote = runtime.block_on(mock_mpsc(0, &runtime, &mut s2b_create_channel_s));
1057 std::thread::sleep(Duration::from_millis(50));
1058
1059 let (mut rs, rr) = remote.split();
1061 runtime
1062 .block_on(rs.send(ProtocolEvent::OpenStream {
1063 sid: Sid::new(1000),
1064 prio: 9u8,
1065 promises: Promises::ORDERED,
1066 guaranteed_bandwidth: 1_000_000,
1067 }))
1068 .unwrap();
1069
1070 let stream = runtime.block_on(b2a_stream_opened_r.recv()).unwrap();
1071 assert_eq!(stream.params().promises, Promises::ORDERED);
1072
1073 let (s, r) = oneshot::channel();
1074 runtime.block_on(async {
1075 drop(s2b_create_channel_s);
1076 s2b_shutdown_bparticipant_s
1077 .send((Duration::from_secs(1), s))
1078 .unwrap();
1079 drop((rs, rr));
1080 r.await.unwrap().unwrap();
1081 });
1082
1083 runtime.block_on(handle).unwrap();
1084
1085 drop((a2b_open_stream_s, b2a_stream_opened_r, b2s_prio_statistic_r));
1086 drop(runtime);
1087 }
1088}