Skip to main content

commonware_resolver/p2p/
ingress.rs

1use crate::{Fetch, Resolver, TargetedResolver, ingress};
2use commonware_actor::{Feedback, mailbox::Sender};
3use commonware_cryptography::PublicKey;
4use commonware_utils::{Span, vec::NonEmptyVec};
5
6/// A key to fetch data for, optionally with target peers.
7pub type FetchKey<K, P, S> = ingress::FetchKey<K, S, Option<NonEmptyVec<P>>>;
8
9/// Messages that can be sent to the peer actor.
10pub 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/// A way to send messages to the peer actor.
21#[derive(Clone)]
22pub struct Mailbox<K: Span, P: Eq, S: Eq = ()> {
23    /// The channel that delivers messages to the peer actor.
24    sender: Sender<Message<K, P, S>>,
25}
26
27impl<K: Span, P: Eq, S: Eq> Mailbox<K, P, S> {
28    /// Create a new mailbox.
29    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    /// Send a fetch to the peer actor.
44    ///
45    /// If a fetch is already in progress for this key, this clears any existing
46    /// targets for that key (the fetch will try any available peer).
47    ///
48    /// If the engine has shut down, this is a no-op.
49    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    /// Send fetches to the peer actor for a batch of keys.
66    ///
67    /// If a fetch is already in progress for any key, this clears any existing
68    /// targets for that key (the fetch will try any available peer).
69    ///
70    /// If the engine has shut down, this is a no-op.
71    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    /// Send a retain request to the peer actor.
83    ///
84    /// If the engine has shut down, this is a no-op.
85    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    /// Send a targeted fetch to the peer actor.
104    ///
105    /// Targets may be any peers of the latest tracked peer set, primary or secondary. A target
106    /// outside that set is skipped until a later peer set includes it, and a target blocked for
107    /// invalid data is skipped until the network unblocks it.
108    ///
109    /// If a fetch is already in progress for this key:
110    /// - If the existing fetch has targets, the new targets are added to the set.
111    /// - If the existing fetch has no targets, it remains unrestricted.
112    ///
113    /// To clear targeting and fall back to any peer of the latest tracked peer set, call
114    /// [`fetch`](Self::fetch). Targets are automatically cleared when the fetch succeeds or is
115    /// canceled.
116    ///
117    /// If the engine has shut down, this is a no-op.
118    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    /// Send targeted fetches to the peer actor for a batch of keys.
136    ///
137    /// If the engine has shut down, this is a no-op.
138    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}