1mod probe;
4#[cfg(test)]
5pub(crate) mod test_support;
6#[cfg(test)]
7mod tests;
8mod transport;
9
10use crate::types::{RunId, 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 MAX_DIRTY_RUNS: usize = 64;
40const VERSION: u32 = 2;
45
46#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
48pub enum Change {
49 StateInvalidated,
52 RunInvalidated(RunId),
55 SessionInvalidated(SessionId),
56 Resync,
58}
59
60#[derive(Debug, Clone, Default, PartialEq, Eq)]
63pub enum Changed {
64 #[default]
65 Nothing,
66 Runs(BTreeSet<RunId>),
67 Everything,
68}
69impl Changed {
70 pub fn add(&mut self, change: &Change) {
71 match change {
72 Change::StateInvalidated | Change::Resync => *self = Self::Everything,
73 Change::RunInvalidated(id) => match self {
74 Self::Nothing => *self = Self::Runs(BTreeSet::from([id.clone()])),
75 Self::Runs(ids) => {
76 ids.insert(id.clone());
77 if ids.len() > MAX_DIRTY_RUNS {
78 *self = Self::Everything;
79 }
80 }
81 Self::Everything => {}
82 },
83 Change::SessionInvalidated(_) => {}
84 }
85 }
86 pub fn merge(&mut self, other: Self) {
88 match other {
89 Self::Nothing => {}
90 Self::Everything => *self = Self::Everything,
91 Self::Runs(ids) => {
92 for id in ids {
93 self.add(&Change::RunInvalidated(id));
94 }
95 }
96 }
97 }
98 pub fn or_written(&self, written: bool) -> Self {
101 match self {
102 Self::Nothing if written => Self::Everything,
103 changed => changed.clone(),
104 }
105 }
106 pub fn run(&self, id: &RunId) -> bool {
108 match self {
109 Self::Nothing => false,
110 Self::Runs(ids) => ids.contains(id),
111 Self::Everything => true,
112 }
113 }
114}
115
116enum Dirty {
119 Pending {
120 state: bool,
121 runs: BTreeSet<RunId>,
122 sessions: BTreeSet<SessionId>,
123 },
124 Resync,
125}
126impl Default for Dirty {
127 fn default() -> Self {
128 Self::Pending {
129 state: false,
130 runs: BTreeSet::new(),
131 sessions: BTreeSet::new(),
132 }
133 }
134}
135impl Dirty {
136 fn add(&mut self, change: Change) {
137 let Self::Pending {
138 state,
139 runs,
140 sessions,
141 } = self
142 else {
143 return;
144 };
145 match change {
146 Change::Resync => *self = Self::Resync,
147 Change::StateInvalidated => {
148 *state = true;
149 runs.clear();
150 }
151 Change::RunInvalidated(id) => {
152 if !*state {
153 runs.insert(id);
154 if runs.len() > MAX_DIRTY_RUNS {
155 *state = true;
156 runs.clear();
157 }
158 }
159 }
160 Change::SessionInvalidated(id) => {
161 sessions.insert(id);
162 if sessions.len() > MAX_DIRTY_SESSIONS {
163 *self = Self::Resync;
164 }
165 }
166 }
167 }
168 fn pop(&mut self) -> Option<Change> {
169 match self {
170 Self::Resync => {
171 *self = Self::default();
172 Some(Change::Resync)
173 }
174 Self::Pending {
175 state,
176 runs,
177 sessions,
178 } => {
179 if std::mem::take(state) {
180 Some(Change::StateInvalidated)
181 } else if let Some(id) = runs.pop_first() {
182 Some(Change::RunInvalidated(id))
183 } else {
184 sessions.pop_first().map(Change::SessionInvalidated)
185 }
186 }
187 }
188 }
189 fn is_empty(&self) -> bool {
190 matches!(self, Self::Pending { state: false, runs, sessions } if runs.is_empty() && sessions.is_empty())
191 }
192}
193struct Inbox {
194 #[cfg(test)]
195 registration: tokio::sync::watch::Sender<Option<crate::config::HubEpoch>>,
196 #[cfg(test)]
198 received: tokio::sync::watch::Sender<Option<u64>>,
199 dirty: Mutex<Dirty>,
200 wake: Notify,
201}
202impl Default for Inbox {
203 fn default() -> Self {
204 Self {
205 #[cfg(test)]
206 registration: tokio::sync::watch::channel(None).0,
207 #[cfg(test)]
208 received: tokio::sync::watch::channel(None).0,
209 dirty: Mutex::default(),
210 wake: Notify::new(),
211 }
212 }
213}
214impl Inbox {
215 fn add(&self, change: Change) {
216 self.dirty.lock().unwrap().add(change);
217 self.wake.notify_one();
218 }
219 async fn next(&self) -> Change {
220 loop {
221 let wake = self.wake.notified();
223 if let Some(change) = self.dirty.lock().unwrap().pop() {
224 return change;
225 }
226 wake.await;
227 }
228 }
229}
230
231pub struct Subscription {
234 inbox: Arc<Inbox>,
235 cancel: CancellationToken,
236 task: tokio::task::JoinHandle<Result<(), String>>,
237}
238impl Subscription {
239 pub async fn new(state: &Path) -> Self {
240 let inbox = Arc::new(Inbox::default());
241 let cancel = CancellationToken::new();
242 let (ready, registered) = oneshot::channel();
243 let state = state.to_path_buf();
244 let task = tokio::spawn(watch(state, inbox.clone(), cancel.clone(), ready));
245 let subscription = Self {
247 inbox,
248 cancel,
249 task,
250 };
251 let _ = registered.await;
253 subscription
254 }
255 pub async fn close(mut self) -> Result<(), String> {
258 self.cancel.cancel();
259 crate::task::joined((&mut self.task).await).map_err(|error| error.to_string())?
260 }
261 pub async fn next_reconciled(&mut self) -> Change {
264 tokio::time::timeout(RECONCILE, self.next_state())
265 .await
266 .unwrap_or(Change::Resync)
267 }
268 pub async fn next(&mut self) -> Change {
269 self.inbox.next().await
270 }
271 pub async fn next_state(&mut self) -> Change {
274 loop {
275 let change = self.next().await;
276 if !matches!(change, Change::SessionInvalidated(_)) {
277 return change;
278 }
279 }
280 }
281}
282impl Drop for Subscription {
283 fn drop(&mut self) {
284 self.cancel.cancel();
285 }
286}
287
288#[derive(Default)]
289struct Outbox {
290 dirty: Dirty,
291 submitted: u64,
292 delivered: u64,
293}
294struct Publishing {
295 endpoint: Endpoint,
296 pending: Mutex<Outbox>,
297 wake: Notify,
298 drained: Notify,
299 #[cfg(test)]
301 delivery_error: Mutex<Option<String>>,
302}
303#[derive(Clone, Default)]
305pub(crate) struct Publisher(Option<Arc<Publishing>>);
306static PUBLISHERS: OnceLock<Mutex<BTreeMap<crate::config::EndpointId, Weak<Publishing>>>> =
307 OnceLock::new();
308impl Publisher {
309 pub(crate) fn new(state: &Path) -> Self {
310 let Ok(runtime) = tokio::runtime::Handle::try_current() else {
311 return Self::default();
312 };
313 let Ok(endpoint) = Endpoint::new(state) else {
314 return Self::default();
315 };
316 let mut publishers = PUBLISHERS.get_or_init(Default::default).lock().unwrap();
317 publishers.retain(|_, publisher| publisher.strong_count() != 0);
318 if let Some(publisher) = publishers.get(&endpoint.identity).and_then(Weak::upgrade) {
319 return Self(Some(publisher));
320 }
321 if publishers.len() >= MAX_STATES {
322 return Self::default();
323 }
324 let publisher = Arc::new(Publishing {
325 endpoint,
326 pending: Mutex::new(Outbox::default()),
327 wake: Notify::new(),
328 drained: Notify::new(),
329 #[cfg(test)]
330 delivery_error: Mutex::default(),
331 });
332 publishers.insert(
333 publisher.endpoint.identity.clone(),
334 Arc::downgrade(&publisher),
335 );
336 runtime.spawn(publish(publisher.clone()));
337 Self(Some(publisher))
338 }
339 pub(crate) fn notify(&self, change: Change) {
340 if let Some(publisher) = &self.0 {
341 let mut pending = publisher.pending.lock().unwrap();
342 pending.dirty.add(change);
343 pending.submitted = pending.submitted.wrapping_add(1);
344 drop(pending);
345 publisher.wake.notify_one();
346 }
347 }
348}
349
350pub async fn drain() {
352 let publishers: Vec<_> = PUBLISHERS
353 .get()
354 .map(|publishers| {
355 publishers
356 .lock()
357 .unwrap()
358 .values()
359 .filter_map(Weak::upgrade)
360 .collect()
361 })
362 .unwrap_or_default();
363 let wait = async {
364 for publisher in publishers {
365 let target = publisher.pending.lock().unwrap().submitted;
366 loop {
367 let drained = publisher.drained.notified();
368 if publisher.pending.lock().unwrap().delivered >= target {
369 break;
370 }
371 drained.await;
372 }
373 }
374 };
375 let _ = tokio::time::timeout(DELIVERY_TIMEOUT, wait).await;
376}
377
378#[derive(Serialize, Deserialize)]
379enum Frame {
380 Hello {
381 version: u32,
382 identity: crate::config::EndpointId,
383 subscriber: bool,
384 },
385 Registered {
386 epoch: crate::config::HubEpoch,
387 },
388 Publish(Change),
389 Delivered,
390 Hint {
391 epoch: crate::config::HubEpoch,
392 sequence: u64,
393 change: Change,
394 },
395}
396async fn read_frame<R: AsyncRead + Unpin>(reader: &mut R) -> std::io::Result<Frame> {
397 let bytes = reader.read_u32().await? as usize;
398 if bytes == 0 || bytes > MAX_FRAME_BYTES {
399 return Err(std::io::ErrorKind::InvalidData.into());
400 }
401 let mut data = vec![0; bytes];
402 reader.read_exact(&mut data).await?;
403 serde_json::from_slice(&data).map_err(std::io::Error::other)
404}
405async fn write_frame<W: AsyncWrite + Unpin>(writer: &mut W, frame: &Frame) -> std::io::Result<()> {
406 let data = serde_json::to_vec(frame).map_err(std::io::Error::other)?;
407 if data.len() > MAX_FRAME_BYTES {
408 return Err(std::io::ErrorKind::InvalidData.into());
409 }
410 writer.write_u32(data.len() as u32).await?;
411 writer.write_all(&data).await?;
412 writer.flush().await
413}
414async fn connect(
415 endpoint: &Endpoint,
416 subscriber: bool,
417) -> std::io::Result<(transport::Stream, crate::config::HubEpoch)> {
418 let mut stream = endpoint.connect().await?;
419 write_frame(
420 &mut stream,
421 &Frame::Hello {
422 version: VERSION,
423 identity: endpoint.identity.clone(),
424 subscriber,
425 },
426 )
427 .await?;
428 match read_frame(&mut stream).await? {
429 Frame::Registered { epoch } if epoch.len() <= crate::types::MAX_ID_BYTES => {
430 Ok((stream, epoch))
431 }
432 _ => Err(std::io::ErrorKind::InvalidData.into()),
433 }
434}
435async fn publish(publisher: Arc<Publishing>) {
436 loop {
437 let wake = publisher.wake.notified();
438 if publisher.pending.lock().unwrap().dirty.is_empty() {
439 tokio::select! {
440 _ = wake => {},
441 _ = tokio::time::sleep(RECONNECT) => {
442 if Arc::strong_count(&publisher) == 1 {
444 break;
445 }
446 continue;
447 }
448 }
449 }
450 let (mut dirty, target) = {
451 let mut pending = publisher.pending.lock().unwrap();
452 (std::mem::take(&mut pending.dirty), pending.submitted)
453 };
454 let delivery = async {
456 let (mut stream, _) = connect(&publisher.endpoint, false).await?;
457 while let Some(change) = dirty.pop() {
458 write_frame(&mut stream, &Frame::Publish(change)).await?;
459 if !matches!(read_frame(&mut stream).await?, Frame::Delivered) {
460 return Err(std::io::ErrorKind::InvalidData.into());
461 }
462 }
463 Ok::<_, std::io::Error>(())
464 };
465 let _delivery = tokio::time::timeout(DELIVERY_TIMEOUT, delivery).await;
466 #[cfg(test)]
467 if !matches!(&_delivery, Ok(Ok(()))) {
468 *publisher.delivery_error.lock().unwrap() = Some(format!("{_delivery:?}"));
469 }
470 publisher.pending.lock().unwrap().delivered = target;
471 publisher.drained.notify_one();
472 }
473}
474
475async fn watch(
476 state: std::path::PathBuf,
477 inbox: Arc<Inbox>,
478 cancel: CancellationToken,
479 ready: oneshot::Sender<()>,
480) -> Result<(), String> {
481 let mut ready = Some(ready);
482 let endpoint = Endpoint::new(&state).ok();
483 let mut hub = None;
484 let mut probe = probe::Probe::new(state);
485 let mut reconcile =
488 tokio::time::interval_at(tokio::time::Instant::now() + RECONCILE, RECONCILE);
489 reconcile.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
490 let _ = probe.changed().await;
492 loop {
493 if cancel.is_cancelled() {
494 break;
495 }
496 if let Some(endpoint) = &endpoint {
497 if hub
498 .as_ref()
499 .is_some_and(|task: &tokio::task::JoinHandle<()>| task.is_finished())
500 {
501 hub = None;
502 }
503 if hub.is_none()
504 && let Ok(Some(owner)) = endpoint.elect()
505 && let Ok(listener) = endpoint.listen()
506 {
507 let endpoint = endpoint.clone();
508 let token = cancel.child_token();
509 hub = Some(tokio::spawn(async move {
510 let _ = serve(endpoint, listener, owner, token).await;
511 }));
512 }
513 let connection = tokio::time::timeout(DELIVERY_TIMEOUT, connect(endpoint, true));
514 tokio::select! {
515 _ = cancel.cancelled() => break,
516 result = connection => {
517 if let Ok(Ok((stream, epoch))) = result {
518 inbox.add(Change::Resync);
519 #[cfg(test)]
520 inbox.registration.send_replace(Some(epoch.clone()));
521 if let Some(ready) = ready.take() {
522 let _ = ready.send(());
523 }
524 let receive = receive(stream, epoch, inbox.clone());
525 tokio::pin!(receive);
526 loop {
527 tokio::select! {
528 _ = cancel.cancelled() => break,
529 _ = &mut receive => {
530 #[cfg(test)]
531 inbox.registration.send_replace(None);
532 inbox.add(Change::Resync);
533 break;
534 }
535 _ = reconcile.tick() => match probe.changed().await {
538 probe::Probed::Same => {}
539 probe::Probed::Written => inbox.add(Change::StateInvalidated),
540 probe::Probed::Replaced => inbox.add(Change::Resync),
541 }
542 }
543 }
544 }
545 }
546 }
547 }
548 if let Some(ready) = ready.take() {
549 inbox.add(Change::Resync);
550 let _ = ready.send(());
551 }
552 tokio::select! {
553 _ = cancel.cancelled() => break,
554 _ = reconcile.tick() => {
555 let _ = probe.changed().await;
556 inbox.add(Change::Resync);
559 },
560 _ = tokio::time::sleep(RECONNECT) => {},
561 }
562 }
563 if let Some(task) = hub {
564 let _ = crate::task::joined(task.await);
566 }
567 probe.close().await
568}
569async fn receive(
570 mut stream: transport::Stream,
571 epoch: crate::config::HubEpoch,
572 inbox: Arc<Inbox>,
573) -> std::io::Result<()> {
574 let mut previous = None;
575 loop {
576 match read_frame(&mut stream).await? {
577 Frame::Hint {
578 epoch: incoming,
579 sequence,
580 change,
581 } if incoming == epoch => {
582 if previous.is_some_and(|previous: u64| sequence != previous.wrapping_add(1)) {
583 inbox.add(Change::Resync);
584 }
585 previous = Some(sequence);
586 inbox.add(change);
587 #[cfg(test)]
588 inbox.received.send_replace(Some(sequence));
589 }
590 _ => return Err(std::io::ErrorKind::InvalidData.into()),
591 }
592 }
593}
594
595#[derive(Clone)]
596struct Hint {
597 sequence: u64,
598 change: Change,
599}
600async fn serve(
601 endpoint: Endpoint,
602 mut listener: transport::Listener,
603 _owner: std::fs::File,
604 cancel: CancellationToken,
605) -> std::io::Result<()> {
606 let epoch: crate::config::HubEpoch = crate::agent::uuid()
607 .map_err(std::io::Error::other)?
608 .parse()
609 .expect("UUID hub epoch");
610 let (sender, _) = broadcast::channel(MAX_BACKLOG);
611 let sequence = Arc::new(Mutex::new(0_u64));
612 let mut clients = JoinSet::new();
613 loop {
614 tokio::select! {
615 _ = cancel.cancelled() => break,
616 Some(result) = clients.join_next(), if !clients.is_empty() => {
617 let _ = crate::task::joined(result);
618 },
619 result = listener.accept(), if clients.len() < MAX_CLIENTS => {
620 let stream = result?;
621 let (identity, epoch, sender, sequence) = (
622 endpoint.identity.clone(),
623 epoch.clone(),
624 sender.clone(),
625 sequence.clone(),
626 );
627 clients.spawn(async move {
628 let _ = client(stream, identity, epoch, sender, sequence).await;
629 });
630 }
631 }
632 }
633 clients.abort_all();
634 while let Some(result) = clients.join_next().await {
635 let _ = crate::task::joined(result);
636 }
637 Ok(())
638}
639async fn client(
640 mut stream: transport::Stream,
641 identity: crate::config::EndpointId,
642 epoch: crate::config::HubEpoch,
643 sender: broadcast::Sender<Hint>,
644 sequence: Arc<Mutex<u64>>,
645) -> std::io::Result<()> {
646 transport::validate_peer(&stream)?;
649 let hello = tokio::time::timeout(DELIVERY_TIMEOUT, read_frame(&mut stream)).await??;
650 let Frame::Hello {
651 version: VERSION,
652 identity: incoming,
653 subscriber,
654 } = hello
655 else {
656 return Err(std::io::ErrorKind::InvalidData.into());
657 };
658 if incoming != identity {
659 return Err(std::io::ErrorKind::PermissionDenied.into());
660 }
661 let mut receiver = sender.subscribe();
663 tokio::time::timeout(
664 DELIVERY_TIMEOUT,
665 write_frame(
666 &mut stream,
667 &Frame::Registered {
668 epoch: epoch.clone(),
669 },
670 ),
671 )
672 .await??;
673 if subscriber {
674 let (mut reader, mut writer) = tokio::io::split(stream);
675 let outgoing = async {
676 loop {
677 let hint = match receiver.recv().await {
678 Ok(hint) => hint,
679 Err(broadcast::error::RecvError::Lagged(_)) => Hint {
680 sequence: *sequence.lock().unwrap(),
681 change: Change::Resync,
682 },
683 Err(_) => return Ok::<_, std::io::Error>(()),
684 };
685 tokio::time::timeout(
687 DELIVERY_TIMEOUT,
688 write_frame(
689 &mut writer,
690 &Frame::Hint {
691 epoch: epoch.clone(),
692 sequence: hint.sequence,
693 change: hint.change,
694 },
695 ),
696 )
697 .await??;
698 }
699 };
700 tokio::select! { result = outgoing => result, _ = reader.read_u8() => Ok(()) }
701 } else {
702 loop {
703 let frame = tokio::time::timeout(DELIVERY_TIMEOUT, read_frame(&mut stream)).await??;
704 let Frame::Publish(change) = frame else {
705 return Err(std::io::ErrorKind::InvalidData.into());
706 };
707 {
708 let mut sequence = sequence.lock().unwrap();
710 *sequence = sequence.wrapping_add(1);
711 let _ = sender.send(Hint {
712 sequence: *sequence,
713 change,
714 });
715 }
716 tokio::time::timeout(
717 DELIVERY_TIMEOUT,
718 write_frame(&mut stream, &Frame::Delivered),
719 )
720 .await??;
721 }
722 }
723}