1use std::sync::{Arc, Mutex};
22use tokio::sync::watch;
23
24pub 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 pub fn update(&self, f: impl FnOnce(&T) -> T) {
55 self.tx.send_modify(|cur| {
56 *cur = f(cur);
57 });
58 }
59
60 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
69pub 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 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 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}