brokapi_runtime/
shutdown.rs

1//! Graceful / immediate shutdown.
2
3use std::time::Instant;
4
5use brokapi_core::{ShutdownMode, ShutdownReport};
6use tokio_util::sync::CancellationToken;
7
8use crate::command_gate::CommandGate;
9use crate::orders::OrderRuntime;
10
11/// Shutdown coordinator for one connector instance: it closes the command gate, gives up waiting for
12/// in-flight commands and cancels the background tasks.
13///
14/// The controller does no network I/O. The connector supplies the work that needs the network (the
15/// unsubscribe calls) as a closure. Orders at the broker are never cancelled by a shutdown: only
16/// the local wait for an answer is given up (orders in `InFlight` become `Unknown`).
17///
18/// # Examples
19///
20/// ```
21/// use std::time::{Duration, Instant};
22/// use brokapi_core::{ClientOrderId, Error, ShutdownMode};
23/// use brokapi_runtime::{CommandGate, OrderRuntime, ShutdownController};
24///
25/// let mut gate = CommandGate::new();
26/// gate.mark_reconciled();
27/// let mut ctl = ShutdownController::new(gate);
28/// let task_token = ctl.child_token();
29///
30/// let mut orders = OrderRuntime::new();
31/// let now = Instant::now();
32/// let cid = ClientOrderId::new_v4();
33/// orders.register_place(cid, now);
34///
35/// // Two subscriptions were closed, none failed.
36/// let report = ctl.shutdown_local(
37///     ShutdownMode::Graceful { deadline: Duration::from_secs(5) },
38///     now,
39///     &mut orders,
40///     || (2, 0),
41/// );
42///
43/// assert_eq!((report.unsubscribed, report.unsubscribe_failed), (2, 0));
44/// assert_eq!(report.inflight_abandoned, vec![cid]);
45/// assert!(task_token.is_cancelled());
46/// assert!(matches!(ctl.gate().check_command(), Err(Error::ShuttingDown)));
47/// ```
48#[derive(Debug)]
49pub struct ShutdownController {
50    /// Parent cancellation token. Cancelled by [`shutdown_local`](Self::shutdown_local).
51    pub root: CancellationToken,
52    gate: CommandGate,
53}
54
55impl ShutdownController {
56    /// Creates a controller that owns the connector's command gate.
57    pub fn new(gate: CommandGate) -> Self {
58        Self {
59            root: CancellationToken::new(),
60            gate,
61        }
62    }
63
64    /// A token for one background task. It is cancelled when the controller shuts down.
65    pub fn child_token(&self) -> CancellationToken {
66        self.root.child_token()
67    }
68
69    /// The command gate, for checking commands.
70    pub fn gate(&self) -> &CommandGate {
71        &self.gate
72    }
73
74    /// The command gate, for marking it reconciled or degraded.
75    pub fn gate_mut(&mut self) -> &mut CommandGate {
76        &mut self.gate
77    }
78
79    /// Runs the local shutdown steps in order and returns a report.
80    ///
81    /// 1. Closes the command gate: new commands fail with [`brokapi_core::Error::ShuttingDown`].
82    /// 2. In [`ShutdownMode::Graceful`] calls `unsubscribe` once. It returns how many
83    ///    subscriptions were closed and how many failed. In [`ShutdownMode::Immediate`] it is not
84    ///    called and both counts are zero.
85    /// 3. Moves in-flight orders to `Unknown` (see [`OrderRuntime::abandon_inflight`]).
86    /// 4. Cancels the root token, which stops every task made with
87    ///    [`child_token`](Self::child_token).
88    ///
89    /// The `deadline` of `Graceful` is not applied here. This method runs to completion. Use
90    /// [`wait_graceful_deadline`](Self::wait_graceful_deadline) to wait for it. Calling the method
91    /// twice is harmless: the second report has no abandoned orders.
92    pub fn shutdown_local(
93        &mut self,
94        mode: ShutdownMode,
95        now: Instant,
96        orders: &mut OrderRuntime,
97        mut unsubscribe: impl FnMut() -> (usize, usize),
98    ) -> ShutdownReport {
99        self.gate.begin_shutdown();
100
101        let (unsubscribed, unsubscribe_failed) = match mode {
102            ShutdownMode::Graceful { .. } => unsubscribe(),
103            // `Immediate`, and a mode of a later version: close nothing in the protocol.
104            _ => (0, 0),
105        };
106
107        let inflight_abandoned = orders.abandon_inflight(now);
108        self.root.cancel();
109
110        ShutdownReport::new(unsubscribed, unsubscribe_failed, inflight_abandoned)
111    }
112
113    /// Sleeps for the `deadline` of a `Graceful` mode, and returns at once for `Immediate`.
114    ///
115    /// It always sleeps the full deadline. It does not return early when the work is done, so use
116    /// it only to bound a wait together with your own completion signal.
117    pub async fn wait_graceful_deadline(&self, mode: ShutdownMode) {
118        if let ShutdownMode::Graceful { deadline } = mode {
119            tokio::time::sleep(deadline).await;
120        }
121    }
122}
123
124#[cfg(test)]
125mod tests {
126    use super::*;
127    use brokapi_core::{ClientOrderId, OrderState, ShutdownMode};
128
129    fn cid(n: u8) -> ClientOrderId {
130        ClientOrderId::from_bytes([n; 16])
131    }
132
133    #[tokio::test(start_paused = true)]
134    async fn shutdown_abandons_inflight_orders() {
135        let gate = CommandGate::new();
136        let mut ctl = ShutdownController::new(gate);
137        let mut orders = OrderRuntime::new();
138        let now = Instant::now();
139        orders.register_place(cid(1), now);
140        orders.register_place(cid(2), now);
141
142        let report = ctl.shutdown_local(ShutdownMode::Immediate, now, &mut orders, || (0, 0));
143
144        assert_eq!(report.inflight_abandoned.len(), 2);
145        assert!(matches!(
146            orders.state(cid(1)),
147            Some(OrderState::Unknown { .. })
148        ));
149    }
150}
brokapi_runtime/shutdown.rs

brokapi_runtime/
shutdown.rs
