Skip to main content

sova_core/
share.rs

1//! Lightweight handles for sharing values across tasks / handlers.
2//!
3//! Inspired by the idea of thread-share (convenient cross-task handles), but only
4//! two primitives — no worker managers or macros.
5//!
6//! | Type | Use for | Notes |
7//! |------|---------|--------|
8//! | [`Cell`] | `Clone` data (counters, flags, config) | `get` / `set` / `update` / [`Cell::changed`] |
9//! | [`Slot`] | ownership handoff (sockets, streams, …) | `put` / `take` / `try_take` |
10//!
11//! Both are cheap [`Clone`] handles (`Arc` inside). Typical wiring:
12//!
13//! ```ignore
14//! let inbox = Slot::<TcpStream>::new();
15//! let n = Cell::new(0u64);
16//! app.state(inbox.clone()).state(n);
17//! // BackgroundService: inbox.put(stream);
18//! // Handler: let s = req.state::<Slot<TcpStream>>().take().await;
19//! ```
20
21use std::sync::{Arc, Mutex};
22use tokio::sync::watch;
23
24/// Shared `Clone` value with change notifications (`tokio::sync::watch` under the hood).
25///
26/// Prefer this for counters and config. For non-[`Clone`] values (sockets), use [`Slot`].
27pub struct Cell<T: Clone> {
28    tx: watch::Sender<T>,
29}
30
31impl<T: Clone> Clone for Cell<T> {
32    fn clone(&self) -> Self {
33        Self {
34            tx: self.tx.clone(),
35        }
36    }
37}
38
39impl<T: Clone> Cell<T> {
40    pub fn new(init: T) -> Self {
41        let (tx, _) = watch::channel(init);
42        Self { tx }
43    }
44
45    pub fn get(&self) -> T {
46        self.tx.borrow().clone()
47    }
48
49    pub fn set(&self, value: T) {
50        self.tx.send_modify(|cur| *cur = value);
51    }
52
53    /// Replace the value with `f(&current)`.
54    pub fn update(&self, f: impl FnOnce(&T) -> T) {
55        self.tx.send_modify(|cur| {
56            *cur = f(cur);
57        });
58    }
59
60    /// Wait until the value changes from the moment this future starts, then return it.
61    pub async fn changed(&self) -> T {
62        let mut rx = self.tx.subscribe();
63        let _ = rx.changed().await;
64        let value = rx.borrow().clone();
65        value
66    }
67}
68
69/// One-item ownership handoff between tasks (sockets, streams, anything non-`Clone`).
70///
71/// [`Slot::put`] stores a value. If the slot was already full, the previous value is
72/// **dropped** (no queue). [`Slot::take`] waits until a value is available.
73pub struct Slot<T> {
74    inner: Arc<SlotInner<T>>,
75}
76
77struct SlotInner<T> {
78    slot: Mutex<Option<T>>,
79    tick: watch::Sender<u64>,
80}
81
82impl<T> Clone for Slot<T> {
83    fn clone(&self) -> Self {
84        Self {
85            inner: Arc::clone(&self.inner),
86        }
87    }
88}
89
90impl<T> Slot<T> {
91    pub fn new() -> Self {
92        let (tick, _) = watch::channel(0u64);
93        Self {
94            inner: Arc::new(SlotInner {
95                slot: Mutex::new(None),
96                tick,
97            }),
98        }
99    }
100
101    /// Store `value`. Replaces (and drops) any unread previous value.
102    pub fn put(&self, value: T) {
103        {
104            let mut g = self.inner.slot.lock().unwrap();
105            *g = Some(value);
106        }
107        self.inner.tick.send_modify(|n| *n = n.wrapping_add(1));
108    }
109
110    pub fn try_take(&self) -> Option<T> {
111        self.inner.slot.lock().unwrap().take()
112    }
113
114    /// Wait until a value is available, then take it.
115    pub async fn take(&self) -> T {
116        let mut rx = self.inner.tick.subscribe();
117        loop {
118            if let Some(v) = self.try_take() {
119                return v;
120            }
121            if rx.changed().await.is_err() {
122                if let Some(v) = self.try_take() {
123                    return v;
124                }
125                std::future::pending::<()>().await;
126            }
127        }
128    }
129}
130
131impl<T> Default for Slot<T> {
132    fn default() -> Self {
133        Self::new()
134    }
135}
136
137#[cfg(test)]
138mod tests {
139    use super::*;
140    use std::time::Duration;
141
142    #[tokio::test]
143    async fn cell_get_set_update() {
144        let c = Cell::new(1u32);
145        assert_eq!(c.get(), 1);
146        c.set(2);
147        assert_eq!(c.get(), 2);
148        c.update(|n| n + 3);
149        assert_eq!(c.get(), 5);
150    }
151
152    #[tokio::test]
153    async fn cell_changed_wakes() {
154        let c = Cell::new(0u32);
155        let c2 = c.clone();
156        let h = tokio::spawn(async move { c2.changed().await });
157        tokio::time::sleep(Duration::from_millis(20)).await;
158        c.set(7);
159        assert_eq!(h.await.unwrap(), 7);
160    }
161
162    #[tokio::test]
163    async fn slot_put_take() {
164        let s = Slot::new();
165        s.put(String::from("hi"));
166        assert_eq!(s.try_take().as_deref(), Some("hi"));
167        assert!(s.try_take().is_none());
168    }
169
170    #[tokio::test]
171    async fn slot_take_waits() {
172        let s = Slot::new();
173        let s2 = s.clone();
174        let h = tokio::spawn(async move { s2.take().await });
175        tokio::time::sleep(Duration::from_millis(20)).await;
176        s.put(42u8);
177        assert_eq!(h.await.unwrap(), 42);
178    }
179
180    #[tokio::test]
181    async fn slot_put_replaces() {
182        let s = Slot::new();
183        s.put(1u8);
184        s.put(2u8);
185        assert_eq!(s.try_take(), Some(2));
186    }
187}