rocketmq_rust/
shutdown.rs

1/*
2 * Licensed to the Apache Software Foundation (ASF) under one or more
3 * contributor license agreements.  See the NOTICE file distributed with
4 * this work for additional information regarding copyright ownership.
5 * The ASF licenses this file to You under the Apache License, Version 2.0
6 * (the "License"); you may not use this file except in compliance with
7 * the License.  You may obtain a copy of the License at
8 *
9 *     http://www.apache.org/licenses/LICENSE-2.0
10 *
11 * Unless required by applicable law or agreed to in writing, software
12 * distributed under the License is distributed on an "AS IS" BASIS,
13 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 * See the License for the specific language governing permissions and
15 * limitations under the License.
16 */
17use tokio::sync::broadcast;
18use tracing::warn;
19
20pub struct Shutdown<T> {
21    /// `true` if the shutdown signal has been received
22    is_shutdown: bool,
23
24    /// The receiver half of the channel used to listen for shutdown.
25    notify: broadcast::Receiver<T>,
26}
27
28impl<T> Shutdown<T>
29where
30    T: Clone,
31{
32    /// Create a new `Shutdown` backed by the given `broadcast::Receiver`.
33    pub fn new(capacity: usize) -> (Shutdown<T>, broadcast::Sender<T>) {
34        let (tx, _) = broadcast::channel(capacity);
35        let shutdown = Shutdown {
36            is_shutdown: false,
37            notify: tx.subscribe(),
38        };
39        (shutdown, tx)
40    }
41
42    /// Returns `true` if the shutdown signal has been received.
43    pub fn is_shutdown(&self) -> bool {
44        self.is_shutdown
45    }
46
47    /// Receive the shutdown notice, waiting if necessary.
48    pub async fn recv(&mut self) {
49        // If the shutdown signal has already been received, then return
50        // immediately.
51        if self.is_shutdown {
52            return;
53        }
54
55        // Cannot receive a "lag error" as only one value is ever sent.
56        let result = self.notify.recv().await;
57        if result.is_err() {
58            warn!("Failed to receive shutdown signal");
59        }
60
61        // Remember that the signal has been received.
62        self.is_shutdown = true;
63    }
64}
65
66#[cfg(test)]
67mod tests {
68    use super::*;
69
70    #[tokio::test]
71    async fn shutdown_signal_received() {
72        let (mut shutdown, sender) = Shutdown::new(1);
73        sender.send(()).unwrap();
74        shutdown.recv().await;
75        assert!(shutdown.is_shutdown());
76    }
77
78    #[tokio::test]
79    async fn shutdown_signal_not_received() {
80        let (shutdown, _) = Shutdown::<()>::new(1);
81        assert!(!shutdown.is_shutdown());
82    }
83
84    #[tokio::test]
85    async fn shutdown_signal_multiple_receivers() {
86        let (mut shutdown1, sender) = Shutdown::new(1);
87        sender.send(()).unwrap();
88        shutdown1.recv().await;
89
90        assert!(shutdown1.is_shutdown());
91    }
92
93    #[tokio::test]
94    async fn shutdown_signal_already_received() {
95        let (mut shutdown, sender) = Shutdown::new(1);
96        sender.send(()).unwrap();
97        shutdown.recv().await;
98        shutdown.recv().await; // Call recv again
99        assert!(shutdown.is_shutdown());
100    }
101}
rocketmq_rust/shutdown.rs

rocketmq_rust/
shutdown.rs
