pulses 0.2.0

A robust, high-performance background job processing library for Rust.
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
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130

//! Routes polled/reclaimed messages to the handler pools that subscribe to
//! their stream, based on a precomputed bitmask routing table.

use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;

use crate::core::Envelope;
use crate::handler_pool::HandlerTask;
use crate::handler_set::RoutingTable;

/// Errors raised while routing a message.
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub(crate) enum RouterError {
    /// A handler pool's mailbox has been closed (the pool stopped).
    #[error("handler channel {index} is closed")]
    ChannelClosed {
        /// Index of the closed handler channel.
        index: usize,
    },
}

/// Maps each stream to the set of handler pools subscribing to it, and forwards
/// messages to those pools' mailboxes.
///
/// Replaces a central dispatcher actor: the reader and reclaimer each hold a
/// (cheaply cloned) `Router` and forward directly, avoiding an extra hop and a
/// lock-guarded shared receiver.
pub(crate) struct Router<Tok> {
    routing_table: RoutingTable,
    senders: Vec<mpsc::Sender<HandlerTask<Tok>>>,
}

impl<Tok> Clone for Router<Tok> {
    fn clone(&self) -> Self { Self { routing_table: self.routing_table.clone(), senders: self.senders.clone() } }
}

impl<Tok: Clone> Router<Tok> {
    pub(crate) fn new(routing_table: RoutingTable, senders: Vec<mpsc::Sender<HandlerTask<Tok>>>) -> Self {
        Self { routing_table, senders }
    }

    /// Forward one message to every subscribing pool.
    ///
    /// The common single-subscriber case moves the message without cloning;
    /// fan-out clones the (cheap, refcounted) envelope and token per target.
    /// Returns immediately if cancellation fires mid-send.
    pub(crate) async fn route(
        &self, envelope: Envelope, token: Tok, cancellation: &CancellationToken,
    ) -> Result<(), RouterError> {
        let Some(&mask) = self.routing_table.get(&envelope.stream) else {
            tracing::warn!(stream = %envelope.stream, "no handler subscribes to stream; dropping message");
            return Ok(());
        };

        let target_count = mask.count_ones();
        if target_count == 0 {
            return Ok(());
        }

        if target_count == 1 {
            let index = mask.trailing_zeros() as usize;
            return self.send(index, HandlerTask { envelope, token }, cancellation).await;
        }

        let mut remaining = mask;
        while remaining != 0 {
            let index = remaining.trailing_zeros() as usize;
            remaining &= remaining - 1;
            let task = HandlerTask { envelope: envelope.clone(), token: token.clone() };
            self.send(index, task, cancellation).await?;
        }
        Ok(())
    }

    async fn send(
        &self, index: usize, task: HandlerTask<Tok>, cancellation: &CancellationToken,
    ) -> Result<(), RouterError> {
        let Some(sender) = self.senders.get(index) else {
            tracing::error!(index, "routing table referenced a missing handler channel");
            return Ok(());
        };

        tokio::select! {
            _ = cancellation.cancelled() => Ok(()),
            result = sender.send(task) => result.map_err(|_| RouterError::ChannelClosed { index }),
        }
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use tokio::sync::mpsc;
    use tokio_util::sync::CancellationToken;

    use super::Router;
    use crate::handler_set::RoutingTable;
    use crate::test_support::make_envelope;

    #[tokio::test]
    async fn routes_single_and_fanout() {
        let (tx0, mut rx0) = mpsc::channel(8);
        let (tx1, mut rx1) = mpsc::channel(8);
        let mut table = RoutingTable::new();
        table.insert(Arc::from("orders"), 1 << 0);
        table.insert(Arc::from("shared"), (1 << 0) | (1 << 1));
        let router = Router::new(table, vec![tx0, tx1]);
        let cancel = CancellationToken::new();

        router.route(make_envelope("orders", "1-1"), "t1".to_owned(), &cancel).await.unwrap();
        router.route(make_envelope("shared", "1-2"), "t2".to_owned(), &cancel).await.unwrap();

        assert_eq!(rx0.recv().await.unwrap().envelope.id.as_ref(), "1-1");
        assert_eq!(rx0.recv().await.unwrap().envelope.id.as_ref(), "1-2");
        assert_eq!(rx1.recv().await.unwrap().envelope.id.as_ref(), "1-2");
    }

    #[tokio::test]
    async fn unknown_stream_is_dropped() {
        let (tx0, mut rx0) = mpsc::channel(8);
        let mut table = RoutingTable::new();
        table.insert(Arc::from("orders"), 1 << 0);
        let router = Router::new(table, vec![tx0]);
        let cancel = CancellationToken::new();

        router.route(make_envelope("unknown", "1-1"), "t".to_owned(), &cancel).await.unwrap();
        assert!(rx0.try_recv().is_err());
    }
}