thedes_async_util/non_blocking/spsc/
bounded.rs

1use std::{
2    cell::{Cell, UnsafeCell},
3    fmt,
4    mem::MaybeUninit,
5    ops::Deref,
6    ptr::NonNull,
7    sync::atomic::{AtomicBool, AtomicUsize, Ordering::*},
8};
9
10use thiserror::Error;
11
12pub fn channel<T>(buf_size: usize) -> (Sender<T>, Receiver<T>) {
13    if buf_size == 0 {
14        panic!(
15            "non-blocking SPSC bounded channel cannot have zero-sized buffer"
16        );
17    }
18    let (sender_shared, receiver_shared) = SharedPtr::new(buf_size);
19    (Sender::new(sender_shared), Receiver::new(receiver_shared))
20}
21
22#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)]
23pub enum SendErrorKind {
24    #[error("Buffer is full")]
25    Overflow,
26    #[error("Receiver disconnected")]
27    Disconnected,
28}
29
30#[derive(Debug, Clone, Error)]
31#[error("Receiver disconnected")]
32pub struct SendError<T> {
33    #[source]
34    kind: SendErrorKind,
35    message: T,
36}
37
38impl<T> SendError<T> {
39    fn new(kind: SendErrorKind, message: T) -> Self {
40        Self { kind, message }
41    }
42
43    pub fn kind(&self) -> SendErrorKind {
44        self.kind
45    }
46
47    pub fn message(&self) -> &T {
48        &self.message
49    }
50
51    pub fn into_message(self) -> T {
52        self.message
53    }
54}
55
56#[derive(Debug, Clone, Error)]
57#[error("Senders disconnected")]
58pub struct RecvError {
59    _private: (),
60}
61
62impl RecvError {
63    fn new() -> Self {
64        Self { _private: () }
65    }
66}
67
68#[repr(C)]
69struct Shared<T> {
70    connected: AtomicBool,
71    unread: AtomicUsize,
72    front: Cell<usize>,
73    back: Cell<usize>,
74    buf: Box<[UnsafeCell<MaybeUninit<T>>]>,
75}
76
77impl<T> fmt::Debug for Shared<T> {
78    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
79        f.debug_struct("Shared")
80            .field("connected", &self.connected)
81            .field("unread", &self.unread)
82            .field("buf", &self.buf)
83            .finish()
84    }
85}
86
87impl<T> Shared<T> {
88    pub fn new_connected(buf_size: usize) -> Self {
89        Self {
90            connected: AtomicBool::new(true),
91            unread: AtomicUsize::new(0),
92            front: Cell::new(0),
93            back: Cell::new(0),
94            buf: (0 .. buf_size)
95                .map(|_| UnsafeCell::new(MaybeUninit::uninit()))
96                .collect(),
97        }
98    }
99
100    pub fn is_connected_weak(&self) -> bool {
101        self.connected.load(Relaxed)
102    }
103
104    pub unsafe fn send(&self, message: T) -> Result<(), SendError<T>> {
105        let unread = self.unread.load(Acquire);
106        let connected = self.connected.load(Acquire);
107        if !connected {
108            return Err(SendError::new(SendErrorKind::Disconnected, message));
109        }
110        if unread == self.buf.len() {
111            return Err(SendError::new(SendErrorKind::Overflow, message));
112        }
113        unsafe {
114            let back = self.back.get();
115            (*self.buf[back].get()).write(message);
116            self.back.set((back + 1) % self.buf.len());
117            self.unread.fetch_add(1, Release);
118        }
119        Ok(())
120    }
121
122    pub unsafe fn recv_one(&self) -> Result<Option<T>, RecvError> {
123        let unread = self.unread.load(Acquire);
124        if unread == 0 {
125            if self.connected.load(Acquire) {
126                return Ok(None);
127            }
128            return Err(RecvError::new());
129        }
130        unsafe {
131            let front = self.front.get();
132            let message = (*self.buf[front].get()).as_ptr().read();
133            self.front.set((front + 1) % self.buf.len());
134            self.unread.fetch_sub(1, Release);
135            Ok(Some(message))
136        }
137    }
138
139    pub unsafe fn recv_many<'a>(
140        &'a self,
141    ) -> Result<RecvMany<'a, T>, RecvError> {
142        let unread = self.unread.load(Acquire);
143        if unread == 0 && !self.connected.load(Acquire) {
144            Err(RecvError::new())?
145        }
146        Ok(RecvMany { shared: self, count: unread })
147    }
148}
149
150impl<T> Drop for Shared<T> {
151    fn drop(&mut self) {
152        unsafe { while let Ok(Some(_)) = self.recv_one() {} }
153    }
154}
155
156struct SharedPtr<T> {
157    inner: NonNull<Shared<T>>,
158}
159
160impl<T> SharedPtr<T> {
161    pub fn new(buf_size: usize) -> (Self, Self) {
162        let shared = Shared::new_connected(buf_size);
163        let shared_boxed = Box::new(shared);
164
165        let inner =
166            unsafe { NonNull::new_unchecked(Box::into_raw(shared_boxed)) };
167
168        (Self { inner }, Self { inner })
169    }
170}
171
172impl<T> fmt::Debug for SharedPtr<T> {
173    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
174        f.debug_struct("SharedPtr").field("inner", &**self).finish()
175    }
176}
177
178impl<T> Deref for SharedPtr<T> {
179    type Target = Shared<T>;
180
181    fn deref(&self) -> &Self::Target {
182        unsafe { self.inner.as_ref() }
183    }
184}
185
186impl<T> Drop for SharedPtr<T> {
187    fn drop(&mut self) {
188        unsafe {
189            let last = self
190                .connected
191                .compare_exchange(true, false, AcqRel, Acquire)
192                .is_err();
193            if last {
194                let _ = Box::from_raw(self.inner.as_ptr());
195            }
196        }
197    }
198}
199
200#[derive(Debug)]
201pub struct RecvMany<'a, T> {
202    shared: &'a Shared<T>,
203    count: usize,
204}
205
206impl<'a, T> Iterator for RecvMany<'a, T> {
207    type Item = T;
208
209    fn next(&mut self) -> Option<Self::Item> {
210        if self.count == 0 {
211            None?
212        }
213        self.count -= 1;
214        unsafe { self.shared.recv_one().ok().flatten() }
215    }
216}
217
218#[derive(Debug)]
219pub struct Sender<T> {
220    shared: SharedPtr<T>,
221}
222
223impl<T> Sender<T> {
224    fn new(shared: SharedPtr<T>) -> Self {
225        Self { shared }
226    }
227
228    pub fn is_connected(&self) -> bool {
229        self.shared.is_connected_weak()
230    }
231
232    pub fn send(&mut self, message: T) -> Result<(), SendError<T>> {
233        unsafe { self.shared.send(message) }
234    }
235}
236
237unsafe impl<T> Send for Sender<T> where T: Send {}
238unsafe impl<T> Sync for Sender<T> where T: Send {}
239
240#[derive(Debug)]
241pub struct Receiver<T> {
242    shared: SharedPtr<T>,
243}
244
245impl<T> Receiver<T> {
246    fn new(shared: SharedPtr<T>) -> Self {
247        Self { shared }
248    }
249
250    pub fn is_connected(&self) -> bool {
251        self.shared.is_connected_weak()
252    }
253
254    pub fn recv_one(&mut self) -> Result<Option<T>, RecvError> {
255        unsafe { self.shared.recv_one() }
256    }
257
258    pub fn recv_many<'a>(&'a mut self) -> Result<RecvMany<'a, T>, RecvError> {
259        unsafe { self.shared.recv_many() }
260    }
261}
262
263unsafe impl<T> Send for Receiver<T> where T: Send {}
264unsafe impl<T> Sync for Receiver<T> where T: Send {}
265
266#[cfg(test)]
267mod test {
268    use crate::non_blocking::spsc::bounded::SendErrorKind;
269
270    use super::channel;
271
272    #[test]
273    #[should_panic]
274    fn zero_sized_buf_panics() {
275        channel::<u64>(0);
276    }
277
278    #[test]
279    fn recv_empty() {
280        let (_sender, mut receiver) = channel::<u64>(3);
281        assert_eq!(receiver.recv_one().unwrap(), None);
282    }
283
284    #[test]
285    fn simple_send_recv_ok() {
286        let (mut sender, mut receiver) = channel::<u64>(3);
287        sender.send(1).unwrap();
288        sender.send(2).unwrap();
289        sender.send(3).unwrap();
290        assert_eq!(receiver.recv_one().unwrap().unwrap(), 1);
291        assert_eq!(receiver.recv_one().unwrap().unwrap(), 2);
292        assert_eq!(receiver.recv_one().unwrap().unwrap(), 3);
293    }
294
295    #[test]
296    fn simple_send_recv_ok_interleaved() {
297        let (mut sender, mut receiver) = channel::<u64>(3);
298        sender.send(1).unwrap();
299        assert_eq!(receiver.recv_one().unwrap().unwrap(), 1);
300        sender.send(2).unwrap();
301        assert_eq!(receiver.recv_one().unwrap().unwrap(), 2);
302        sender.send(3).unwrap();
303        assert_eq!(receiver.recv_one().unwrap().unwrap(), 3);
304    }
305
306    #[test]
307    fn send_recv_ok_interleaved_wrap_around() {
308        let (mut sender, mut receiver) = channel::<u64>(3);
309        sender.send(1).unwrap();
310        assert_eq!(receiver.recv_one().unwrap().unwrap(), 1);
311        sender.send(2).unwrap();
312        assert_eq!(receiver.recv_one().unwrap().unwrap(), 2);
313        sender.send(3).unwrap();
314        assert_eq!(receiver.recv_one().unwrap().unwrap(), 3);
315        sender.send(4).unwrap();
316        assert_eq!(receiver.recv_one().unwrap().unwrap(), 4);
317        sender.send(5).unwrap();
318        assert_eq!(receiver.recv_one().unwrap().unwrap(), 5);
319    }
320
321    #[test]
322    fn send_overflow() {
323        let (mut sender, mut receiver) = channel::<u64>(3);
324        sender.send(1).unwrap();
325        sender.send(2).unwrap();
326        sender.send(3).unwrap();
327        assert_eq!(sender.send(4).unwrap_err().kind(), SendErrorKind::Overflow);
328        assert_eq!(receiver.recv_one().unwrap().unwrap(), 1);
329        assert_eq!(receiver.recv_one().unwrap().unwrap(), 2);
330        assert_eq!(receiver.recv_one().unwrap().unwrap(), 3);
331        assert_eq!(receiver.recv_one().unwrap(), None);
332        sender.send(5).unwrap();
333        assert_eq!(receiver.recv_one().unwrap().unwrap(), 5);
334    }
335
336    #[test]
337    fn sender_disconnected() {
338        let (_, mut receiver) = channel::<u64>(3);
339        receiver.recv_one().unwrap_err();
340    }
341
342    #[test]
343    fn receiver_disconnected() {
344        let (mut sender, _) = channel::<u64>(3);
345        assert_eq!(
346            sender.send(4).unwrap_err().kind(),
347            SendErrorKind::Disconnected
348        );
349    }
350
351    #[test]
352    fn recv_many() {
353        let (mut sender, mut receiver) = channel::<u64>(6);
354        sender.send(1).unwrap();
355        sender.send(2).unwrap();
356        sender.send(3).unwrap();
357
358        let recv_many = receiver.recv_many().unwrap();
359
360        sender.send(4).unwrap();
361
362        let messages: Vec<_> = recv_many.collect();
363
364        assert_eq!(messages, vec![1, 2, 3]);
365    }
366
367    #[test]
368    fn recv_many_but_sender_dropped() {
369        let (mut sender, mut receiver) = channel::<u64>(3);
370        sender.send(1).unwrap();
371        assert_eq!(receiver.recv_one().unwrap().unwrap(), 1);
372        drop(sender);
373        receiver.recv_many().unwrap_err();
374    }
375}
thedes_async_util/non_blocking/spsc/bounded.rs

thedes_async_util/non_blocking/spsc/
bounded.rs
