rust-zero-core 0.1.0-alpha.1

Runtime primitives for rust-zero services
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100

use std::{
    collections::HashMap,
    hash::Hash,
    sync::{Arc, Mutex},
};

use tokio::sync::broadcast;

/// A typed, topic-based in-process message broker.
///
/// Slow subscribers receive [`broadcast::error::RecvError::Lagged`] rather than blocking
/// publishers, matching the fail-fast behavior expected from service event buses.
pub struct Broker<Topic, Message> {
    capacity: usize,
    topics: Arc<Mutex<HashMap<Topic, broadcast::Sender<Message>>>>,
}

impl<Topic, Message> Clone for Broker<Topic, Message> {
    fn clone(&self) -> Self {
        Self {
            capacity: self.capacity,
            topics: Arc::clone(&self.topics),
        }
    }
}

impl<Topic, Message> Broker<Topic, Message>
where
    Topic: Clone + Eq + Hash,
    Message: Clone,
{
    pub fn new(capacity: usize) -> Self {
        assert!(capacity > 0, "broker capacity must be greater than zero");
        Self {
            capacity,
            topics: Arc::new(Mutex::new(HashMap::new())),
        }
    }

    pub fn subscribe(&self, topic: Topic) -> Subscription<Message> {
        let mut topics = self.topics.lock().expect("broker lock poisoned");
        let sender = topics.entry(topic).or_insert_with(|| {
            let (sender, _) = broadcast::channel(self.capacity);
            sender
        });
        Subscription {
            receiver: sender.subscribe(),
        }
    }

    /// Publishes a message and returns the number of active subscribers.
    pub fn publish(&self, topic: Topic, message: Message) -> usize {
        let mut topics = self.topics.lock().expect("broker lock poisoned");
        let sender = topics.entry(topic).or_insert_with(|| {
            let (sender, _) = broadcast::channel(self.capacity);
            sender
        });
        sender.send(message).unwrap_or(0)
    }

    pub fn subscriber_count(&self, topic: &Topic) -> usize {
        self.topics
            .lock()
            .expect("broker lock poisoned")
            .get(topic)
            .map_or(0, broadcast::Sender::receiver_count)
    }
}

/// A subscription to one broker topic.
pub struct Subscription<Message> {
    receiver: broadcast::Receiver<Message>,
}

impl<Message> Subscription<Message>
where
    Message: Clone,
{
    pub async fn recv(&mut self) -> Result<Message, broadcast::error::RecvError> {
        self.receiver.recv().await
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn fans_messages_out_to_topic_subscribers() {
        let broker = Broker::new(8);
        let mut first = broker.subscribe("orders");
        let mut second = broker.subscribe("orders");
        let mut unrelated = broker.subscribe("users");

        assert_eq!(broker.publish("orders", 42), 2);
        assert_eq!(first.recv().await.unwrap(), 42);
        assert_eq!(second.recv().await.unwrap(), 42);
        assert!(unrelated.receiver.try_recv().is_err());
    }
}