macula_rust/station_link/
pubsub.rs

1//! PubSub, as macula 12's link does it. A PUBLISH carries a publication
2//! signed with the link's identity key; SUBSCRIBE and UNSUBSCRIBE are control
3//! frames naming the realm, the topic and this node, with no
4//! acknowledgement; an EVENT carries a publication that is verified before it
5//! is delivered, and delivered once however many copies arrive, until it
6//! expires.
7
8use std::collections::HashMap;
9use std::sync::atomic::{AtomicU64, Ordering};
10use std::sync::{Arc, Mutex, Weak};
11
12use tokio::sync::mpsc;
13
14use crate::cbor::Value;
15use crate::frame::{self, PublicationSpec};
16use crate::node_key::NodeKey;
17
18use super::{now_ms, Inner, Link, LinkError};
19
20/// How many events a subscription holds that its reader has not taken; an
21/// event arriving at a full subscription is dropped and counted.
22const EVENT_BUFFER: usize = 64;
23
24/// What [`Link::publish`] sends: the realm and topic, the payload, and its
25/// time to live, `None` for macula's 10 minutes.
26#[derive(Debug, Clone, PartialEq)]
27pub struct Publication {
28    pub realm: [u8; 32],
29    pub topic: String,
30    pub payload: Value,
31    pub ttl_ms: Option<u64>,
32}
33
34/// A publication a subscription heard, verified: who published it (the key
35/// id its signature verified under), where, its seq and time, the payload,
36/// and how it arrived (direct or plumtree).
37#[derive(Debug, Clone, PartialEq)]
38pub struct Event {
39    pub publisher: [u8; 32],
40    pub realm: [u8; 32],
41    pub topic: String,
42    pub seq: u64,
43    pub published_at: u64,
44    pub payload: Value,
45    pub delivered_via: String,
46}
47
48/// Numbers one publisher's publications: the first is the wall clock in
49/// microseconds, and each after is one more than the last, or the clock,
50/// whichever is later, as macula_publication_seq does. Links of one identity
51/// key share one, so their seqs never repeat.
52#[derive(Debug, Default)]
53pub struct PublicationSeq {
54    last: Mutex<u64>,
55}
56
57impl PublicationSeq {
58    /// The seq for the next publication.
59    pub fn next(&self) -> u64 {
60        let now = std::time::SystemTime::now()
61            .duration_since(std::time::UNIX_EPOCH)
62            .map(|d| d.as_micros() as u64)
63            .unwrap_or(0);
64        let mut last = self.last.lock().unwrap_or_else(|p| p.into_inner());
65        *last = now.max(*last + 1);
66        *last
67    }
68}
69
70/// Remembers the publications a node delivered, by publication hash, until
71/// each expires, so an event heard on several links, or twice on one, is
72/// delivered once. The links of one node share one.
73#[derive(Debug, Default)]
74pub struct EventDedup {
75    seen: Mutex<HashMap<[u8; 48], u64>>,
76}
77
78impl EventDedup {
79    /// Whether `hash` is new at `now_ms`, remembering it until `expires_at`.
80    fn first(&self, hash: [u8; 48], expires_at: u64, now_ms: i64) -> bool {
81        let mut seen = self.seen.lock().unwrap_or_else(|p| p.into_inner());
82        if seen.contains_key(&hash) {
83            return false;
84        }
85        seen.retain(|_, until| *until >= now_ms as u64);
86        seen.insert(hash, expires_at);
87        true
88    }
89}
90
91/// A PUBLISH signed once, to be sent on several links of one node: every copy
92/// is the same publication, so a subscriber hearing it on several links
93/// delivers it once.
94#[derive(Debug, Clone, PartialEq)]
95pub struct SignedPublication {
96    frame: Value,
97}
98
99impl SignedPublication {
100    /// Signs `p` with `key` at the next seq of `seq`.
101    pub fn sign(
102        key: &NodeKey,
103        seq: &PublicationSeq,
104        p: Publication,
105    ) -> Result<SignedPublication, LinkError> {
106        let frame = frame::sign_publish(
107            &PublicationSpec {
108                realm: p.realm,
109                topic: p.topic,
110                seq: seq.next(),
111                published_at: now_ms() as u64,
112                payload: p.payload,
113                ttl_ms: p.ttl_ms,
114            },
115            key,
116        )?;
117        Ok(SignedPublication { frame })
118    }
119}
120
121static NEXT_SUBSCRIBER: AtomicU64 = AtomicU64::new(1);
122
123/// One subscriber of a realm and topic on a link.
124pub(super) struct SubscriberSlot {
125    id: u64,
126    events: mpsc::Sender<Event>,
127}
128
129/// Ends every subscription to `key` on the link when the station refused
130/// its SUBSCRIBE: their events stop, and each one's `recv` gives `None`.
131pub(super) fn subscribe_refused(inner: &Inner, key: &([u8; 32], String)) {
132    inner.lock().subs.remove(key);
133}
134
135/// One subscription to a realm and topic on a link, until
136/// [`Subscription::unsubscribe`] or the link ends.
137pub struct Subscription {
138    link: Weak<Inner>,
139    key: ([u8; 32], String),
140    id: u64,
141    events: mpsc::Receiver<Event>,
142    unsubscribed: bool,
143}
144
145impl Subscription {
146    /// The next event, or `None` once the subscription has ended.
147    pub async fn recv(&mut self) -> Option<Event> {
148        self.events.recv().await
149    }
150
151    /// Ends the subscription, and sends UNSUBSCRIBE once no other
152    /// subscription on the link holds its realm and topic.
153    pub async fn unsubscribe(mut self) -> Result<(), LinkError> {
154        self.unsubscribed = true;
155        let Some(inner) = self.link.upgrade() else {
156            return Ok(());
157        };
158        if drop_subscriber(&inner, &self.key, self.id) {
159            let frame =
160                frame::unsubscribe_frame(self.key.1.as_bytes(), &self.key.0, &inner.self_id)?;
161            inner.send_control(&frame).await?;
162        }
163        Ok(())
164    }
165}
166
167impl Drop for Subscription {
168    /// A subscription dropped without unsubscribing unsubscribes as it goes.
169    fn drop(&mut self) {
170        if self.unsubscribed {
171            return;
172        }
173        let Some(inner) = self.link.upgrade() else {
174            return;
175        };
176        if !drop_subscriber(&inner, &self.key, self.id) {
177            return;
178        }
179        let key = self.key.clone();
180        if let Ok(runtime) = tokio::runtime::Handle::try_current() {
181            runtime.spawn(unsubscribe_dropped(inner, key));
182        }
183    }
184}
185
186/// Sends the UNSUBSCRIBE of a subscription dropped without unsubscribing,
187/// for whatever it is worth now.
188async fn unsubscribe_dropped(inner: Arc<Inner>, key: ([u8; 32], String)) {
189    if let Ok(frame) = frame::unsubscribe_frame(key.1.as_bytes(), &key.0, &inner.self_id) {
190        let _ = inner.send_control(&frame).await;
191    }
192}
193
194/// Adds a subscriber, and whether it is the first on its realm and topic;
195/// the link's end when it has ended.
196fn add_subscriber(
197    inner: &Inner,
198    key: &([u8; 32], String),
199    id: u64,
200    events: mpsc::Sender<Event>,
201) -> Result<bool, LinkError> {
202    let mut state = inner.lock();
203    if let Some(e) = &state.ended {
204        return Err(e.clone());
205    }
206    let slots = state.subs.entry(key.clone()).or_default();
207    slots.push(SubscriberSlot { id, events });
208    Ok(slots.len() == 1)
209}
210
211/// Sends the SUBSCRIBE for `topic` in `realm`.
212async fn send_subscribe(inner: &Inner, realm: &[u8; 32], topic: &str) -> Result<(), LinkError> {
213    let frame = frame::subscribe_frame(topic.as_bytes(), realm, &inner.self_id);
214    match frame {
215        Ok(frame) => inner.send_control(&frame).await,
216        Err(e) => Err(e.into()),
217    }
218}
219
220/// Removes a subscriber, and whether it was the last on its realm and topic.
221fn drop_subscriber(inner: &Inner, key: &([u8; 32], String), id: u64) -> bool {
222    let mut state = inner.lock();
223    let Some(slots) = state.subs.get_mut(key) else {
224        return false;
225    };
226    slots.retain(|s| s.id != id);
227    if slots.is_empty() {
228        state.subs.remove(key);
229        return true;
230    }
231    false
232}
233
234impl Link {
235    /// Signs `p` as a PUBLISH and sends it.
236    pub async fn publish(&self, p: Publication) -> Result<(), LinkError> {
237        let signed = SignedPublication::sign(&self.inner.key, &self.inner.publication_seq, p)?;
238        self.publish_signed(&signed).await
239    }
240
241    /// Sends a publication signed once for several links, as it is.
242    pub async fn publish_signed(&self, p: &SignedPublication) -> Result<(), LinkError> {
243        self.inner.send_control(&p.frame).await
244    }
245
246    /// Subscribes to `topic` in `realm`: the first subscription to a realm
247    /// and topic on the link sends SUBSCRIBE.
248    pub async fn subscribe(
249        &self,
250        realm: &[u8; 32],
251        topic: &str,
252    ) -> Result<Subscription, LinkError> {
253        let key = (*realm, topic.to_string());
254        let (events_tx, events) = mpsc::channel(EVENT_BUFFER);
255        let id = NEXT_SUBSCRIBER.fetch_add(1, Ordering::Relaxed);
256        let first = add_subscriber(&self.inner, &key, id, events_tx)?;
257        let subscription = Subscription {
258            link: Arc::downgrade(&self.inner),
259            key: key.clone(),
260            id,
261            events,
262            unsubscribed: false,
263        };
264        let sent = if first {
265            send_subscribe(&self.inner, realm, topic).await
266        } else {
267            Ok(())
268        };
269        if let Err(e) = sent {
270            drop_subscriber(&self.inner, &key, id);
271            let mut subscription = subscription;
272            subscription.unsubscribed = true;
273            return Err(e);
274        }
275        Ok(subscription)
276    }
277}
278
279/// An EVENT's publication, once it verifies, delivered to every subscription
280/// on its realm and topic, unless it was delivered before.
281pub(super) fn evented(inner: &Arc<Inner>, v: &Value) {
282    let now = now_ms();
283    let Ok(publication) = frame::verify_publication(v, inner.profile, now) else {
284        inner.count("event_unverified");
285        return;
286    };
287    let delivered_via = frame::delivered_via(v).to_string();
288    if !inner
289        .dedup
290        .first(publication.publication_hash, publication.expires_at, now)
291    {
292        inner.count("event_duplicate");
293        return;
294    }
295    let event = Event {
296        publisher: publication.publisher,
297        realm: publication.realm,
298        topic: publication.topic.clone(),
299        seq: publication.seq,
300        published_at: publication.published_at,
301        payload: publication.payload,
302        delivered_via,
303    };
304    let mut state = inner.lock();
305    let Some(slots) = state.subs.get(&(publication.realm, publication.topic)) else {
306        *state
307            .unrouted
308            .entry("event_unsubscribed".into())
309            .or_default() += 1;
310        return;
311    };
312    let overflowed = slots
313        .iter()
314        .filter(|s| s.events.try_send(event.clone()).is_err())
315        .count();
316    if overflowed > 0 {
317        *state.unrouted.entry("event_overflow".into()).or_default() += overflowed as u64;
318    }
319}
macula_rust/station_link/pubsub.rs

macula_rust/station_link/
pubsub.rs
