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 fn namespace(state: &Path) -> std::io::Result<std::path::PathBuf> {
234 Endpoint::new(state).map(|endpoint| endpoint.directory().to_owned())
235}
236
237pub struct Subscription {
240 inbox: Arc<Inbox>,
241 cancel: CancellationToken,
242 task: tokio::task::JoinHandle<Result<(), String>>,
243}
244impl Subscription {
245 pub async fn new(state: &Path) -> Self {
246 let inbox = Arc::new(Inbox::default());
247 let cancel = CancellationToken::new();
248 let (ready, registered) = oneshot::channel();
249 let state = state.to_path_buf();
250 let task = tokio::spawn(watch(state, inbox.clone(), cancel.clone(), ready));
251 let subscription = Self {
253 inbox,
254 cancel,
255 task,
256 };
257 let _ = registered.await;
259 subscription
260 }
261 pub async fn close(mut self) -> Result<(), String> {
264 self.cancel.cancel();
265 crate::task::joined((&mut self.task).await).map_err(|error| error.to_string())?
266 }
267 pub async fn next_reconciled(&mut self) -> Change {
270 tokio::time::timeout(RECONCILE, self.next_state())
271 .await
272 .unwrap_or(Change::Resync)
273 }
274 pub async fn next(&mut self) -> Change {
275 self.inbox.next().await
276 }
277 pub async fn next_state(&mut self) -> Change {
280 loop {
281 let change = self.next().await;
282 if !matches!(change, Change::SessionInvalidated(_)) {
283 return change;
284 }
285 }
286 }
287}
288impl Drop for Subscription {
289 fn drop(&mut self) {
290 self.cancel.cancel();
291 }
292}
293
294#[derive(Default)]
295struct Outbox {
296 dirty: Dirty,
297 submitted: u64,
298 delivered: u64,
299}
300struct Publishing {
301 endpoint: Endpoint,
302 pending: Mutex<Outbox>,
303 wake: Notify,
304 drained: Notify,
305 #[cfg(test)]
307 delivery_error: Mutex<Option<String>>,
308}
309#[derive(Clone, Default)]
311pub(crate) struct Publisher(Option<Arc<Publishing>>);
312static PUBLISHERS: OnceLock<Mutex<BTreeMap<crate::config::EndpointId, Weak<Publishing>>>> =
313 OnceLock::new();
314impl Publisher {
315 pub(crate) fn new(state: &Path) -> Self {
316 let Ok(runtime) = tokio::runtime::Handle::try_current() else {
317 return Self::default();
318 };
319 let Ok(endpoint) = Endpoint::new(state) else {
320 return Self::default();
321 };
322 let mut publishers = PUBLISHERS.get_or_init(Default::default).lock().unwrap();
323 publishers.retain(|_, publisher| publisher.strong_count() != 0);
324 if let Some(publisher) = publishers.get(&endpoint.identity).and_then(Weak::upgrade) {
325 return Self(Some(publisher));
326 }
327 if publishers.len() >= MAX_STATES {
328 return Self::default();
329 }
330 let publisher = Arc::new(Publishing {
331 endpoint,
332 pending: Mutex::new(Outbox::default()),
333 wake: Notify::new(),
334 drained: Notify::new(),
335 #[cfg(test)]
336 delivery_error: Mutex::default(),
337 });
338 publishers.insert(
339 publisher.endpoint.identity.clone(),
340 Arc::downgrade(&publisher),
341 );
342 runtime.spawn(publish(publisher.clone()));
343 Self(Some(publisher))
344 }
345 pub(crate) fn notify(&self, change: Change) {
346 if let Some(publisher) = &self.0 {
347 let mut pending = publisher.pending.lock().unwrap();
348 pending.dirty.add(change);
349 pending.submitted = pending.submitted.wrapping_add(1);
350 drop(pending);
351 publisher.wake.notify_one();
352 }
353 }
354}
355
356pub async fn drain() {
358 let publishers: Vec<_> = PUBLISHERS
359 .get()
360 .map(|publishers| {
361 publishers
362 .lock()
363 .unwrap()
364 .values()
365 .filter_map(Weak::upgrade)
366 .collect()
367 })
368 .unwrap_or_default();
369 let wait = async {
370 for publisher in publishers {
371 let target = publisher.pending.lock().unwrap().submitted;
372 loop {
373 let drained = publisher.drained.notified();
374 if publisher.pending.lock().unwrap().delivered >= target {
375 break;
376 }
377 drained.await;
378 }
379 }
380 };
381 let _ = tokio::time::timeout(DELIVERY_TIMEOUT, wait).await;
382}
383
384#[derive(Serialize, Deserialize)]
385enum Frame {
386 Hello {
387 version: u32,
388 identity: crate::config::EndpointId,
389 subscriber: bool,
390 },
391 Registered {
392 epoch: crate::config::HubEpoch,
393 },
394 Publish(Change),
395 Delivered,
396 Hint {
397 epoch: crate::config::HubEpoch,
398 sequence: u64,
399 change: Change,
400 },
401}
402async fn read_frame<R: AsyncRead + Unpin>(reader: &mut R) -> std::io::Result<Frame> {
403 let bytes = reader.read_u32().await? as usize;
404 if bytes == 0 || bytes > MAX_FRAME_BYTES {
405 return Err(std::io::ErrorKind::InvalidData.into());
406 }
407 let mut data = vec![0; bytes];
408 reader.read_exact(&mut data).await?;
409 serde_json::from_slice(&data).map_err(std::io::Error::other)
410}
411async fn write_frame<W: AsyncWrite + Unpin>(writer: &mut W, frame: &Frame) -> std::io::Result<()> {
412 let data = serde_json::to_vec(frame).map_err(std::io::Error::other)?;
413 if data.len() > MAX_FRAME_BYTES {
414 return Err(std::io::ErrorKind::InvalidData.into());
415 }
416 writer.write_u32(data.len() as u32).await?;
417 writer.write_all(&data).await?;
418 writer.flush().await
419}
420async fn connect(
421 endpoint: &Endpoint,
422 subscriber: bool,
423) -> std::io::Result<(transport::Stream, crate::config::HubEpoch)> {
424 let mut stream = endpoint.connect().await?;
425 write_frame(
426 &mut stream,
427 &Frame::Hello {
428 version: VERSION,
429 identity: endpoint.identity.clone(),
430 subscriber,
431 },
432 )
433 .await?;
434 match read_frame(&mut stream).await? {
435 Frame::Registered { epoch } if epoch.len() <= crate::types::MAX_ID_BYTES => {
436 Ok((stream, epoch))
437 }
438 _ => Err(std::io::ErrorKind::InvalidData.into()),
439 }
440}
441async fn publish(publisher: Arc<Publishing>) {
442 loop {
443 let wake = publisher.wake.notified();
444 if publisher.pending.lock().unwrap().dirty.is_empty() {
445 tokio::select! {
446 _ = wake => {},
447 _ = tokio::time::sleep(RECONNECT) => {
448 if Arc::strong_count(&publisher) == 1 {
450 break;
451 }
452 continue;
453 }
454 }
455 }
456 let (mut dirty, target) = {
457 let mut pending = publisher.pending.lock().unwrap();
458 (std::mem::take(&mut pending.dirty), pending.submitted)
459 };
460 let delivery = async {
462 let (mut stream, _) = connect(&publisher.endpoint, false).await?;
463 while let Some(change) = dirty.pop() {
464 write_frame(&mut stream, &Frame::Publish(change)).await?;
465 if !matches!(read_frame(&mut stream).await?, Frame::Delivered) {
466 return Err(std::io::ErrorKind::InvalidData.into());
467 }
468 }
469 Ok::<_, std::io::Error>(())
470 };
471 let _delivery = tokio::time::timeout(DELIVERY_TIMEOUT, delivery).await;
472 #[cfg(test)]
473 if !matches!(&_delivery, Ok(Ok(()))) {
474 *publisher.delivery_error.lock().unwrap() = Some(format!("{_delivery:?}"));
475 }
476 publisher.pending.lock().unwrap().delivered = target;
477 publisher.drained.notify_one();
478 }
479}
480
481async fn watch(
482 state: std::path::PathBuf,
483 inbox: Arc<Inbox>,
484 cancel: CancellationToken,
485 ready: oneshot::Sender<()>,
486) -> Result<(), String> {
487 let mut ready = Some(ready);
488 let endpoint = Endpoint::new(&state).ok();
489 let mut hub = None;
490 let mut probe = probe::Probe::new(state);
491 let mut reconcile =
494 tokio::time::interval_at(tokio::time::Instant::now() + RECONCILE, RECONCILE);
495 reconcile.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
496 let _ = probe.changed().await;
498 loop {
499 if cancel.is_cancelled() {
500 break;
501 }
502 if let Some(endpoint) = &endpoint {
503 if hub
504 .as_ref()
505 .is_some_and(|task: &tokio::task::JoinHandle<()>| task.is_finished())
506 {
507 hub = None;
508 }
509 if hub.is_none()
510 && let Ok(Some(owner)) = endpoint.elect()
511 && let Ok(listener) = endpoint.listen()
512 {
513 let endpoint = endpoint.clone();
514 let token = cancel.child_token();
515 hub = Some(tokio::spawn(async move {
516 let _ = serve(endpoint, listener, owner, token).await;
517 }));
518 }
519 let connection = tokio::time::timeout(DELIVERY_TIMEOUT, connect(endpoint, true));
520 tokio::select! {
521 _ = cancel.cancelled() => break,
522 result = connection => {
523 if let Ok(Ok((stream, epoch))) = result {
524 inbox.add(Change::Resync);
525 #[cfg(test)]
526 inbox.registration.send_replace(Some(epoch.clone()));
527 if let Some(ready) = ready.take() {
528 let _ = ready.send(());
529 }
530 let receive = receive(stream, epoch, inbox.clone());
531 tokio::pin!(receive);
532 loop {
533 tokio::select! {
534 _ = cancel.cancelled() => break,
535 _ = &mut receive => {
536 #[cfg(test)]
537 inbox.registration.send_replace(None);
538 inbox.add(Change::Resync);
539 break;
540 }
541 _ = reconcile.tick() => match probe.changed().await {
544 probe::Probed::Same => {}
545 probe::Probed::Written => inbox.add(Change::StateInvalidated),
546 probe::Probed::Replaced => inbox.add(Change::Resync),
547 }
548 }
549 }
550 }
551 }
552 }
553 }
554 if let Some(ready) = ready.take() {
555 inbox.add(Change::Resync);
556 let _ = ready.send(());
557 }
558 tokio::select! {
559 _ = cancel.cancelled() => break,
560 _ = reconcile.tick() => {
561 let _ = probe.changed().await;
562 inbox.add(Change::Resync);
565 },
566 _ = tokio::time::sleep(RECONNECT) => {},
567 }
568 }
569 if let Some(task) = hub {
570 let _ = crate::task::joined(task.await);
572 }
573 probe.close().await
574}
575async fn receive(
576 mut stream: transport::Stream,
577 epoch: crate::config::HubEpoch,
578 inbox: Arc<Inbox>,
579) -> std::io::Result<()> {
580 let mut previous = None;
581 loop {
582 match read_frame(&mut stream).await? {
583 Frame::Hint {
584 epoch: incoming,
585 sequence,
586 change,
587 } if incoming == epoch => {
588 if previous.is_some_and(|previous: u64| sequence != previous.wrapping_add(1)) {
589 inbox.add(Change::Resync);
590 }
591 previous = Some(sequence);
592 inbox.add(change);
593 #[cfg(test)]
594 inbox.received.send_replace(Some(sequence));
595 }
596 _ => return Err(std::io::ErrorKind::InvalidData.into()),
597 }
598 }
599}
600
601#[derive(Clone)]
602struct Hint {
603 sequence: u64,
604 change: Change,
605}
606async fn serve(
607 endpoint: Endpoint,
608 mut listener: transport::Listener,
609 _owner: std::fs::File,
610 cancel: CancellationToken,
611) -> std::io::Result<()> {
612 let epoch: crate::config::HubEpoch = crate::agent::uuid()
613 .map_err(std::io::Error::other)?
614 .parse()
615 .expect("UUID hub epoch");
616 let (sender, _) = broadcast::channel(MAX_BACKLOG);
617 let sequence = Arc::new(Mutex::new(0_u64));
618 let mut clients = JoinSet::new();
619 loop {
620 tokio::select! {
621 _ = cancel.cancelled() => break,
622 Some(result) = clients.join_next(), if !clients.is_empty() => {
623 let _ = crate::task::joined(result);
624 },
625 result = listener.accept(), if clients.len() < MAX_CLIENTS => {
626 let stream = result?;
627 let (identity, epoch, sender, sequence) = (
628 endpoint.identity.clone(),
629 epoch.clone(),
630 sender.clone(),
631 sequence.clone(),
632 );
633 clients.spawn(async move {
634 let _ = client(stream, identity, epoch, sender, sequence).await;
635 });
636 }
637 }
638 }
639 clients.abort_all();
640 while let Some(result) = clients.join_next().await {
641 let _ = crate::task::joined(result);
642 }
643 Ok(())
644}
645async fn client(
646 mut stream: transport::Stream,
647 identity: crate::config::EndpointId,
648 epoch: crate::config::HubEpoch,
649 sender: broadcast::Sender<Hint>,
650 sequence: Arc<Mutex<u64>>,
651) -> std::io::Result<()> {
652 transport::validate_peer(&stream)?;
655 let hello = tokio::time::timeout(DELIVERY_TIMEOUT, read_frame(&mut stream)).await??;
656 let Frame::Hello {
657 version: VERSION,
658 identity: incoming,
659 subscriber,
660 } = hello
661 else {
662 return Err(std::io::ErrorKind::InvalidData.into());
663 };
664 if incoming != identity {
665 return Err(std::io::ErrorKind::PermissionDenied.into());
666 }
667 let mut receiver = sender.subscribe();
669 tokio::time::timeout(
670 DELIVERY_TIMEOUT,
671 write_frame(
672 &mut stream,
673 &Frame::Registered {
674 epoch: epoch.clone(),
675 },
676 ),
677 )
678 .await??;
679 if subscriber {
680 let (mut reader, mut writer) = tokio::io::split(stream);
681 let outgoing = async {
682 loop {
683 let hint = match receiver.recv().await {
684 Ok(hint) => hint,
685 Err(broadcast::error::RecvError::Lagged(_)) => Hint {
686 sequence: *sequence.lock().unwrap(),
687 change: Change::Resync,
688 },
689 Err(_) => return Ok::<_, std::io::Error>(()),
690 };
691 tokio::time::timeout(
693 DELIVERY_TIMEOUT,
694 write_frame(
695 &mut writer,
696 &Frame::Hint {
697 epoch: epoch.clone(),
698 sequence: hint.sequence,
699 change: hint.change,
700 },
701 ),
702 )
703 .await??;
704 }
705 };
706 tokio::select! { result = outgoing => result, _ = reader.read_u8() => Ok(()) }
707 } else {
708 loop {
709 let frame = tokio::time::timeout(DELIVERY_TIMEOUT, read_frame(&mut stream)).await??;
710 let Frame::Publish(change) = frame else {
711 return Err(std::io::ErrorKind::InvalidData.into());
712 };
713 {
714 let mut sequence = sequence.lock().unwrap();
716 *sequence = sequence.wrapping_add(1);
717 let _ = sender.send(Hint {
718 sequence: *sequence,
719 change,
720 });
721 }
722 tokio::time::timeout(
723 DELIVERY_TIMEOUT,
724 write_frame(&mut stream, &Frame::Delivered),
725 )
726 .await??;
727 }
728 }
729}