1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
|
//! An async notification channel.
//!
//! This channel allows one task to notify another. No data is passed from the sender to receiver. A
//! [`NotifySender`] may send multiple notifications and a [`NotifyReceiver`] may receive multiple
//! notifications. Notifications will be coalesced, so if a `NotifySender` sends multiple
//! notifications, the `NotifyReceiver` may or may not receive all of the notifications. If there
//! are multiple `NotifyReceiver`s, each will be notified.
//!
//! An optional type can be attached to the `NotifySender` and `NotifyReceiver` to identify the
//! purpose of the notifications and to provide type checking.
// TODO(arti#534): we expect to use this for flow control, so we should remove this later
#![cfg_attr(not(test), expect(dead_code))]
use std::marker::PhantomData;
use std::pin::Pin;
use std::task::{Context, Poll};
use educe::Educe;
use futures::stream::{Fuse, FusedStream};
use futures::{Stream, StreamExt};
use pin_project::pin_project;
use postage::watch;
/// A [`NotifySender`] which can notify [`NotifyReceiver`]s.
///
/// See the [module documentation](self) for details.
#[derive(Educe)]
#[educe(Debug)]
pub(crate) struct NotifySender<T = ()> {
/// The "sender" we use to implement the async behaviour.
sender: watch::Sender<()>,
/// Allows the user to optionally attach a type marker to identify the purpose of the
/// notifications.
#[educe(Debug(ignore))]
_marker: PhantomData<fn() -> T>,
}
/// A [`NotifyReceiver`] which can receive notifications from a [`NotifySender`].
///
/// See the [module documentation](self) for details.
// We should theoretically be able to impl `Clone`, but `Fuse` does not implement `Clone` so we'd
// have to implement something manually. If we do want `Clone` in the future, be careful about the
// initial state of the new `NotifyReceiver` (see the `try_recv` in `NotifySender::subscribe`).
#[derive(Educe)]
#[educe(Debug)]
#[pin_project]
pub(crate) struct NotifyReceiver<T = ()> {
/// The "receiver" we use to implement the async behaviour.
#[pin]
receiver: Fuse<watch::Receiver<()>>,
/// Allows the user to optionally attach a type marker to identify the purpose of the
/// notifications.
#[educe(Debug(ignore))]
_marker: PhantomData<fn() -> T>,
}
impl NotifySender {
/// Create a new untyped [`NotifySender`].
pub(crate) fn new() -> Self {
Self::new_typed()
}
}
impl<T> NotifySender<T> {
/// Create a new typed [`NotifySender<T>`].
pub(crate) fn new_typed() -> Self {
let (sender, _receiver) = watch::channel();
Self {
sender,
_marker: Default::default(),
}
}
/// Notify all [`NotifyReceiver`]s.
pub(crate) fn notify(&mut self) {
// from `postage::watch::Sender`:
// > Mutably borrows the contained value, blocking the channel while the borrow is held.
// > After the borrow is released, receivers will be notified of a new value.
self.sender.borrow_mut();
}
/// Create a new [`NotifyReceiver`] for this [`NotifySender`].
///
/// A new `NotifyReceiver` will not see any past notifications.
pub(crate) fn subscribe(&mut self) -> NotifyReceiver<T> {
let mut receiver = self.sender.subscribe();
// a `watch::Receiver` will always return the existing status of the `watch::Sender` as the
// first stream item, so we need to recv and discard it so that this `NotifyReceiver` begins
// in the "pending" state
use postage::stream::Stream as PostageStream;
use postage::stream::TryRecvError;
assert_eq!(PostageStream::try_recv(&mut receiver), Ok(()));
assert_eq!(
PostageStream::try_recv(&mut receiver),
Err(TryRecvError::Pending),
);
NotifyReceiver {
receiver: receiver.fuse(),
_marker: Default::default(),
}
}
}
impl<T> Stream for NotifyReceiver<T> {
type Item = ();
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
self.project().receiver.poll_next(cx)
}
}
// the `NotifyReceiver` stores a `Fuse`
impl<T> FusedStream for NotifyReceiver<T> {
fn is_terminated(&self) -> bool {
self.receiver.is_terminated()
}
}
#[cfg(test)]
mod test {
#![allow(clippy::unwrap_used)]
use super::*;
use futures::FutureExt;
#[test]
fn notify() {
tor_rtmock::MockRuntime::test_with_various(|_rt| async move {
let mut sender = NotifySender::new();
let mut receiver = sender.subscribe();
// receivers should initially wait for a notification
assert_eq!(receiver.next().now_or_never(), None);
assert_eq!(receiver.next().now_or_never(), None);
sender.notify();
// we should receive a single notification
assert_eq!(receiver.next().now_or_never(), Some(Some(())));
assert_eq!(receiver.next().now_or_never(), None);
sender.notify();
sender.notify();
sender.notify();
// we should still receive a single notification
assert_eq!(receiver.next().now_or_never(), Some(Some(())));
assert_eq!(receiver.next().now_or_never(), None);
sender.notify();
drop(sender);
// we should see the last notification, and then since we dropped the sender, the stream
// should indicate that it's finished
assert_eq!(receiver.next().now_or_never(), Some(Some(())));
assert_eq!(receiver.next().now_or_never(), Some(None));
assert_eq!(receiver.next().now_or_never(), Some(None));
});
}
#[test]
fn notify_multiple_receivers() {
tor_rtmock::MockRuntime::test_with_various(|_rt| async move {
let mut sender = NotifySender::new();
let mut receiver_1 = sender.subscribe();
let mut receiver_2 = sender.subscribe();
sender.notify();
let mut receiver_3 = sender.subscribe();
// first two receivers should each receive a notification
assert_eq!(receiver_1.next().now_or_never(), Some(Some(())));
assert_eq!(receiver_2.next().now_or_never(), Some(Some(())));
// third receiver should not receive a notification since it was created after the
// notification was generated
assert_eq!(receiver_3.next().now_or_never(), None);
});
}
}
|
