mod probe;
#[cfg(test)]
pub(crate) mod test_support;
#[cfg(test)]
mod tests;
mod transport;
use crate::types::SessionId;
use serde::{Deserialize, Serialize};
use std::{
collections::{BTreeMap, BTreeSet},
path::Path,
sync::{Arc, Mutex, OnceLock, Weak},
time::Duration,
};
use tokio::{
io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt},
sync::{Notify, broadcast, oneshot},
task::JoinSet,
};
use tokio_util::sync::CancellationToken;
use transport::Endpoint;
pub const RECONCILE: Duration = Duration::from_secs(5);
const RECONNECT: Duration = Duration::from_millis(500);
const DELIVERY_TIMEOUT: Duration = Duration::from_millis(300);
const MAX_FRAME_BYTES: usize = 1024;
const MAX_STATES: usize = 32;
const MAX_CLIENTS: usize = 64;
const MAX_BACKLOG: usize = 64;
const MAX_DIRTY_SESSIONS: usize = 64;
const VERSION: u32 = 1;
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum Change {
StateInvalidated,
SessionInvalidated(SessionId),
Resync,
}
enum Dirty {
Pending {
state: bool,
sessions: BTreeSet<SessionId>,
},
Resync,
}
impl Default for Dirty {
fn default() -> Self {
Self::Pending {
state: false,
sessions: BTreeSet::new(),
}
}
}
impl Dirty {
fn add(&mut self, change: Change) {
let Self::Pending { state, sessions } = self else {
return;
};
match change {
Change::Resync => *self = Self::Resync,
Change::StateInvalidated => *state = true,
Change::SessionInvalidated(id) => {
sessions.insert(id);
if sessions.len() > MAX_DIRTY_SESSIONS {
*self = Self::Resync;
}
}
}
}
fn pop(&mut self) -> Option<Change> {
match self {
Self::Resync => {
*self = Self::default();
Some(Change::Resync)
}
Self::Pending { state, sessions } => {
if std::mem::take(state) {
Some(Change::StateInvalidated)
} else {
sessions.pop_first().map(Change::SessionInvalidated)
}
}
}
}
fn is_empty(&self) -> bool {
matches!(self, Self::Pending { state: false, sessions } if sessions.is_empty())
}
}
struct Inbox {
#[cfg(test)]
registration: tokio::sync::watch::Sender<Option<crate::config::HubEpoch>>,
#[cfg(test)]
received: tokio::sync::watch::Sender<Option<u64>>,
dirty: Mutex<Dirty>,
wake: Notify,
}
impl Default for Inbox {
fn default() -> Self {
Self {
#[cfg(test)]
registration: tokio::sync::watch::channel(None).0,
#[cfg(test)]
received: tokio::sync::watch::channel(None).0,
dirty: Mutex::default(),
wake: Notify::new(),
}
}
}
impl Inbox {
fn add(&self, change: Change) {
self.dirty.lock().unwrap().add(change);
self.wake.notify_one();
}
async fn next(&self) -> Change {
loop {
let wake = self.wake.notified();
if let Some(change) = self.dirty.lock().unwrap().pop() {
return change;
}
wake.await;
}
}
}
pub struct Subscription {
inbox: Arc<Inbox>,
cancel: CancellationToken,
task: tokio::task::JoinHandle<Result<(), String>>,
}
impl Subscription {
pub async fn new(state: &Path) -> Self {
let inbox = Arc::new(Inbox::default());
let cancel = CancellationToken::new();
let (ready, registered) = oneshot::channel();
let state = state.to_path_buf();
let task = tokio::spawn(watch(state, inbox.clone(), cancel.clone(), ready));
let subscription = Self {
inbox,
cancel,
task,
};
let _ = registered.await;
subscription
}
pub async fn close(mut self) -> Result<(), String> {
self.cancel.cancel();
crate::task::joined((&mut self.task).await).map_err(|error| error.to_string())?
}
pub async fn next_reconciled(&mut self) -> Change {
tokio::time::timeout(RECONCILE, self.next_state())
.await
.unwrap_or(Change::Resync)
}
pub async fn next(&mut self) -> Change {
self.inbox.next().await
}
pub async fn next_state(&mut self) -> Change {
loop {
let change = self.next().await;
if !matches!(change, Change::SessionInvalidated(_)) {
return change;
}
}
}
}
impl Drop for Subscription {
fn drop(&mut self) {
self.cancel.cancel();
}
}
#[derive(Default)]
struct Outbox {
dirty: Dirty,
submitted: u64,
delivered: u64,
}
struct Publishing {
endpoint: Endpoint,
pending: Mutex<Outbox>,
wake: Notify,
drained: Notify,
#[cfg(test)]
delivery_error: Mutex<Option<String>>,
}
#[derive(Clone, Default)]
pub(crate) struct Publisher(Option<Arc<Publishing>>);
static PUBLISHERS: OnceLock<Mutex<BTreeMap<crate::config::EndpointId, Weak<Publishing>>>> =
OnceLock::new();
impl Publisher {
pub(crate) fn new(state: &Path) -> Self {
let Ok(runtime) = tokio::runtime::Handle::try_current() else {
return Self::default();
};
let Ok(endpoint) = Endpoint::new(state) else {
return Self::default();
};
let mut publishers = PUBLISHERS.get_or_init(Default::default).lock().unwrap();
publishers.retain(|_, publisher| publisher.strong_count() != 0);
if let Some(publisher) = publishers.get(&endpoint.identity).and_then(Weak::upgrade) {
return Self(Some(publisher));
}
if publishers.len() >= MAX_STATES {
return Self::default();
}
let publisher = Arc::new(Publishing {
endpoint,
pending: Mutex::new(Outbox::default()),
wake: Notify::new(),
drained: Notify::new(),
#[cfg(test)]
delivery_error: Mutex::default(),
});
publishers.insert(
publisher.endpoint.identity.clone(),
Arc::downgrade(&publisher),
);
runtime.spawn(publish(publisher.clone()));
Self(Some(publisher))
}
pub(crate) fn notify(&self, change: Change) {
if let Some(publisher) = &self.0 {
let mut pending = publisher.pending.lock().unwrap();
pending.dirty.add(change);
pending.submitted = pending.submitted.wrapping_add(1);
drop(pending);
publisher.wake.notify_one();
}
}
}
pub async fn drain() {
let publishers: Vec<_> = PUBLISHERS
.get()
.map(|publishers| {
publishers
.lock()
.unwrap()
.values()
.filter_map(Weak::upgrade)
.collect()
})
.unwrap_or_default();
let wait = async {
for publisher in publishers {
let target = publisher.pending.lock().unwrap().submitted;
loop {
let drained = publisher.drained.notified();
if publisher.pending.lock().unwrap().delivered >= target {
break;
}
drained.await;
}
}
};
let _ = tokio::time::timeout(DELIVERY_TIMEOUT, wait).await;
}
#[derive(Serialize, Deserialize)]
enum Frame {
Hello {
version: u32,
identity: crate::config::EndpointId,
subscriber: bool,
},
Registered {
epoch: crate::config::HubEpoch,
},
Publish(Change),
Delivered,
Hint {
epoch: crate::config::HubEpoch,
sequence: u64,
change: Change,
},
}
async fn read_frame<R: AsyncRead + Unpin>(reader: &mut R) -> std::io::Result<Frame> {
let bytes = reader.read_u32().await? as usize;
if bytes == 0 || bytes > MAX_FRAME_BYTES {
return Err(std::io::ErrorKind::InvalidData.into());
}
let mut data = vec![0; bytes];
reader.read_exact(&mut data).await?;
serde_json::from_slice(&data).map_err(std::io::Error::other)
}
async fn write_frame<W: AsyncWrite + Unpin>(writer: &mut W, frame: &Frame) -> std::io::Result<()> {
let data = serde_json::to_vec(frame).map_err(std::io::Error::other)?;
if data.len() > MAX_FRAME_BYTES {
return Err(std::io::ErrorKind::InvalidData.into());
}
writer.write_u32(data.len() as u32).await?;
writer.write_all(&data).await?;
writer.flush().await
}
async fn connect(
endpoint: &Endpoint,
subscriber: bool,
) -> std::io::Result<(transport::Stream, crate::config::HubEpoch)> {
let mut stream = endpoint.connect().await?;
write_frame(
&mut stream,
&Frame::Hello {
version: VERSION,
identity: endpoint.identity.clone(),
subscriber,
},
)
.await?;
match read_frame(&mut stream).await? {
Frame::Registered { epoch } if epoch.len() <= crate::types::MAX_ID_BYTES => {
Ok((stream, epoch))
}
_ => Err(std::io::ErrorKind::InvalidData.into()),
}
}
async fn publish(publisher: Arc<Publishing>) {
loop {
let wake = publisher.wake.notified();
if publisher.pending.lock().unwrap().dirty.is_empty() {
tokio::select! {
_ = wake => {},
_ = tokio::time::sleep(RECONNECT) => {
if Arc::strong_count(&publisher) == 1 {
break;
}
continue;
}
}
}
let (mut dirty, target) = {
let mut pending = publisher.pending.lock().unwrap();
(std::mem::take(&mut pending.dirty), pending.submitted)
};
let delivery = async {
let (mut stream, _) = connect(&publisher.endpoint, false).await?;
while let Some(change) = dirty.pop() {
write_frame(&mut stream, &Frame::Publish(change)).await?;
if !matches!(read_frame(&mut stream).await?, Frame::Delivered) {
return Err(std::io::ErrorKind::InvalidData.into());
}
}
Ok::<_, std::io::Error>(())
};
let _delivery = tokio::time::timeout(DELIVERY_TIMEOUT, delivery).await;
#[cfg(test)]
if !matches!(&_delivery, Ok(Ok(()))) {
*publisher.delivery_error.lock().unwrap() = Some(format!("{_delivery:?}"));
}
publisher.pending.lock().unwrap().delivered = target;
publisher.drained.notify_one();
}
}
async fn watch(
state: std::path::PathBuf,
inbox: Arc<Inbox>,
cancel: CancellationToken,
ready: oneshot::Sender<()>,
) -> Result<(), String> {
let mut ready = Some(ready);
let endpoint = Endpoint::new(&state).ok();
let mut hub = None;
let mut probe = probe::Probe::new(state);
let mut reconcile =
tokio::time::interval_at(tokio::time::Instant::now() + RECONCILE, RECONCILE);
reconcile.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
let _ = probe.changed().await;
loop {
if cancel.is_cancelled() {
break;
}
if let Some(endpoint) = &endpoint {
if hub
.as_ref()
.is_some_and(|task: &tokio::task::JoinHandle<()>| task.is_finished())
{
hub = None;
}
if hub.is_none()
&& let Ok(Some(owner)) = endpoint.elect()
&& let Ok(listener) = endpoint.listen()
{
let endpoint = endpoint.clone();
let token = cancel.child_token();
hub = Some(tokio::spawn(async move {
let _ = serve(endpoint, listener, owner, token).await;
}));
}
let connection = tokio::time::timeout(DELIVERY_TIMEOUT, connect(endpoint, true));
tokio::select! {
_ = cancel.cancelled() => break,
result = connection => {
if let Ok(Ok((stream, epoch))) = result {
inbox.add(Change::Resync);
#[cfg(test)]
inbox.registration.send_replace(Some(epoch.clone()));
if let Some(ready) = ready.take() {
let _ = ready.send(());
}
let receive = receive(stream, epoch, inbox.clone());
tokio::pin!(receive);
loop {
tokio::select! {
_ = cancel.cancelled() => break,
_ = &mut receive => {
#[cfg(test)]
inbox.registration.send_replace(None);
inbox.add(Change::Resync);
break;
}
_ = reconcile.tick() => {
if probe.changed().await {
inbox.add(Change::Resync);
}
}
}
}
}
}
}
}
if let Some(ready) = ready.take() {
inbox.add(Change::Resync);
let _ = ready.send(());
}
tokio::select! {
_ = cancel.cancelled() => break,
_ = reconcile.tick() => {
let _ = probe.changed().await;
inbox.add(Change::Resync);
},
_ = tokio::time::sleep(RECONNECT) => {},
}
}
if let Some(task) = hub {
let _ = crate::task::joined(task.await);
}
probe.close().await
}
async fn receive(
mut stream: transport::Stream,
epoch: crate::config::HubEpoch,
inbox: Arc<Inbox>,
) -> std::io::Result<()> {
let mut previous = None;
loop {
match read_frame(&mut stream).await? {
Frame::Hint {
epoch: incoming,
sequence,
change,
} if incoming == epoch => {
if previous.is_some_and(|previous: u64| sequence != previous.wrapping_add(1)) {
inbox.add(Change::Resync);
}
previous = Some(sequence);
inbox.add(change);
#[cfg(test)]
inbox.received.send_replace(Some(sequence));
}
_ => return Err(std::io::ErrorKind::InvalidData.into()),
}
}
}
#[derive(Clone)]
struct Hint {
sequence: u64,
change: Change,
}
async fn serve(
endpoint: Endpoint,
mut listener: transport::Listener,
_owner: std::fs::File,
cancel: CancellationToken,
) -> std::io::Result<()> {
let epoch: crate::config::HubEpoch = crate::agent::uuid()
.map_err(std::io::Error::other)?
.parse()
.expect("UUID hub epoch");
let (sender, _) = broadcast::channel(MAX_BACKLOG);
let sequence = Arc::new(Mutex::new(0_u64));
let mut clients = JoinSet::new();
loop {
tokio::select! {
_ = cancel.cancelled() => break,
Some(result) = clients.join_next(), if !clients.is_empty() => {
let _ = crate::task::joined(result);
},
result = listener.accept(), if clients.len() < MAX_CLIENTS => {
let stream = result?;
let (identity, epoch, sender, sequence) = (
endpoint.identity.clone(),
epoch.clone(),
sender.clone(),
sequence.clone(),
);
clients.spawn(async move {
let _ = client(stream, identity, epoch, sender, sequence).await;
});
}
}
}
clients.abort_all();
while let Some(result) = clients.join_next().await {
let _ = crate::task::joined(result);
}
Ok(())
}
async fn client(
mut stream: transport::Stream,
identity: crate::config::EndpointId,
epoch: crate::config::HubEpoch,
sender: broadcast::Sender<Hint>,
sequence: Arc<Mutex<u64>>,
) -> std::io::Result<()> {
transport::validate_peer(&stream)?;
let hello = tokio::time::timeout(DELIVERY_TIMEOUT, read_frame(&mut stream)).await??;
let Frame::Hello {
version: VERSION,
identity: incoming,
subscriber,
} = hello
else {
return Err(std::io::ErrorKind::InvalidData.into());
};
if incoming != identity {
return Err(std::io::ErrorKind::PermissionDenied.into());
}
let mut receiver = sender.subscribe();
tokio::time::timeout(
DELIVERY_TIMEOUT,
write_frame(
&mut stream,
&Frame::Registered {
epoch: epoch.clone(),
},
),
)
.await??;
if subscriber {
let (mut reader, mut writer) = tokio::io::split(stream);
let outgoing = async {
loop {
let hint = match receiver.recv().await {
Ok(hint) => hint,
Err(broadcast::error::RecvError::Lagged(_)) => Hint {
sequence: *sequence.lock().unwrap(),
change: Change::Resync,
},
Err(_) => return Ok::<_, std::io::Error>(()),
};
tokio::time::timeout(
DELIVERY_TIMEOUT,
write_frame(
&mut writer,
&Frame::Hint {
epoch: epoch.clone(),
sequence: hint.sequence,
change: hint.change,
},
),
)
.await??;
}
};
tokio::select! { result = outgoing => result, _ = reader.read_u8() => Ok(()) }
} else {
loop {
let frame = tokio::time::timeout(DELIVERY_TIMEOUT, read_frame(&mut stream)).await??;
let Frame::Publish(change) = frame else {
return Err(std::io::ErrorKind::InvalidData.into());
};
{
let mut sequence = sequence.lock().unwrap();
*sequence = sequence.wrapping_add(1);
let _ = sender.send(Hint {
sequence: *sequence,
change,
});
}
tokio::time::timeout(
DELIVERY_TIMEOUT,
write_frame(&mut stream, &Frame::Delivered),
)
.await??;
}
}
}