1mod probe;
4#[cfg(test)]
5pub(crate) mod test_support;
6#[cfg(test)]
7mod tests;
8mod transport;
9
10use crate::types::SessionId;
11use serde::{Deserialize, Serialize};
12use std::{
13 collections::{BTreeMap, BTreeSet},
14 path::Path,
15 sync::{Arc, Mutex, OnceLock, Weak},
16 time::Duration,
17};
18use tokio::{
19 io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt},
20 sync::{Notify, broadcast, oneshot},
21 task::JoinSet,
22};
23use tokio_util::sync::CancellationToken;
24use transport::Endpoint;
25
26pub const RECONCILE: Duration = Duration::from_secs(5);
28const RECONNECT: Duration = Duration::from_millis(500);
30const DELIVERY_TIMEOUT: Duration = Duration::from_millis(300);
32const MAX_FRAME_BYTES: usize = 1024;
34const MAX_STATES: usize = 32;
36const MAX_CLIENTS: usize = 64;
37const MAX_BACKLOG: usize = 64;
38const MAX_DIRTY_SESSIONS: usize = 64;
39const VERSION: u32 = 1;
43
44#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
46pub enum Change {
47 StateInvalidated,
48 SessionInvalidated(SessionId),
49 Resync,
51}
52
53enum Dirty {
55 Pending {
56 state: bool,
57 sessions: BTreeSet<SessionId>,
58 },
59 Resync,
60}
61impl Default for Dirty {
62 fn default() -> Self {
63 Self::Pending {
64 state: false,
65 sessions: BTreeSet::new(),
66 }
67 }
68}
69impl Dirty {
70 fn add(&mut self, change: Change) {
71 let Self::Pending { state, sessions } = self else {
72 return;
73 };
74 match change {
75 Change::Resync => *self = Self::Resync,
76 Change::StateInvalidated => *state = true,
77 Change::SessionInvalidated(id) => {
78 sessions.insert(id);
79 if sessions.len() > MAX_DIRTY_SESSIONS {
80 *self = Self::Resync;
81 }
82 }
83 }
84 }
85 fn pop(&mut self) -> Option<Change> {
86 match self {
87 Self::Resync => {
88 *self = Self::default();
89 Some(Change::Resync)
90 }
91 Self::Pending { state, sessions } => {
92 if std::mem::take(state) {
93 Some(Change::StateInvalidated)
94 } else {
95 sessions.pop_first().map(Change::SessionInvalidated)
96 }
97 }
98 }
99 }
100 fn is_empty(&self) -> bool {
101 matches!(self, Self::Pending { state: false, sessions } if sessions.is_empty())
102 }
103}
104struct Inbox {
105 #[cfg(test)]
106 registration: tokio::sync::watch::Sender<Option<crate::config::HubEpoch>>,
107 #[cfg(test)]
109 received: tokio::sync::watch::Sender<Option<u64>>,
110 dirty: Mutex<Dirty>,
111 wake: Notify,
112}
113impl Default for Inbox {
114 fn default() -> Self {
115 Self {
116 #[cfg(test)]
117 registration: tokio::sync::watch::channel(None).0,
118 #[cfg(test)]
119 received: tokio::sync::watch::channel(None).0,
120 dirty: Mutex::default(),
121 wake: Notify::new(),
122 }
123 }
124}
125impl Inbox {
126 fn add(&self, change: Change) {
127 self.dirty.lock().unwrap().add(change);
128 self.wake.notify_one();
129 }
130 async fn next(&self) -> Change {
131 loop {
132 let wake = self.wake.notified();
134 if let Some(change) = self.dirty.lock().unwrap().pop() {
135 return change;
136 }
137 wake.await;
138 }
139 }
140}
141
142pub struct Subscription {
145 inbox: Arc<Inbox>,
146 cancel: CancellationToken,
147 task: tokio::task::JoinHandle<Result<(), String>>,
148}
149impl Subscription {
150 pub async fn new(state: &Path) -> Self {
151 let inbox = Arc::new(Inbox::default());
152 let cancel = CancellationToken::new();
153 let (ready, registered) = oneshot::channel();
154 let state = state.to_path_buf();
155 let task = tokio::spawn(watch(state, inbox.clone(), cancel.clone(), ready));
156 let subscription = Self {
158 inbox,
159 cancel,
160 task,
161 };
162 let _ = registered.await;
164 subscription
165 }
166 pub async fn close(mut self) -> Result<(), String> {
169 self.cancel.cancel();
170 crate::task::joined((&mut self.task).await).map_err(|error| error.to_string())?
171 }
172 pub async fn next_reconciled(&mut self) -> Change {
175 tokio::time::timeout(RECONCILE, self.next_state())
176 .await
177 .unwrap_or(Change::Resync)
178 }
179 pub async fn next(&mut self) -> Change {
180 self.inbox.next().await
181 }
182 pub async fn next_state(&mut self) -> Change {
184 loop {
185 let change = self.next().await;
186 if !matches!(change, Change::SessionInvalidated(_)) {
187 return change;
188 }
189 }
190 }
191}
192impl Drop for Subscription {
193 fn drop(&mut self) {
194 self.cancel.cancel();
195 }
196}
197
198#[derive(Default)]
199struct Outbox {
200 dirty: Dirty,
201 submitted: u64,
202 delivered: u64,
203}
204struct Publishing {
205 endpoint: Endpoint,
206 pending: Mutex<Outbox>,
207 wake: Notify,
208 drained: Notify,
209 #[cfg(test)]
211 delivery_error: Mutex<Option<String>>,
212}
213#[derive(Clone, Default)]
215pub(crate) struct Publisher(Option<Arc<Publishing>>);
216static PUBLISHERS: OnceLock<Mutex<BTreeMap<crate::config::EndpointId, Weak<Publishing>>>> =
217 OnceLock::new();
218impl Publisher {
219 pub(crate) fn new(state: &Path) -> Self {
220 let Ok(runtime) = tokio::runtime::Handle::try_current() else {
221 return Self::default();
222 };
223 let Ok(endpoint) = Endpoint::new(state) else {
224 return Self::default();
225 };
226 let mut publishers = PUBLISHERS.get_or_init(Default::default).lock().unwrap();
227 publishers.retain(|_, publisher| publisher.strong_count() != 0);
228 if let Some(publisher) = publishers.get(&endpoint.identity).and_then(Weak::upgrade) {
229 return Self(Some(publisher));
230 }
231 if publishers.len() >= MAX_STATES {
232 return Self::default();
233 }
234 let publisher = Arc::new(Publishing {
235 endpoint,
236 pending: Mutex::new(Outbox::default()),
237 wake: Notify::new(),
238 drained: Notify::new(),
239 #[cfg(test)]
240 delivery_error: Mutex::default(),
241 });
242 publishers.insert(
243 publisher.endpoint.identity.clone(),
244 Arc::downgrade(&publisher),
245 );
246 runtime.spawn(publish(publisher.clone()));
247 Self(Some(publisher))
248 }
249 pub(crate) fn notify(&self, change: Change) {
250 if let Some(publisher) = &self.0 {
251 let mut pending = publisher.pending.lock().unwrap();
252 pending.dirty.add(change);
253 pending.submitted = pending.submitted.wrapping_add(1);
254 drop(pending);
255 publisher.wake.notify_one();
256 }
257 }
258}
259
260pub async fn drain() {
262 let publishers: Vec<_> = PUBLISHERS
263 .get()
264 .map(|publishers| {
265 publishers
266 .lock()
267 .unwrap()
268 .values()
269 .filter_map(Weak::upgrade)
270 .collect()
271 })
272 .unwrap_or_default();
273 let wait = async {
274 for publisher in publishers {
275 let target = publisher.pending.lock().unwrap().submitted;
276 loop {
277 let drained = publisher.drained.notified();
278 if publisher.pending.lock().unwrap().delivered >= target {
279 break;
280 }
281 drained.await;
282 }
283 }
284 };
285 let _ = tokio::time::timeout(DELIVERY_TIMEOUT, wait).await;
286}
287
288#[derive(Serialize, Deserialize)]
289enum Frame {
290 Hello {
291 version: u32,
292 identity: crate::config::EndpointId,
293 subscriber: bool,
294 },
295 Registered {
296 epoch: crate::config::HubEpoch,
297 },
298 Publish(Change),
299 Delivered,
300 Hint {
301 epoch: crate::config::HubEpoch,
302 sequence: u64,
303 change: Change,
304 },
305}
306async fn read_frame<R: AsyncRead + Unpin>(reader: &mut R) -> std::io::Result<Frame> {
307 let bytes = reader.read_u32().await? as usize;
308 if bytes == 0 || bytes > MAX_FRAME_BYTES {
309 return Err(std::io::ErrorKind::InvalidData.into());
310 }
311 let mut data = vec![0; bytes];
312 reader.read_exact(&mut data).await?;
313 serde_json::from_slice(&data).map_err(std::io::Error::other)
314}
315async fn write_frame<W: AsyncWrite + Unpin>(writer: &mut W, frame: &Frame) -> std::io::Result<()> {
316 let data = serde_json::to_vec(frame).map_err(std::io::Error::other)?;
317 if data.len() > MAX_FRAME_BYTES {
318 return Err(std::io::ErrorKind::InvalidData.into());
319 }
320 writer.write_u32(data.len() as u32).await?;
321 writer.write_all(&data).await?;
322 writer.flush().await
323}
324async fn connect(
325 endpoint: &Endpoint,
326 subscriber: bool,
327) -> std::io::Result<(transport::Stream, crate::config::HubEpoch)> {
328 let mut stream = endpoint.connect().await?;
329 write_frame(
330 &mut stream,
331 &Frame::Hello {
332 version: VERSION,
333 identity: endpoint.identity.clone(),
334 subscriber,
335 },
336 )
337 .await?;
338 match read_frame(&mut stream).await? {
339 Frame::Registered { epoch } if epoch.len() <= crate::types::MAX_ID_BYTES => {
340 Ok((stream, epoch))
341 }
342 _ => Err(std::io::ErrorKind::InvalidData.into()),
343 }
344}
345async fn publish(publisher: Arc<Publishing>) {
346 loop {
347 let wake = publisher.wake.notified();
348 if publisher.pending.lock().unwrap().dirty.is_empty() {
349 tokio::select! {
350 _ = wake => {},
351 _ = tokio::time::sleep(RECONNECT) => {
352 if Arc::strong_count(&publisher) == 1 {
354 break;
355 }
356 continue;
357 }
358 }
359 }
360 let (mut dirty, target) = {
361 let mut pending = publisher.pending.lock().unwrap();
362 (std::mem::take(&mut pending.dirty), pending.submitted)
363 };
364 let delivery = async {
366 let (mut stream, _) = connect(&publisher.endpoint, false).await?;
367 while let Some(change) = dirty.pop() {
368 write_frame(&mut stream, &Frame::Publish(change)).await?;
369 if !matches!(read_frame(&mut stream).await?, Frame::Delivered) {
370 return Err(std::io::ErrorKind::InvalidData.into());
371 }
372 }
373 Ok::<_, std::io::Error>(())
374 };
375 let _delivery = tokio::time::timeout(DELIVERY_TIMEOUT, delivery).await;
376 #[cfg(test)]
377 if !matches!(&_delivery, Ok(Ok(()))) {
378 *publisher.delivery_error.lock().unwrap() = Some(format!("{_delivery:?}"));
379 }
380 publisher.pending.lock().unwrap().delivered = target;
381 publisher.drained.notify_one();
382 }
383}
384
385async fn watch(
386 state: std::path::PathBuf,
387 inbox: Arc<Inbox>,
388 cancel: CancellationToken,
389 ready: oneshot::Sender<()>,
390) -> Result<(), String> {
391 let mut ready = Some(ready);
392 let endpoint = Endpoint::new(&state).ok();
393 let mut hub = None;
394 let mut probe = probe::Probe::new(state);
395 let mut reconcile =
398 tokio::time::interval_at(tokio::time::Instant::now() + RECONCILE, RECONCILE);
399 reconcile.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
400 let _ = probe.changed().await;
402 loop {
403 if cancel.is_cancelled() {
404 break;
405 }
406 if let Some(endpoint) = &endpoint {
407 if hub
408 .as_ref()
409 .is_some_and(|task: &tokio::task::JoinHandle<()>| task.is_finished())
410 {
411 hub = None;
412 }
413 if hub.is_none()
414 && let Ok(Some(owner)) = endpoint.elect()
415 && let Ok(listener) = endpoint.listen()
416 {
417 let endpoint = endpoint.clone();
418 let token = cancel.child_token();
419 hub = Some(tokio::spawn(async move {
420 let _ = serve(endpoint, listener, owner, token).await;
421 }));
422 }
423 let connection = tokio::time::timeout(DELIVERY_TIMEOUT, connect(endpoint, true));
424 tokio::select! {
425 _ = cancel.cancelled() => break,
426 result = connection => {
427 if let Ok(Ok((stream, epoch))) = result {
428 inbox.add(Change::Resync);
429 #[cfg(test)]
430 inbox.registration.send_replace(Some(epoch.clone()));
431 if let Some(ready) = ready.take() {
432 let _ = ready.send(());
433 }
434 let receive = receive(stream, epoch, inbox.clone());
435 tokio::pin!(receive);
436 loop {
437 tokio::select! {
438 _ = cancel.cancelled() => break,
439 _ = &mut receive => {
440 #[cfg(test)]
441 inbox.registration.send_replace(None);
442 inbox.add(Change::Resync);
443 break;
444 }
445 _ = reconcile.tick() => {
446 if probe.changed().await {
447 inbox.add(Change::Resync);
448 }
449 }
450 }
451 }
452 }
453 }
454 }
455 }
456 if let Some(ready) = ready.take() {
457 inbox.add(Change::Resync);
458 let _ = ready.send(());
459 }
460 tokio::select! {
461 _ = cancel.cancelled() => break,
462 _ = reconcile.tick() => {
463 let _ = probe.changed().await;
464 inbox.add(Change::Resync);
467 },
468 _ = tokio::time::sleep(RECONNECT) => {},
469 }
470 }
471 if let Some(task) = hub {
472 let _ = crate::task::joined(task.await);
474 }
475 probe.close().await
476}
477async fn receive(
478 mut stream: transport::Stream,
479 epoch: crate::config::HubEpoch,
480 inbox: Arc<Inbox>,
481) -> std::io::Result<()> {
482 let mut previous = None;
483 loop {
484 match read_frame(&mut stream).await? {
485 Frame::Hint {
486 epoch: incoming,
487 sequence,
488 change,
489 } if incoming == epoch => {
490 if previous.is_some_and(|previous: u64| sequence != previous.wrapping_add(1)) {
491 inbox.add(Change::Resync);
492 }
493 previous = Some(sequence);
494 inbox.add(change);
495 #[cfg(test)]
496 inbox.received.send_replace(Some(sequence));
497 }
498 _ => return Err(std::io::ErrorKind::InvalidData.into()),
499 }
500 }
501}
502
503#[derive(Clone)]
504struct Hint {
505 sequence: u64,
506 change: Change,
507}
508async fn serve(
509 endpoint: Endpoint,
510 mut listener: transport::Listener,
511 _owner: std::fs::File,
512 cancel: CancellationToken,
513) -> std::io::Result<()> {
514 let epoch: crate::config::HubEpoch = crate::agent::uuid()
515 .map_err(std::io::Error::other)?
516 .parse()
517 .expect("UUID hub epoch");
518 let (sender, _) = broadcast::channel(MAX_BACKLOG);
519 let sequence = Arc::new(Mutex::new(0_u64));
520 let mut clients = JoinSet::new();
521 loop {
522 tokio::select! {
523 _ = cancel.cancelled() => break,
524 Some(result) = clients.join_next(), if !clients.is_empty() => {
525 let _ = crate::task::joined(result);
526 },
527 result = listener.accept(), if clients.len() < MAX_CLIENTS => {
528 let stream = result?;
529 let (identity, epoch, sender, sequence) = (
530 endpoint.identity.clone(),
531 epoch.clone(),
532 sender.clone(),
533 sequence.clone(),
534 );
535 clients.spawn(async move {
536 let _ = client(stream, identity, epoch, sender, sequence).await;
537 });
538 }
539 }
540 }
541 clients.abort_all();
542 while let Some(result) = clients.join_next().await {
543 let _ = crate::task::joined(result);
544 }
545 Ok(())
546}
547async fn client(
548 mut stream: transport::Stream,
549 identity: crate::config::EndpointId,
550 epoch: crate::config::HubEpoch,
551 sender: broadcast::Sender<Hint>,
552 sequence: Arc<Mutex<u64>>,
553) -> std::io::Result<()> {
554 transport::validate_peer(&stream)?;
557 let hello = tokio::time::timeout(DELIVERY_TIMEOUT, read_frame(&mut stream)).await??;
558 let Frame::Hello {
559 version: VERSION,
560 identity: incoming,
561 subscriber,
562 } = hello
563 else {
564 return Err(std::io::ErrorKind::InvalidData.into());
565 };
566 if incoming != identity {
567 return Err(std::io::ErrorKind::PermissionDenied.into());
568 }
569 let mut receiver = sender.subscribe();
571 tokio::time::timeout(
572 DELIVERY_TIMEOUT,
573 write_frame(
574 &mut stream,
575 &Frame::Registered {
576 epoch: epoch.clone(),
577 },
578 ),
579 )
580 .await??;
581 if subscriber {
582 let (mut reader, mut writer) = tokio::io::split(stream);
583 let outgoing = async {
584 loop {
585 let hint = match receiver.recv().await {
586 Ok(hint) => hint,
587 Err(broadcast::error::RecvError::Lagged(_)) => Hint {
588 sequence: *sequence.lock().unwrap(),
589 change: Change::Resync,
590 },
591 Err(_) => return Ok::<_, std::io::Error>(()),
592 };
593 tokio::time::timeout(
595 DELIVERY_TIMEOUT,
596 write_frame(
597 &mut writer,
598 &Frame::Hint {
599 epoch: epoch.clone(),
600 sequence: hint.sequence,
601 change: hint.change,
602 },
603 ),
604 )
605 .await??;
606 }
607 };
608 tokio::select! { result = outgoing => result, _ = reader.read_u8() => Ok(()) }
609 } else {
610 loop {
611 let frame = tokio::time::timeout(DELIVERY_TIMEOUT, read_frame(&mut stream)).await??;
612 let Frame::Publish(change) = frame else {
613 return Err(std::io::ErrorKind::InvalidData.into());
614 };
615 {
616 let mut sequence = sequence.lock().unwrap();
618 *sequence = sequence.wrapping_add(1);
619 let _ = sender.send(Hint {
620 sequence: *sequence,
621 change,
622 });
623 }
624 tokio::time::timeout(
625 DELIVERY_TIMEOUT,
626 write_frame(&mut stream, &Frame::Delivered),
627 )
628 .await??;
629 }
630 }
631}