Skip to main content

nautilus_network/
writer.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16//! Admission limits shared by writer queues and reconnect buffers.
17
18use std::sync::Arc;
19
20use crate::error::{NetworkConfigError, NetworkConfigResult, SendError};
21
22pub(crate) const DEFAULT_WRITER_CAPACITY: usize = 1_024;
23
24/// A writer command sender with a shared limit on outstanding messages.
25///
26/// Clones share capacity. A message holds its slot until written or discarded, including
27/// while it waits for reconnect replay.
28#[derive(Debug)]
29pub struct WriterSender<T> {
30    tx: tokio::sync::mpsc::UnboundedSender<(T, tokio::sync::OwnedSemaphorePermit)>,
31    slots: Arc<tokio::sync::Semaphore>,
32    slots_control: Arc<tokio::sync::Semaphore>,
33    slots_update: Arc<tokio::sync::Semaphore>,
34}
35
36impl<T> Clone for WriterSender<T> {
37    fn clone(&self) -> Self {
38        Self {
39            tx: self.tx.clone(),
40            slots: Arc::clone(&self.slots),
41            slots_control: Arc::clone(&self.slots_control),
42            slots_update: Arc::clone(&self.slots_update),
43        }
44    }
45}
46
47impl<T> WriterSender<T> {
48    /// Enqueues a command without waiting for capacity.
49    ///
50    /// # Errors
51    ///
52    /// Returns an error if all slots are occupied, or if the writer has stopped.
53    pub fn send(&self, command: T) -> Result<(), SendError> {
54        self.send_with_slots(command, &self.slots)
55    }
56
57    // Auth recovery and keepalives must not compete with messages held for replay
58    pub(crate) fn send_control(&self, command: T) -> Result<(), SendError> {
59        self.send_with_slots(command, &self.slots_control)
60    }
61
62    // Reconnect timeouts can abandon an update before the writer processes it,
63    // reserve one independent slot so retries stay bounded even when replay is full.
64    pub(crate) fn send_update(&self, command: T) -> Result<(), SendError> {
65        self.send_with_slots(command, &self.slots_update)
66    }
67
68    fn send_with_slots(
69        &self,
70        command: T,
71        slots: &Arc<tokio::sync::Semaphore>,
72    ) -> Result<(), SendError> {
73        if self.tx.is_closed() {
74            return Err(SendError::BrokenPipe("writer channel closed".to_string()));
75        }
76
77        let permit = Arc::clone(slots)
78            .try_acquire_owned()
79            .map_err(|_| SendError::BufferFull)?;
80        self.tx
81            .send((command, permit))
82            .map_err(|_| SendError::BrokenPipe("writer channel closed".to_string()))
83    }
84}
85
86pub(crate) type WriterReceiver<T> =
87    tokio::sync::mpsc::UnboundedReceiver<(T, tokio::sync::OwnedSemaphorePermit)>;
88
89pub(crate) fn channel<T>(capacity: usize) -> (WriterSender<T>, WriterReceiver<T>) {
90    let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
91    (
92        WriterSender {
93            tx,
94            slots: Arc::new(tokio::sync::Semaphore::new(capacity)),
95            slots_control: Arc::new(tokio::sync::Semaphore::new(capacity)),
96            slots_update: Arc::new(tokio::sync::Semaphore::new(1)),
97        },
98        rx,
99    )
100}
101
102pub(crate) fn validate_capacity(capacity: Option<usize>) -> NetworkConfigResult<()> {
103    if let Some(capacity) = capacity
104        && (capacity == 0 || capacity > tokio::sync::Semaphore::MAX_PERMITS)
105    {
106        return Err(NetworkConfigError::invalid(
107            "writer_capacity",
108            "must be positive and no greater than Semaphore::MAX_PERMITS",
109        ));
110    }
111
112    Ok(())
113}
114
115#[cfg(test)]
116mod tests {
117    use rstest::rstest;
118
119    use super::*;
120
121    #[rstest]
122    #[tokio::test]
123    async fn test_capacity_follows_message_until_discarded() {
124        let (sender, mut receiver) = channel(1);
125        sender.send("accepted").unwrap();
126        let (message, permit) = receiver.recv().await.unwrap();
127        let clone = sender.clone();
128
129        assert_eq!(message, "accepted");
130        assert!(matches!(clone.send("overflow"), Err(SendError::BufferFull)));
131        sender.send_control("authentication").unwrap();
132        assert!(matches!(
133            clone.send_control("overflow"),
134            Err(SendError::BufferFull)
135        ));
136        let (control, control_permit) = receiver.recv().await.unwrap();
137        assert_eq!(control, "authentication");
138        assert!(matches!(
139            clone.send_control("overflow"),
140            Err(SendError::BufferFull)
141        ));
142        drop(control_permit);
143        sender.send_update("replacement").unwrap();
144        let (update, update_permit) = receiver.recv().await.unwrap();
145        assert_eq!(update, "replacement");
146        assert!(matches!(
147            sender.send_update("overflow"),
148            Err(SendError::BufferFull)
149        ));
150        drop(update_permit);
151        sender.send_update("retry").unwrap();
152        let (retry, _) = receiver.recv().await.unwrap();
153        assert_eq!(retry, "retry");
154        assert!(matches!(
155            sender.send("overflow"),
156            Err(SendError::BufferFull)
157        ));
158
159        drop(permit);
160        clone.send("after discard").unwrap();
161        let (message, _) = receiver.recv().await.unwrap();
162        assert_eq!(message, "after discard");
163
164        drop(receiver);
165        assert!(matches!(
166            sender.send("closed"),
167            Err(SendError::BrokenPipe(_))
168        ));
169    }
170}