nautilus_network/
writer.rs1use std::sync::Arc;
19
20use crate::error::{NetworkConfigError, NetworkConfigResult, SendError};
21
22pub(crate) const DEFAULT_WRITER_CAPACITY: usize = 1_024;
23
24#[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 pub fn send(&self, command: T) -> Result<(), SendError> {
54 self.send_with_slots(command, &self.slots)
55 }
56
57 pub(crate) fn send_control(&self, command: T) -> Result<(), SendError> {
59 self.send_with_slots(command, &self.slots_control)
60 }
61
62 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}