1mod error;
2
3pub use error::*;
4use parking_lot::{RwLock, RwLockReadGuard, RwLockUpgradableReadGuard};
5use std::{
6 collections::BTreeMap,
7 mem,
8 pin::Pin,
9 sync::Arc,
10 task::{Context, Poll, Waker},
11};
12
13pub fn channel<T>(value: T) -> (Sender<T>, Receiver<T>) {
14 let state = Arc::new(RwLock::new(State {
15 value,
16 wakers: BTreeMap::new(),
17 next_waker_id: WakerId::default(),
18 version: 0,
19 closed: false,
20 }));
21
22 (
23 Sender {
24 state: state.clone(),
25 },
26 Receiver { state, version: 0 },
27 )
28}
29
30#[derive(Default, Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
31struct WakerId(usize);
32
33impl WakerId {
34 fn post_inc(&mut self) -> Self {
35 let id = *self;
36 self.0 = id.0.wrapping_add(1);
37 *self
38 }
39}
40
41struct State<T> {
42 value: T,
43 wakers: BTreeMap<WakerId, Waker>,
44 next_waker_id: WakerId,
45 version: usize,
46 closed: bool,
47}
48
49pub struct Sender<T> {
50 state: Arc<RwLock<State<T>>>,
51}
52
53impl<T> Sender<T> {
54 pub fn receiver(&self) -> Receiver<T> {
55 let version = self.state.read().version;
56 Receiver {
57 state: self.state.clone(),
58 version,
59 }
60 }
61
62 pub fn send(&mut self, value: T) -> Result<(), NoReceiverError> {
63 if let Some(state) = Arc::get_mut(&mut self.state) {
64 let state = state.get_mut();
65 state.value = value;
66 debug_assert_eq!(state.wakers.len(), 0);
67 Err(NoReceiverError)
68 } else {
69 let mut state = self.state.write();
70 state.value = value;
71 state.version = state.version.wrapping_add(1);
72 let wakers = mem::take(&mut state.wakers);
73 drop(state);
74
75 for (_, waker) in wakers {
76 waker.wake();
77 }
78
79 Ok(())
80 }
81 }
82}
83
84impl<T> Drop for Sender<T> {
85 fn drop(&mut self) {
86 let mut state = self.state.write();
87 state.closed = true;
88 for (_, waker) in mem::take(&mut state.wakers) {
89 waker.wake();
90 }
91 }
92}
93
94#[derive(Clone)]
95pub struct Receiver<T> {
96 state: Arc<RwLock<State<T>>>,
97 version: usize,
98}
99
100struct Changed<'a, T> {
101 receiver: &'a mut Receiver<T>,
102 pending_waker_id: Option<WakerId>,
103}
104
105impl<T> Future for Changed<'_, T> {
106 type Output = Result<(), NoSenderError>;
107
108 fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
109 let this = &mut *self;
110
111 let state = this.receiver.state.upgradable_read();
112 if state.version != this.receiver.version {
113 // The sender produced a new value. Avoid unregistering the pending
114 // waker, because the sender has already done so.
115 this.pending_waker_id = None;
116 this.receiver.version = state.version;
117 Poll::Ready(Ok(()))
118 } else if state.closed {
119 Poll::Ready(Err(NoSenderError))
120 } else {
121 let mut state = RwLockUpgradableReadGuard::upgrade(state);
122
123 // Unregister the pending waker. This should happen automatically
124 // when the waker gets awoken by the sender, but if this future was
125 // polled again without an explicit call to `wake` (e.g., a spurious
126 // wake by the executor), we need to remove it manually.
127 if let Some(pending_waker_id) = this.pending_waker_id.take() {
128 state.wakers.remove(&pending_waker_id);
129 }
130
131 // Register the waker for this future.
132 let waker_id = state.next_waker_id.post_inc();
133 state.wakers.insert(waker_id, cx.waker().clone());
134 this.pending_waker_id = Some(waker_id);
135
136 Poll::Pending
137 }
138 }
139}
140
141impl<T> Drop for Changed<'_, T> {
142 fn drop(&mut self) {
143 // If this future gets dropped before the waker has a chance of being
144 // awoken, we need to clear it to avoid a memory leak.
145 if let Some(waker_id) = self.pending_waker_id {
146 let mut state = self.receiver.state.write();
147 state.wakers.remove(&waker_id);
148 }
149 }
150}
151
152impl<T> Receiver<T> {
153 pub fn borrow(&mut self) -> parking_lot::MappedRwLockReadGuard<'_, T> {
154 let state = self.state.read();
155 self.version = state.version;
156 RwLockReadGuard::map(state, |state| &state.value)
157 }
158
159 pub fn changed(&mut self) -> impl Future<Output = Result<(), NoSenderError>> {
160 Changed {
161 receiver: self,
162 pending_waker_id: None,
163 }
164 }
165}
166
167impl<T: Clone> Receiver<T> {
168 pub async fn recv(&mut self) -> Result<T, NoSenderError> {
169 self.changed().await?;
170 Ok(self.borrow().clone())
171 }
172}
173
174#[cfg(test)]
175mod tests {
176 use super::*;
177 use futures::{FutureExt, select_biased};
178 use gpui::{AppContext, TestAppContext};
179 use std::{
180 pin::pin,
181 sync::atomic::{AtomicBool, AtomicUsize, Ordering::SeqCst},
182 };
183
184 #[gpui::test]
185 async fn test_basic_watch() {
186 let (mut sender, mut receiver) = channel(0);
187 assert_eq!(sender.send(1), Ok(()));
188 assert_eq!(receiver.recv().await, Ok(1));
189
190 assert_eq!(sender.send(2), Ok(()));
191 assert_eq!(sender.send(3), Ok(()));
192 assert_eq!(receiver.recv().await, Ok(3));
193
194 drop(receiver);
195 assert_eq!(sender.send(4), Err(NoReceiverError));
196
197 let mut receiver = sender.receiver();
198 assert_eq!(sender.send(5), Ok(()));
199 assert_eq!(receiver.recv().await, Ok(5));
200
201 // Ensure `changed` doesn't resolve if we just read the latest value
202 // using `borrow`.
203 assert_eq!(sender.send(6), Ok(()));
204 assert_eq!(*receiver.borrow(), 6);
205 assert_eq!(receiver.changed().now_or_never(), None);
206
207 assert_eq!(sender.send(7), Ok(()));
208 drop(sender);
209 assert_eq!(receiver.recv().await, Ok(7));
210 assert_eq!(receiver.recv().await, Err(NoSenderError));
211 }
212
213 #[gpui::test(iterations = 1000)]
214 async fn test_watch_random(cx: &mut TestAppContext) {
215 let next_id = Arc::new(AtomicUsize::new(1));
216 let closed = Arc::new(AtomicBool::new(false));
217 let (mut tx, rx) = channel(0);
218 let mut tasks = Vec::new();
219
220 tasks.push(cx.background_spawn({
221 let executor = cx.executor();
222 let next_id = next_id.clone();
223 let closed = closed.clone();
224 async move {
225 for _ in 0..16 {
226 executor.simulate_random_delay().await;
227 let id = next_id.fetch_add(1, SeqCst);
228 zlog::info!("sending {}", id);
229 tx.send(id).ok();
230 }
231 closed.store(true, SeqCst);
232 }
233 }));
234
235 for receiver_id in 0..16 {
236 let executor = cx.executor().clone();
237 let next_id = next_id.clone();
238 let closed = closed.clone();
239 let mut rx = rx.clone();
240 let mut prev_observed_value = *rx.borrow();
241 tasks.push(cx.background_spawn(async move {
242 for _ in 0..16 {
243 executor.simulate_random_delay().await;
244
245 zlog::info!("{}: receiving", receiver_id);
246 let mut timeout = executor.simulate_random_delay().fuse();
247 let mut recv = pin!(rx.recv().fuse());
248 select_biased! {
249 _ = timeout => {
250 zlog::info!("{}: dropping recv future", receiver_id);
251 }
252 result = recv => {
253 match result {
254 Ok(value) => {
255 zlog::info!("{}: received {}", receiver_id, value);
256 assert_eq!(value, next_id.load(SeqCst) - 1);
257 assert_ne!(value, prev_observed_value);
258 prev_observed_value = value;
259 }
260 Err(NoSenderError) => {
261 zlog::info!("{}: closed", receiver_id);
262 assert!(closed.load(SeqCst));
263 break;
264 }
265 }
266 }
267 }
268 }
269 }));
270 }
271
272 futures::future::join_all(tasks).await;
273 }
274
275 #[ctor::ctor]
276 fn init_logger() {
277 zlog::init_test();
278 }
279}