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(
116 &mut self,
117 key: impl Into<Fetch<Self::Key, Self::Subscriber>> + Send,
118 targets: NonEmptyVec<Self::PublicKey>,
119 ) -> Feedback {
120 let Fetch {
121 key,
122 subscriber,
123 span,
124 } = key.into();
125 self.sender.enqueue(Message::Fetch(vec![FetchKey {
126 key,
127 subscribers: NonEmptyVec::new((subscriber, span)),
128 metadata: Some(targets),
129 }]))
130 }
131
132 fn fetch_all_targeted<D>(&mut self, keys: Vec<(D, NonEmptyVec<Self::PublicKey>)>) -> Feedback
136 where
137 D: Into<Fetch<Self::Key, Self::Subscriber>> + Send,
138 {
139 self.sender.enqueue(Message::Fetch(
140 keys.into_iter()
141 .map(|(key, targets)| fetch_key(key.into(), Some(targets)))
142 .collect(),
143 ))
144 }
145}
146
147#[cfg(test)]
148mod tests {
149 use super::*;
150 use commonware_actor::mailbox::{Overflow, Policy};
151
152 type TestMessage = Message<u8, u8, u16>;
153 type TestPending = ingress::Pending<u8, u16, Option<NonEmptyVec<u8>>>;
154
155 fn fetch(key: u8, subscriber: u16, targets: Option<NonEmptyVec<u8>>) -> TestMessage {
156 Message::Fetch(vec![FetchKey {
157 key,
158 subscribers: NonEmptyVec::new((subscriber, tracing::Span::none())),
159 metadata: targets,
160 }])
161 }
162
163 fn fetch_with_subscribers(
164 key: u8,
165 subscribers: Vec<u16>,
166 targets: Option<NonEmptyVec<u8>>,
167 ) -> TestMessage {
168 Message::Fetch(vec![FetchKey {
169 key,
170 subscribers: NonEmptyVec::from_unchecked(
171 subscribers
172 .into_iter()
173 .map(|subscriber| (subscriber, tracing::Span::none()))
174 .collect(),
175 ),
176 metadata: targets,
177 }])
178 }
179
180 fn subscriber_is(value: u16) -> impl Fn(&u8, &u16) -> bool + Send {
181 move |_, subscriber| *subscriber == value
182 }
183
184 fn targets(values: &[u8]) -> NonEmptyVec<u8> {
185 NonEmptyVec::from_unchecked(values.to_vec())
186 }
187
188 fn drain(pending: &mut TestPending) -> Vec<TestMessage> {
189 let mut messages = Vec::new();
190 Overflow::drain(pending, |message| {
191 messages.push(message);
192 None
193 });
194 messages
195 }
196
197 fn assert_fetch(message: &TestMessage, expected_key: u8, expected_targets: Option<&[u8]>) {
198 let Message::Fetch(keys) = message else {
199 panic!("expected fetch");
200 };
201 assert_eq!(keys.len(), 1);
202 assert_eq!(keys[0].key, expected_key);
203 match (&keys[0].metadata, expected_targets) {
204 (None, None) => {}
205 (Some(actual), Some(expected)) => assert_eq!(&actual[..], expected),
206 _ => panic!("unexpected targets"),
207 }
208 }
209
210 fn assert_fetch_keys(message: &TestMessage, expected: &[u8]) {
211 let Message::Fetch(keys) = message else {
212 panic!("expected fetch");
213 };
214 let actual: Vec<_> = keys.iter().map(|key| key.key).collect();
215 assert_eq!(actual, expected);
216 }
217
218 fn assert_fetch_subscribers(
219 message: &TestMessage,
220 expected_key: u8,
221 expected_subscribers: &[u16],
222 ) {
223 let Message::Fetch(keys) = message else {
224 panic!("expected fetch");
225 };
226 assert_eq!(keys.len(), 1);
227 assert_eq!(keys[0].key, expected_key);
228 let actual: Vec<_> = keys[0]
229 .subscribers
230 .iter()
231 .map(|(subscriber, _)| *subscriber)
232 .collect();
233 assert_eq!(actual, expected_subscribers);
234 }
235
236 #[test]
237 fn targeted_fetches_for_same_key_are_merged() {
238 let mut pending = TestPending::default();
239
240 Policy::handle(&mut pending, fetch(1, 10, Some(targets(&[2, 3]))));
241 Policy::handle(&mut pending, fetch(1, 11, Some(targets(&[3, 4]))));
242
243 let messages = drain(&mut pending);
244 assert_eq!(messages.len(), 1);
245 assert_fetch(&messages[0], 1, Some(&[2, 3, 4]));
246 assert_fetch_subscribers(&messages[0], 1, &[10, 11]);
247 }
248
249 #[test]
250 fn duplicate_fetches_for_same_key_merge_subscribers() {
251 let mut pending = TestPending::default();
252
253 Policy::handle(&mut pending, fetch_with_subscribers(1, vec![10, 11], None));
254 Policy::handle(&mut pending, fetch_with_subscribers(1, vec![11, 12], None));
255
256 let messages = drain(&mut pending);
257 assert_eq!(messages.len(), 1);
258 assert_fetch_subscribers(&messages[0], 1, &[10, 11, 12]);
259 }
260
261 #[test]
262 fn unrestricted_fetch_dominates_targeted_fetches() {
263 let mut pending = TestPending::default();
264
265 Policy::handle(&mut pending, fetch(1, 10, Some(targets(&[2]))));
266 Policy::handle(&mut pending, fetch(1, 11, None));
267 Policy::handle(&mut pending, fetch(1, 12, Some(targets(&[3]))));
268
269 let messages = drain(&mut pending);
270 assert_eq!(messages.len(), 1);
271 assert_fetch(&messages[0], 1, None);
272 }
273
274 #[test]
275 fn retain_removes_fetches_for_dropped_subscribers() {
276 let mut pending = TestPending::default();
277
278 Policy::handle(&mut pending, fetch(1, 10, None));
279 Policy::handle(&mut pending, fetch(2, 11, None));
280 Policy::handle(
281 &mut pending,
282 Message::Retain {
283 predicate: Box::new(subscriber_is(11)),
284 },
285 );
286
287 let messages = drain(&mut pending);
288 assert_eq!(messages.len(), 2);
289 assert!(matches!(messages[0], Message::Retain { .. }));
290 assert_fetch(&messages[1], 2, None);
291 }
292
293 #[test]
294 fn retain_prunes_pending_fetch_subscribers() {
295 let mut pending = TestPending::default();
296
297 Policy::handle(&mut pending, fetch_with_subscribers(1, vec![10, 11], None));
298 Policy::handle(
299 &mut pending,
300 Message::Retain {
301 predicate: Box::new(subscriber_is(11)),
302 },
303 );
304
305 let messages = drain(&mut pending);
306 assert_eq!(messages.len(), 2);
307 assert!(matches!(messages[0], Message::Retain { .. }));
308 assert_fetch_subscribers(&messages[1], 1, &[11]);
309 }
310
311 #[test]
312 fn retain_drops_pending_fetch_when_all_subscribers_are_dropped() {
313 let mut pending = TestPending::default();
314
315 Policy::handle(&mut pending, fetch_with_subscribers(1, vec![10, 11], None));
316 Policy::handle(
317 &mut pending,
318 Message::Retain {
319 predicate: Box::new(subscriber_is(12)),
320 },
321 );
322
323 let messages = drain(&mut pending);
324 assert_eq!(messages.len(), 1);
325 assert!(matches!(messages[0], Message::Retain { .. }));
326 }
327
328 #[test]
329 fn fetch_after_retain_is_retained_when_subscriber_is_dropped() {
330 let mut pending = TestPending::default();
331
332 Policy::handle(
333 &mut pending,
334 Message::Retain {
335 predicate: Box::new(|_, subscriber| *subscriber != 10),
336 },
337 );
338 Policy::handle(&mut pending, fetch(1, 10, None));
339 Policy::handle(&mut pending, fetch(2, 11, None));
340
341 let messages = drain(&mut pending);
342 assert_eq!(messages.len(), 2);
343 assert!(matches!(messages[0], Message::Retain { .. }));
344 assert_fetch_keys(&messages[1], &[1, 2]);
345 }
346}