1use crate::{Fetch, Resolver, TargetedResolver, ingress};
2use commonware_actor::{Feedback, mailbox::Sender};
3use commonware_cryptography::PublicKey;
4use commonware_utils::{Span, vec::NonEmptyVec};
5
6pub type FetchKey<K, P, S> = ingress::FetchKey<K, S, Option<NonEmptyVec<P>>>;
8
9pub type Message<K, P, S> = ingress::Message<K, S, Option<NonEmptyVec<P>>>;
11
12fn fetch_key<K, P, S>(fetch: Fetch<K, S>, targets: Option<NonEmptyVec<P>>) -> FetchKey<K, P, S> {
13 FetchKey {
14 key: fetch.key,
15 subscribers: NonEmptyVec::new((fetch.subscriber, fetch.span)),
16 metadata: targets,
17 }
18}
19
20#[derive(Clone)]
22pub struct Mailbox<K: Span, P: Eq, S: Eq = ()> {
23 sender: Sender<Message<K, P, S>>,
25}
26
27impl<K: Span, P: Eq, S: Eq> Mailbox<K, P, S> {
28 pub(super) const fn new(sender: Sender<Message<K, P, S>>) -> Self {
30 Self { sender }
31 }
32}
33
34impl<K, P, S> Resolver for Mailbox<K, P, S>
35where
36 K: Span,
37 P: PublicKey,
38 S: Clone + Eq + Send + 'static,
39{
40 type Key = K;
41 type Subscriber = S;
42
43 fn fetch<D>(&mut self, key: D) -> Feedback
50 where
51 D: Into<Fetch<Self::Key, Self::Subscriber>> + Send,
52 {
53 let Fetch {
54 key,
55 subscriber,
56 span,
57 } = key.into();
58 self.sender.enqueue(Message::Fetch(vec![FetchKey {
59 key,
60 subscribers: NonEmptyVec::new((subscriber, span)),
61 metadata: None,
62 }]))
63 }
64
65 fn fetch_all<D>(&mut self, keys: Vec<D>) -> Feedback
72 where
73 D: Into<Fetch<Self::Key, Self::Subscriber>> + Send,
74 {
75 self.sender.enqueue(Message::Fetch(
76 keys.into_iter()
77 .map(|key| fetch_key(key.into(), None))
78 .collect(),
79 ))
80 }
81
82 fn retain(
86 &mut self,
87 predicate: impl Fn(&Self::Key, &Self::Subscriber) -> bool + Send + 'static,
88 ) -> Feedback {
89 self.sender.enqueue(Message::Retain {
90 predicate: Box::new(predicate),
91 })
92 }
93}
94
95impl<K, P, S> TargetedResolver for Mailbox<K, P, S>
96where
97 K: Span,
98 P: PublicKey,
99 S: Clone + Eq + Send + 'static,
100{
101 type PublicKey = P;
102
103 fn fetch_targeted(
119 &mut self,
120 key: impl Into<Fetch<Self::Key, Self::Subscriber>> + Send,
121 targets: NonEmptyVec<Self::PublicKey>,
122 ) -> Feedback {
123 let Fetch {
124 key,
125 subscriber,
126 span,
127 } = key.into();
128 self.sender.enqueue(Message::Fetch(vec![FetchKey {
129 key,
130 subscribers: NonEmptyVec::new((subscriber, span)),
131 metadata: Some(targets),
132 }]))
133 }
134
135 fn fetch_all_targeted<D>(&mut self, keys: Vec<(D, NonEmptyVec<Self::PublicKey>)>) -> Feedback
139 where
140 D: Into<Fetch<Self::Key, Self::Subscriber>> + Send,
141 {
142 self.sender.enqueue(Message::Fetch(
143 keys.into_iter()
144 .map(|(key, targets)| fetch_key(key.into(), Some(targets)))
145 .collect(),
146 ))
147 }
148}
149
150#[cfg(test)]
151mod tests {
152 use super::*;
153 use commonware_actor::mailbox::{Overflow, Policy};
154
155 type TestMessage = Message<u8, u8, u16>;
156 type TestPending = ingress::Pending<u8, u16, Option<NonEmptyVec<u8>>>;
157
158 fn fetch(key: u8, subscriber: u16, targets: Option<NonEmptyVec<u8>>) -> TestMessage {
159 Message::Fetch(vec![FetchKey {
160 key,
161 subscribers: NonEmptyVec::new((subscriber, tracing::Span::none())),
162 metadata: targets,
163 }])
164 }
165
166 fn fetch_with_subscribers(
167 key: u8,
168 subscribers: Vec<u16>,
169 targets: Option<NonEmptyVec<u8>>,
170 ) -> TestMessage {
171 Message::Fetch(vec![FetchKey {
172 key,
173 subscribers: NonEmptyVec::from_unchecked(
174 subscribers
175 .into_iter()
176 .map(|subscriber| (subscriber, tracing::Span::none()))
177 .collect(),
178 ),
179 metadata: targets,
180 }])
181 }
182
183 fn subscriber_is(value: u16) -> impl Fn(&u8, &u16) -> bool + Send {
184 move |_, subscriber| *subscriber == value
185 }
186
187 fn targets(values: &[u8]) -> NonEmptyVec<u8> {
188 NonEmptyVec::from_unchecked(values.to_vec())
189 }
190
191 fn drain(pending: &mut TestPending) -> Vec<TestMessage> {
192 let mut messages = Vec::new();
193 Overflow::drain(pending, |message| {
194 messages.push(message);
195 None
196 });
197 messages
198 }
199
200 fn assert_fetch(message: &TestMessage, expected_key: u8, expected_targets: Option<&[u8]>) {
201 let Message::Fetch(keys) = message else {
202 panic!("expected fetch");
203 };
204 assert_eq!(keys.len(), 1);
205 assert_eq!(keys[0].key, expected_key);
206 match (&keys[0].metadata, expected_targets) {
207 (None, None) => {}
208 (Some(actual), Some(expected)) => assert_eq!(&actual[..], expected),
209 _ => panic!("unexpected targets"),
210 }
211 }
212
213 fn assert_fetch_keys(message: &TestMessage, expected: &[u8]) {
214 let Message::Fetch(keys) = message else {
215 panic!("expected fetch");
216 };
217 let actual: Vec<_> = keys.iter().map(|key| key.key).collect();
218 assert_eq!(actual, expected);
219 }
220
221 fn assert_fetch_subscribers(
222 message: &TestMessage,
223 expected_key: u8,
224 expected_subscribers: &[u16],
225 ) {
226 let Message::Fetch(keys) = message else {
227 panic!("expected fetch");
228 };
229 assert_eq!(keys.len(), 1);
230 assert_eq!(keys[0].key, expected_key);
231 let actual: Vec<_> = keys[0]
232 .subscribers
233 .iter()
234 .map(|(subscriber, _)| *subscriber)
235 .collect();
236 assert_eq!(actual, expected_subscribers);
237 }
238
239 #[test]
240 fn targeted_fetches_for_same_key_are_merged() {
241 let mut pending = TestPending::default();
242
243 Policy::handle(&mut pending, fetch(1, 10, Some(targets(&[2, 3]))));
244 Policy::handle(&mut pending, fetch(1, 11, Some(targets(&[3, 4]))));
245
246 let messages = drain(&mut pending);
247 assert_eq!(messages.len(), 1);
248 assert_fetch(&messages[0], 1, Some(&[2, 3, 4]));
249 assert_fetch_subscribers(&messages[0], 1, &[10, 11]);
250 }
251
252 #[test]
253 fn duplicate_fetches_for_same_key_merge_subscribers() {
254 let mut pending = TestPending::default();
255
256 Policy::handle(&mut pending, fetch_with_subscribers(1, vec![10, 11], None));
257 Policy::handle(&mut pending, fetch_with_subscribers(1, vec![11, 12], None));
258
259 let messages = drain(&mut pending);
260 assert_eq!(messages.len(), 1);
261 assert_fetch_subscribers(&messages[0], 1, &[10, 11, 12]);
262 }
263
264 #[test]
265 fn unrestricted_fetch_dominates_targeted_fetches() {
266 let mut pending = TestPending::default();
267
268 Policy::handle(&mut pending, fetch(1, 10, Some(targets(&[2]))));
269 Policy::handle(&mut pending, fetch(1, 11, None));
270 Policy::handle(&mut pending, fetch(1, 12, Some(targets(&[3]))));
271
272 let messages = drain(&mut pending);
273 assert_eq!(messages.len(), 1);
274 assert_fetch(&messages[0], 1, None);
275 }
276
277 #[test]
278 fn retain_removes_fetches_for_dropped_subscribers() {
279 let mut pending = TestPending::default();
280
281 Policy::handle(&mut pending, fetch(1, 10, None));
282 Policy::handle(&mut pending, fetch(2, 11, None));
283 Policy::handle(
284 &mut pending,
285 Message::Retain {
286 predicate: Box::new(subscriber_is(11)),
287 },
288 );
289
290 let messages = drain(&mut pending);
291 assert_eq!(messages.len(), 2);
292 assert!(matches!(messages[0], Message::Retain { .. }));
293 assert_fetch(&messages[1], 2, None);
294 }
295
296 #[test]
297 fn retain_prunes_pending_fetch_subscribers() {
298 let mut pending = TestPending::default();
299
300 Policy::handle(&mut pending, fetch_with_subscribers(1, vec![10, 11], None));
301 Policy::handle(
302 &mut pending,
303 Message::Retain {
304 predicate: Box::new(subscriber_is(11)),
305 },
306 );
307
308 let messages = drain(&mut pending);
309 assert_eq!(messages.len(), 2);
310 assert!(matches!(messages[0], Message::Retain { .. }));
311 assert_fetch_subscribers(&messages[1], 1, &[11]);
312 }
313
314 #[test]
315 fn retain_drops_pending_fetch_when_all_subscribers_are_dropped() {
316 let mut pending = TestPending::default();
317
318 Policy::handle(&mut pending, fetch_with_subscribers(1, vec![10, 11], None));
319 Policy::handle(
320 &mut pending,
321 Message::Retain {
322 predicate: Box::new(subscriber_is(12)),
323 },
324 );
325
326 let messages = drain(&mut pending);
327 assert_eq!(messages.len(), 1);
328 assert!(matches!(messages[0], Message::Retain { .. }));
329 }
330
331 #[test]
332 fn fetch_after_retain_is_retained_when_subscriber_is_dropped() {
333 let mut pending = TestPending::default();
334
335 Policy::handle(
336 &mut pending,
337 Message::Retain {
338 predicate: Box::new(|_, subscriber| *subscriber != 10),
339 },
340 );
341 Policy::handle(&mut pending, fetch(1, 10, None));
342 Policy::handle(&mut pending, fetch(2, 11, None));
343
344 let messages = drain(&mut pending);
345 assert_eq!(messages.len(), 2);
346 assert!(matches!(messages[0], Message::Retain { .. }));
347 assert_fetch_keys(&messages[1], &[1, 2]);
348 }
349}