1use notify_future::Notify;
2use std::collections::{HashMap, VecDeque};
3use std::fmt::{Display, Formatter};
4use std::future::Future;
5use std::hash::Hash;
6use std::pin::Pin;
7use std::sync::{Arc, Mutex};
8use std::task::{Context, Poll};
9
10#[derive(Debug, Eq, PartialEq, Copy, Clone)]
11pub enum WaiterError {
12 AlreadyExist,
13 Timeout,
14 NoWaiter,
15}
16
17impl Display for WaiterError {
18 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
19 match self {
20 WaiterError::AlreadyExist => write!(f, "AlreadyExist"),
21 WaiterError::Timeout => write!(f, "Timeout"),
22 WaiterError::NoWaiter => write!(f, "NoWaiter"),
23 }
24 }
25}
26
27impl std::error::Error for WaiterError {}
28
29pub type WaiterResult<T> = Result<T, WaiterError>;
30
31pub struct ResultFuture<'a, R> {
32 future: Pin<Box<dyn Future<Output = Result<R, WaiterError>> + 'a + Send>>,
33 drop_cleanup: Option<Box<dyn FnOnce() + 'a + Send>>,
34}
35
36impl<'a, R> ResultFuture<'a, R> {
37 pub fn new(future: Pin<Box<dyn Future<Output = Result<R, WaiterError>> + 'a + Send>>) -> Self {
38 Self {
39 future,
40 drop_cleanup: None,
41 }
42 }
43
44 fn new_with_drop_cleanup(
45 future: Pin<Box<dyn Future<Output = Result<R, WaiterError>> + 'a + Send>>,
46 drop_cleanup: impl FnOnce() + 'a + Send,
47 ) -> Self {
48 Self {
49 future,
50 drop_cleanup: Some(Box::new(drop_cleanup)),
51 }
52 }
53}
54
55impl<R> Future for ResultFuture<'_, R> {
56 type Output = Result<R, WaiterError>;
57
58 fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
59 let result = self.future.as_mut().poll(cx);
60 if result.is_ready() {
61 self.drop_cleanup = None;
62 }
63 result
64 }
65}
66
67impl<R> Drop for ResultFuture<'_, R> {
68 fn drop(&mut self) {
69 if let Some(cleanup) = self.drop_cleanup.take() {
70 cleanup();
71 }
72 }
73}
74
75struct CallbackWaiterState<K, R> {
76 result_notifies: HashMap<K, CallbackRegistration<R>>,
77 result_cache: HashMap<K, VecDeque<R>>,
78}
79
80struct CallbackRegistration<R> {
81 identity: Arc<()>,
82 notify: Option<Notify<R>>,
83}
84
85pub struct CallbackWaiter<K, R> {
86 state: Mutex<CallbackWaiterState<K, R>>,
87}
88
89impl<K: Hash + Eq + Clone + 'static + Send, R: 'static + Send> CallbackWaiter<K, R> {
90 pub fn new() -> Self {
91 Self {
92 state: Mutex::new(CallbackWaiterState {
93 result_notifies: HashMap::new(),
94 result_cache: HashMap::new(),
95 }),
96 }
97 }
98
99 pub fn create_result_future(&self, callback_id: K) -> WaiterResult<ResultFuture<'_, R>> {
100 let waiter = {
101 let mut state = self.state.lock().unwrap();
102 if let Some(registration) = state.result_notifies.get(&callback_id) {
103 if registration
104 .notify
105 .as_ref()
106 .is_some_and(|notify| !notify.is_canceled())
107 {
108 return Err(WaiterError::AlreadyExist);
109 }
110 }
111 if let Some(result) = state.result_cache.get_mut(&callback_id) {
112 if let Some(ret) = result.pop_front() {
113 return Ok(ResultFuture::new(Box::pin(async move { Ok(ret) })));
114 }
115 }
116
117 let (notify, waiter) = Notify::new();
118 let identity = Arc::new(());
119 state.result_notifies.insert(
120 callback_id.clone(),
121 CallbackRegistration {
122 identity: identity.clone(),
123 notify: Some(notify),
124 },
125 );
126 (waiter, identity)
127 };
128 let (waiter, identity) = waiter;
129 let cleanup_id = callback_id.clone();
130 let drop_identity = identity.clone();
131
132 Ok(ResultFuture::new_with_drop_cleanup(
133 Box::pin(async move {
134 let ret = waiter.await;
135 self.remove_registration_if_current(&callback_id, &identity);
136 Ok(ret)
137 }),
138 move || {
139 self.remove_registration_if_current(&cleanup_id, &drop_identity);
140 },
141 ))
142 }
143
144 pub fn create_timeout_result_future(
145 &self,
146 callback_id: K,
147 timeout: std::time::Duration,
148 ) -> WaiterResult<ResultFuture<'_, R>> {
149 let waiter = {
150 let mut state = self.state.lock().unwrap();
151 if let Some(registration) = state.result_notifies.get(&callback_id) {
152 if registration
153 .notify
154 .as_ref()
155 .is_some_and(|notify| !notify.is_canceled())
156 {
157 return Err(WaiterError::AlreadyExist);
158 }
159 }
160 if let Some(result) = state.result_cache.get_mut(&callback_id) {
161 if let Some(ret) = result.pop_front() {
162 return Ok(ResultFuture::new(Box::pin(async move { Ok(ret) })));
163 }
164 }
165
166 let (notify, waiter) = Notify::new();
167 let identity = Arc::new(());
168 state.result_notifies.insert(
169 callback_id.clone(),
170 CallbackRegistration {
171 identity: identity.clone(),
172 notify: Some(notify),
173 },
174 );
175 (waiter, identity)
176 };
177 let (waiter, identity) = waiter;
178 let cleanup_id = callback_id.clone();
179 let drop_identity = identity.clone();
180
181 Ok(ResultFuture::new_with_drop_cleanup(
182 Box::pin(async move {
183 let ret = tokio::time::timeout(timeout, waiter).await;
184 self.remove_registration_if_current(&callback_id, &identity);
185 match ret {
186 Ok(ret) => Ok(ret),
187 Err(_) => Err(WaiterError::Timeout),
188 }
189 }),
190 move || {
191 self.remove_registration_if_current(&cleanup_id, &drop_identity);
192 },
193 ))
194 }
195
196 fn remove_registration_if_current(&self, callback_id: &K, identity: &Arc<()>) {
197 let mut state = self.state.lock().unwrap();
198 let is_current = state
199 .result_notifies
200 .get(callback_id)
201 .is_some_and(|registration| Arc::ptr_eq(®istration.identity, identity));
202 if is_current {
203 state.result_notifies.remove(callback_id);
204 }
205 }
206
207 pub fn set_result(&self, callback_id: K, result: R) -> Result<(), WaiterError> {
208 let mut state = self.state.lock().unwrap();
209 if let Some(registration) = state.result_notifies.get_mut(&callback_id) {
210 if let Some(future) = registration.notify.take() {
211 if !future.is_canceled() {
212 future.notify(result);
213 return Ok(());
214 }
215 }
216 }
217 Err(WaiterError::NoWaiter)
218 }
219
220 pub fn set_result_with_cache(&self, callback_id: K, result: R) {
221 let mut state = self.state.lock().unwrap();
222 if let Some(registration) = state.result_notifies.get_mut(&callback_id) {
223 if let Some(future) = registration.notify.take() {
224 if !future.is_canceled() {
225 future.notify(result);
226 return;
227 }
228 }
229 }
230 state
231 .result_cache
232 .entry(callback_id)
233 .or_default()
234 .push_back(result);
235 }
236}
237
238struct SingleCallbackWaiterState<R> {
239 result_notify: Option<SingleCallbackRegistration<R>>,
240 result_cache: VecDeque<R>,
241}
242
243struct SingleCallbackRegistration<R> {
244 identity: Arc<()>,
245 notify: Option<Notify<R>>,
246}
247
248pub struct SingleCallbackWaiter<R> {
249 state: Mutex<SingleCallbackWaiterState<R>>,
250}
251
252impl<R: 'static + Send> SingleCallbackWaiter<R> {
253 pub fn new() -> Self {
254 Self {
255 state: Mutex::new(SingleCallbackWaiterState {
256 result_notify: None,
257 result_cache: VecDeque::new(),
258 }),
259 }
260 }
261
262 pub fn create_result_future(&self) -> WaiterResult<ResultFuture<'_, R>> {
263 let waiter = {
264 let mut state = self.state.lock().unwrap();
265 if let Some(registration) = state.result_notify.as_ref() {
266 if registration
267 .notify
268 .as_ref()
269 .is_some_and(|notify| !notify.is_canceled())
270 {
271 return Err(WaiterError::AlreadyExist);
272 }
273 }
274
275 if let Some(ret) = state.result_cache.pop_front() {
276 return Ok(ResultFuture::new(Box::pin(async move { Ok(ret) })));
277 }
278 let (notify, waiter) = Notify::new();
279 let identity = Arc::new(());
280 state.result_notify = Some(SingleCallbackRegistration {
281 identity: identity.clone(),
282 notify: Some(notify),
283 });
284 (waiter, identity)
285 };
286 let (waiter, identity) = waiter;
287 let drop_identity = identity.clone();
288
289 Ok(ResultFuture::new_with_drop_cleanup(
290 Box::pin(async move {
291 let ret = waiter.await;
292 self.remove_registration_if_current(&identity);
293 Ok(ret)
294 }),
295 move || {
296 self.remove_registration_if_current(&drop_identity);
297 },
298 ))
299 }
300
301 pub fn create_timeout_result_future(
302 &self,
303 timeout: std::time::Duration,
304 ) -> WaiterResult<ResultFuture<'_, R>> {
305 let waiter = {
306 let mut state = self.state.lock().unwrap();
307 if let Some(registration) = state.result_notify.as_ref() {
308 if registration
309 .notify
310 .as_ref()
311 .is_some_and(|notify| !notify.is_canceled())
312 {
313 return Err(WaiterError::AlreadyExist);
314 }
315 }
316
317 if let Some(ret) = state.result_cache.pop_front() {
318 return Ok(ResultFuture::new(Box::pin(async move { Ok(ret) })));
319 }
320
321 let (notify, waiter) = Notify::new();
322 let identity = Arc::new(());
323 state.result_notify = Some(SingleCallbackRegistration {
324 identity: identity.clone(),
325 notify: Some(notify),
326 });
327 (waiter, identity)
328 };
329 let (waiter, identity) = waiter;
330 let drop_identity = identity.clone();
331
332 Ok(ResultFuture::new_with_drop_cleanup(
333 Box::pin(async move {
334 let ret = tokio::time::timeout(timeout, waiter).await;
335 self.remove_registration_if_current(&identity);
336 match ret {
337 Ok(ret) => Ok(ret),
338 Err(_) => Err(WaiterError::Timeout),
339 }
340 }),
341 move || {
342 self.remove_registration_if_current(&drop_identity);
343 },
344 ))
345 }
346
347 fn remove_registration_if_current(&self, identity: &Arc<()>) {
348 let mut state = self.state.lock().unwrap();
349 let is_current = state
350 .result_notify
351 .as_ref()
352 .is_some_and(|registration| Arc::ptr_eq(®istration.identity, identity));
353 if is_current {
354 state.result_notify = None;
355 }
356 }
357
358 pub fn set_result(&self, result: R) -> Result<(), WaiterError> {
359 let mut state = self.state.lock().unwrap();
360 if let Some(registration) = state.result_notify.as_mut() {
361 if let Some(future) = registration.notify.take() {
362 if !future.is_canceled() {
363 future.notify(result);
364 return Ok(());
365 }
366 }
367 }
368 Err(WaiterError::NoWaiter)
369 }
370
371 pub fn set_result_with_cache(&self, result: R) {
372 let mut state = self.state.lock().unwrap();
373 if let Some(registration) = state.result_notify.as_mut() {
374 if let Some(future) = registration.notify.take() {
375 if !future.is_canceled() {
376 future.notify(result);
377 return;
378 }
379 }
380 }
381 state.result_cache.push_back(result);
382 }
383}
384#[cfg(test)]
385mod test {
386 use super::*;
387 use std::sync::Arc;
388 use tokio::time::{sleep, Duration};
389
390 #[tokio::test]
391 async fn test_waiter() {
392 let waiter = Arc::new(CallbackWaiter::new());
393 let callback_id = 1;
394 let result_future = waiter.create_result_future(callback_id).unwrap();
395 assert!(waiter.create_result_future(callback_id).is_err());
396 let tmp = waiter.clone();
397 tokio::spawn(async move {
398 sleep(Duration::from_millis(1000)).await;
399 let ret = tmp.set_result(callback_id, 1);
400 assert!(ret.is_ok());
401 });
402 let ret = result_future.await.unwrap();
403 assert_eq!(ret, 1);
404 }
405
406 #[tokio::test]
407 async fn test_waiter1() {
408 let waiter = Arc::new(CallbackWaiter::new());
409 let callback_id = 1;
410 let tmp = waiter.clone();
411 tokio::spawn(async move {
412 tmp.set_result_with_cache(callback_id, 1);
413 });
414 let result_future = waiter.create_result_future(callback_id).unwrap();
415 let ret = result_future.await.unwrap();
416 assert_eq!(ret, 1);
417 }
418
419 #[tokio::test]
420 async fn test_waiter_timout() {
421 let waiter = Arc::new(CallbackWaiter::new());
422 let callback_id = 1;
423 let result_future = waiter
424 .create_timeout_result_future(callback_id, Duration::from_secs(2))
425 .unwrap();
426 let tmp = waiter.clone();
427 tokio::spawn(async move {
428 sleep(Duration::from_millis(1000)).await;
429 let ret = tmp.set_result(callback_id, 1);
430 assert!(ret.is_ok());
431 });
432 let ret = result_future.await.unwrap();
433 assert_eq!(ret, 1);
434 }
435
436 #[tokio::test]
437 async fn test_waiter_timout2() {
438 let waiter = Arc::new(CallbackWaiter::new());
439 let callback_id = 1;
440 let result_future = waiter
441 .create_timeout_result_future(callback_id, Duration::from_secs(2))
442 .unwrap();
443 let tmp = waiter.clone();
444 tokio::spawn(async move {
445 sleep(Duration::from_secs(3)).await;
446 let ret = tmp.set_result(callback_id, 1);
447 assert!(ret.is_err());
448 });
449 match result_future.await {
450 Ok(_) => {}
451 Err(e) => {
452 assert_eq!(e, WaiterError::Timeout);
453 }
454 }
455 }
456
457 #[tokio::test]
458 async fn test_waiter_timout3() {
459 let waiter = Arc::new(CallbackWaiter::new());
460 let callback_id = 1;
461 let tmp = waiter.clone();
462 tokio::spawn(async move {
463 let ret = tmp.set_result(callback_id, 1);
464 assert!(ret.is_err());
465 })
466 .await
467 .unwrap();
468 let result_future = waiter
469 .create_timeout_result_future(callback_id, Duration::from_secs(2))
470 .unwrap();
471 assert!(waiter
472 .create_timeout_result_future(callback_id, Duration::from_secs(2))
473 .is_err());
474 match result_future.await {
475 Ok(_) => {}
476 Err(e) => {
477 assert_eq!(e, WaiterError::Timeout);
478 }
479 }
480 }
481
482 #[tokio::test]
483 async fn test_signle_waiter() {
484 let waiter = Arc::new(SingleCallbackWaiter::new());
485 let result_future = waiter.create_result_future().unwrap();
486 assert!(waiter.create_result_future().is_err());
487 let tmp = waiter.clone();
488 tokio::spawn(async move {
489 sleep(Duration::from_millis(1000)).await;
490 let ret = tmp.set_result(1);
491 assert!(ret.is_ok());
492 });
493 let ret = result_future.await.unwrap();
494 assert_eq!(ret, 1);
495 }
496
497 #[tokio::test]
498 async fn test_single_waiter1() {
499 let waiter = Arc::new(SingleCallbackWaiter::new());
500 let tmp = waiter.clone();
501 tokio::spawn(async move {
502 tmp.set_result_with_cache(1);
503 });
504 let result_future = waiter.create_result_future().unwrap();
505 let ret = result_future.await.unwrap();
506 assert_eq!(ret, 1);
507 }
508
509 #[tokio::test]
510 async fn test_single_waiter_timout() {
511 let waiter = Arc::new(SingleCallbackWaiter::new());
512 let result_future = waiter
513 .create_timeout_result_future(Duration::from_secs(2))
514 .unwrap();
515 assert!(waiter
516 .create_timeout_result_future(Duration::from_secs(2))
517 .is_err());
518 let tmp = waiter.clone();
519 tokio::spawn(async move {
520 sleep(Duration::from_millis(1000)).await;
521 let ret = tmp.set_result(1);
522 assert!(ret.is_ok());
523 });
524 let ret = result_future.await.unwrap();
525 assert_eq!(ret, 1);
526 }
527
528 #[tokio::test]
529 async fn test_single_waiter_timout2() {
530 let waiter = Arc::new(SingleCallbackWaiter::new());
531 let result_future = waiter
532 .create_timeout_result_future(Duration::from_secs(2))
533 .unwrap();
534 let tmp = waiter.clone();
535 tokio::spawn(async move {
536 sleep(Duration::from_secs(3)).await;
537 let ret = tmp.set_result(1);
538 assert!(ret.is_err());
539 });
540 match result_future.await {
541 Ok(_) => {}
542 Err(e) => {
543 assert_eq!(e, WaiterError::Timeout);
544 }
545 }
546 }
547
548 #[tokio::test]
549 async fn test_single_waiter_timout3() {
550 let waiter = Arc::new(SingleCallbackWaiter::new());
551 let tmp = waiter.clone();
552 tokio::spawn(async move {
553 let ret = tmp.set_result(1);
554 assert!(ret.is_err());
555 })
556 .await
557 .unwrap();
558 let result_future = waiter
559 .create_timeout_result_future(Duration::from_secs(2))
560 .unwrap();
561 match result_future.await {
562 Ok(_) => {}
563 Err(e) => {
564 assert_eq!(e, WaiterError::Timeout);
565 }
566 }
567 }
568
569 #[tokio::test]
570 async fn test_waiter_reregister_after_future_drop() {
571 let waiter = Arc::new(CallbackWaiter::new());
572 let callback_id = 42;
573 let dropped_future = waiter.create_result_future(callback_id).unwrap();
574 drop(dropped_future);
575
576 sleep(Duration::from_millis(10)).await;
577
578 let result_future = waiter.create_result_future(callback_id).unwrap();
579 let tmp = waiter.clone();
580 tokio::spawn(async move {
581 tmp.set_result(callback_id, 7).unwrap();
582 });
583
584 let ret = result_future.await.unwrap();
585 assert_eq!(ret, 7);
586 }
587
588 #[tokio::test]
589 async fn test_waiter_cache_fifo_under_load() {
590 let waiter = CallbackWaiter::new();
591 let callback_id = 1;
592 let total = 200;
593
594 for i in 0..total {
595 waiter.set_result_with_cache(callback_id, i);
596 }
597
598 for expected in 0..total {
599 let ret = waiter
600 .create_result_future(callback_id)
601 .unwrap()
602 .await
603 .unwrap();
604 assert_eq!(ret, expected);
605 }
606 }
607
608 #[tokio::test]
609 async fn test_waiter_timeout_set_result_race() {
610 for callback_id in 0..50 {
611 let waiter = Arc::new(CallbackWaiter::new());
612 let result_future = waiter
613 .create_timeout_result_future(callback_id, Duration::from_millis(50))
614 .unwrap();
615
616 let tmp = waiter.clone();
617 let set_task = tokio::spawn(async move {
618 sleep(Duration::from_millis(50)).await;
619 tmp.set_result(callback_id, 1)
620 });
621
622 let future_result = result_future.await;
623 let set_result = set_task.await.unwrap();
624
625 match (future_result, set_result) {
626 (Ok(1), Ok(())) => {}
627 (Err(WaiterError::Timeout), Err(WaiterError::NoWaiter)) => {}
628 (other_future, other_set) => {
629 panic!(
630 "unexpected race outcome: {:?}, {:?}",
631 other_future, other_set
632 );
633 }
634 }
635 }
636 }
637}