Skip to main content

ruststream_sqlx/inbox/
threads.rs

1//! A subscription on dedicated threads with pools of their own: how many threads, whether a
2//! partition key keeps to one, and how many connections each thread's pool opens.
3
4use std::num::{NonZeroU32, NonZeroUsize};
5
6use ruststream::runtime::{Declared, SubscriberBuilder, SubscriberSettings, Workers};
7
8use super::queue::InboxQueue;
9
10/// The dedicated threads of an `InboxQueue` subscription, and the connections each one opens for
11/// its handlers' queries: what [`InboxSettings::on_threads`](crate::InboxSettings::on_threads)
12/// mounts.
13///
14/// A handler on one of the threads queries through `Ctx<keys::Pool<..>>` on a pool of the
15/// thread's own, with the service pool's options and [`connections`](Self::connections)
16/// connections at most, one unless set. Its connections open on the thread and close with it.
17/// Knowing both numbers, the subscription counts `threads × connections` against the broker's
18/// [`connection_limit`](crate::SqlxBroker::connection_limit) when it opens.
19///
20/// # Examples
21///
22/// ```
23/// # #[cfg(feature = "postgres")]
24/// # mod demo {
25/// use ruststream::nonzero;
26/// use ruststream_sqlx::prelude::*;
27/// use sqlx::{PgPool, Postgres};
28/// # #[derive(Inbox, sqlx::FromRow)]
29/// # #[inbox(table = "ledger_jobs")]
30/// # pub struct LedgerJob { #[field(id)] id: i64, #[field(partition_key)] account: String, #[field(payload)] payload: Vec<u8> }
31/// # #[derive(serde::Deserialize)]
32/// # struct Entry { amount: i64 }
33///
34/// #[subscriber(InboxQueue::<LedgerJob>::new("entries"))]
35/// async fn post(entry: &Entry, Ctx(pool): Ctx<keys::Pool<Postgres>>) -> HandlerOutcome {
36///     let posted = sqlx::query("INSERT INTO postings (amount) VALUES ($1)")
37///         .bind(entry.amount)
38///         .execute(&pool)
39///         .await;
40///     if posted.is_ok() { HandlerOutcome::ack() } else { HandlerOutcome::retry() }
41/// }
42///
43/// pub fn app(pool: PgPool) -> RustStream {
44///     // Eight threads, an account's entries on one of them, one connection each.
45///     let threads = InboxThreads::new(nonzero!(8)).by_key();
46///     RustStream::new(AppInfo::new("ledger", "0.1.0")).with_broker(SqlxBroker::new(pool), |b| {
47///         b.include(post.on_threads(threads));
48///     })
49/// }
50/// # }
51/// # fn main() {}
52/// ```
53#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
54#[must_use]
55pub struct InboxThreads {
56    count: NonZeroUsize,
57    connections: NonZeroU32,
58    by_key: bool,
59}
60
61impl InboxThreads {
62    /// `count` dedicated threads, as `threads(count)`, each opening one connection.
63    pub const fn new(count: NonZeroUsize) -> Self {
64        Self {
65            count,
66            connections: NonZeroU32::MIN,
67            by_key: false,
68        }
69    }
70
71    /// Sends the deliveries of one partition key to one thread, as `threads(count, by_key)`, so
72    /// a key keeps its order.
73    pub const fn by_key(mut self) -> Self {
74        self.by_key = true;
75        self
76    }
77
78    /// Opens `connections` connections at most on each thread, for its handlers' queries.
79    ///
80    /// A thread runs its deliveries one at a time, so one connection serves a handler that holds
81    /// one connection at a time. A handler that holds two at once, or queries from tasks it
82    /// spawns on its thread, needs one more for each.
83    pub const fn connections(mut self, connections: NonZeroU32) -> Self {
84        self.connections = connections;
85        self
86    }
87
88    /// The connections the threads open at most together: threads times connections each.
89    pub(crate) fn total_connections(&self) -> u64 {
90        u64::try_from(self.count.get())
91            .unwrap_or(u64::MAX)
92            .saturating_mul(u64::from(self.connections.get()))
93    }
94
95    /// The connections each thread's pool opens at most.
96    pub(crate) const fn connections_each(&self) -> NonZeroU32 {
97        self.connections
98    }
99
100    /// The dispatch the core runs the threads with.
101    pub(crate) const fn workers(&self) -> Workers {
102        if self.by_key {
103            Workers::threads_keyed(self.count)
104        } else {
105            Workers::threads(self.count)
106        }
107    }
108}
109
110/// Recording a subscription's dedicated threads on its descriptor. Machinery behind
111/// [`InboxSettings::on_threads`](crate::InboxSettings::on_threads); never named in a service.
112#[doc(hidden)]
113#[diagnostic::on_unimplemented(
114    message = "`.on_threads(..)` runs an `InboxQueue` subscription on threads with pools of their own",
115    label = "this registration does not subscribe through `InboxQueue`",
116    note = "a subscription by name takes the framework's `threads(n)`: subscribe with \
117            `InboxQueue::<Row>::new(..)` for threads with pools of their own"
118)]
119pub trait ThreadsStep {
120    /// The registration with its threads recorded.
121    type Out;
122
123    /// Records `threads` on the registration's descriptor.
124    fn apply_threads(self, threads: InboxThreads) -> Self::Out;
125}
126
127impl<Def, Row, Mode, State, DefCodec> ThreadsStep
128    for SubscriberBuilder<Def, InboxQueue<Row, Mode>, State, DefCodec>
129where
130    Def: Declared,
131{
132    type Out = Self;
133
134    fn apply_threads(self, threads: InboxThreads) -> Self {
135        self.map_source(move |queue| queue.on_threads(threads))
136    }
137}