Skip to main content

moq_net/
session.rs

1use std::{sync::Arc, time::Duration};
2
3use web_async::MaybeSendBoxFuture;
4use web_transport_trait::Stats;
5
6use crate::{BandwidthConsumer, BandwidthProducer, Error, Version};
7
8/// A MoQ transport session, wrapping a WebTransport connection.
9///
10/// Created via:
11/// - [`crate::Client::connect`] for clients.
12/// - [`crate::Server::accept`] for servers.
13#[derive(Clone)]
14pub struct Session {
15	session: Arc<dyn SessionInner>,
16	version: Version,
17	send_bandwidth: Option<BandwidthConsumer>,
18	recv_bandwidth: Option<BandwidthConsumer>,
19	closed: bool,
20}
21
22impl Session {
23	pub(super) fn new<S: web_transport_trait::Session>(
24		session: S,
25		version: Version,
26		recv_bandwidth: Option<BandwidthConsumer>,
27	) -> Self {
28		// Send bandwidth is version-agnostic: it depends on QUIC backend support.
29		let send_bandwidth = if session.stats().estimated_send_rate().is_some() {
30			let producer = BandwidthProducer::new();
31			let consumer = producer.consume();
32
33			let session = session.clone();
34			web_async::spawn(async move {
35				run_send_bandwidth(&session, producer).await;
36			});
37
38			Some(consumer)
39		} else {
40			None
41		};
42
43		Self {
44			session: Arc::new(session),
45			version,
46			send_bandwidth,
47			recv_bandwidth,
48			closed: false,
49		}
50	}
51
52	/// Returns the negotiated protocol version.
53	pub fn version(&self) -> Version {
54		self.version
55	}
56
57	/// Returns a consumer for the estimated send bitrate (from the congestion controller).
58	///
59	/// Returns `None` if the QUIC backend doesn't support bandwidth estimation.
60	pub fn send_bandwidth(&self) -> Option<BandwidthConsumer> {
61		self.send_bandwidth.clone()
62	}
63
64	/// Returns a consumer for the estimated receive bitrate (from PROBE).
65	///
66	/// Returns `None` if the MoQ version doesn't support PROBE (requires moq-lite-03+).
67	pub fn recv_bandwidth(&self) -> Option<BandwidthConsumer> {
68		self.recv_bandwidth.clone()
69	}
70
71	/// Close the underlying transport session.
72	pub fn close(&mut self, err: Error) {
73		if self.closed {
74			return;
75		}
76		self.closed = true;
77		self.session.close(err.to_code(), err.to_string().as_ref());
78	}
79
80	/// Block until the transport session is closed.
81	pub async fn closed(&self) -> Result<(), Error> {
82		let err = self.session.closed().await;
83		Err(Error::Transport(err))
84	}
85}
86
87impl Drop for Session {
88	fn drop(&mut self) {
89		if !self.closed {
90			self.session.close(Error::Cancel.to_code(), "dropped");
91		}
92	}
93}
94
95/// Polls the QUIC congestion controller for estimated send rate.
96///
97/// Exits as soon as the session closes so we don't pin the underlying connection
98/// after the wrapping [`Session`] is dropped.
99async fn run_send_bandwidth<S: web_transport_trait::Session>(session: &S, producer: BandwidthProducer) {
100	tokio::select! {
101		_ = session.closed() => {}
102		_ = producer.closed() => {}
103		_ = run_send_bandwidth_inner(session, &producer) => {}
104	}
105}
106
107/// Toggles between waiting for a consumer and polling stats while one exists.
108/// Returns when the producer channel errors (closed by the consumer side).
109async fn run_send_bandwidth_inner<S: web_transport_trait::Session>(session: &S, producer: &BandwidthProducer) {
110	const POLL_INTERVAL: Duration = Duration::from_millis(100);
111
112	loop {
113		if producer.used().await.is_err() {
114			return;
115		}
116
117		let mut interval = web_async::time::interval(POLL_INTERVAL);
118		loop {
119			tokio::select! {
120				biased;
121				res = producer.unused() => {
122					if res.is_err() {
123						return;
124					}
125					// No more consumers, pause polling.
126					break;
127				}
128				_ = interval.tick() => {
129					let bitrate = session.stats().estimated_send_rate();
130					if producer.set(bitrate).is_err() {
131						return;
132					}
133				}
134			}
135		}
136	}
137}
138
139// We use a wrapper type that is dyn-compatible to remove the generic bounds from Session.
140trait SessionInner: web_transport_trait::MaybeSend + web_transport_trait::MaybeSync {
141	fn close(&self, code: u32, reason: &str);
142	fn closed(&self) -> MaybeSendBoxFuture<'_, String>;
143}
144
145impl<S: web_transport_trait::Session> SessionInner for S {
146	fn close(&self, code: u32, reason: &str) {
147		S::close(self, code, reason);
148	}
149
150	fn closed(&self) -> MaybeSendBoxFuture<'_, String> {
151		Box::pin(async move { S::closed(self).await.to_string() })
152	}
153}