1use crate::{
14 delivery::{Completion as DeliveryCompletion, Tracker as DeliveryTracker},
15 ingress::{self, FetchKey, Message},
16 subscribers, Consumer, Delivery, Fetch, TargetedResolver,
17};
18use commonware_actor::{mailbox, Feedback};
19use commonware_cryptography::PublicKey;
20use commonware_macros::select_loop;
21use commonware_runtime::{spawn_cell, Clock, ContextCell, Handle, Metrics, Spawner};
22use commonware_utils::{
23 futures::{AbortablePool, Aborter},
24 vec::NonEmptyVec,
25 Span,
26};
27use futures::future::{self, Either};
28use std::{
29 collections::{BTreeMap, BTreeSet},
30 future::Future,
31 marker::PhantomData,
32 num::NonZeroUsize,
33 time::{Duration, SystemTime},
34};
35use tracing::{debug, trace, warn};
36
37pub trait Fetcher {
39 type Key: Span;
41
42 type Value;
44
45 fn fetch(&self, key: Self::Key) -> impl Future<Output = Option<Self::Value>> + Send;
51}
52
53pub struct Resolver<K, S, P>
55where
56 K: Span,
57 S: Clone + Eq + Send + 'static,
58 P: PublicKey,
59{
60 mailbox: mailbox::Sender<Message<K, S>>,
61 _peer: PhantomData<P>,
62}
63
64impl<K, S, P> Clone for Resolver<K, S, P>
65where
66 K: Span,
67 S: Clone + Eq + Send + 'static,
68 P: PublicKey,
69{
70 fn clone(&self) -> Self {
71 Self {
72 mailbox: self.mailbox.clone(),
73 _peer: PhantomData,
74 }
75 }
76}
77
78impl<K, S, P> crate::Resolver for Resolver<K, S, P>
79where
80 K: Span,
81 S: Clone + Eq + Send + 'static,
82 P: PublicKey,
83{
84 type Key = K;
85 type Subscriber = S;
86
87 fn fetch<F>(&mut self, fetch: F) -> Feedback
88 where
89 F: Into<Fetch<Self::Key, Self::Subscriber>> + Send,
90 {
91 self.send(Message::Fetch(vec![FetchKey::from(fetch.into())]))
92 }
93
94 fn fetch_all<F>(&mut self, fetches: Vec<F>) -> Feedback
95 where
96 F: Into<Fetch<Self::Key, Self::Subscriber>> + Send,
97 {
98 self.send(Message::Fetch(
99 fetches
100 .into_iter()
101 .map(|fetch| FetchKey::from(fetch.into()))
102 .collect(),
103 ))
104 }
105
106 fn retain(
107 &mut self,
108 predicate: impl Fn(&Self::Key, &Self::Subscriber) -> bool + Send + 'static,
109 ) -> Feedback {
110 self.send(Message::Retain {
111 predicate: Box::new(predicate),
112 })
113 }
114}
115
116impl<K, S, P> TargetedResolver for Resolver<K, S, P>
117where
118 K: Span,
119 S: Clone + Eq + Send + 'static,
120 P: PublicKey,
121{
122 type PublicKey = P;
123
124 fn fetch_targeted(
125 &mut self,
126 fetch: impl Into<Fetch<Self::Key, Self::Subscriber>> + Send,
127 _targets: NonEmptyVec<Self::PublicKey>,
128 ) -> Feedback {
129 <Self as crate::Resolver>::fetch(self, fetch)
130 }
131
132 fn fetch_all_targeted<F>(&mut self, fetches: Vec<(F, NonEmptyVec<Self::PublicKey>)>) -> Feedback
133 where
134 F: Into<Fetch<Self::Key, Self::Subscriber>> + Send,
135 {
136 <Self as crate::Resolver>::fetch_all(
137 self,
138 fetches.into_iter().map(|(fetch, _)| fetch).collect(),
139 )
140 }
141}
142
143impl<K, S, P> Resolver<K, S, P>
144where
145 K: Span,
146 S: Clone + Eq + Send + 'static,
147 P: PublicKey,
148{
149 const fn new(mailbox: mailbox::Sender<Message<K, S>>) -> Self {
150 Self {
151 mailbox,
152 _peer: PhantomData,
153 }
154 }
155
156 fn send(&self, message: Message<K, S>) -> Feedback {
157 self.mailbox.enqueue(message)
158 }
159}
160
161pub fn init<E, F, Con, P>(
163 context: E,
164 fetcher: F,
165 consumer: Con,
166 mailbox_size: NonZeroUsize,
167 fetch_retry_timeout: Duration,
168) -> Resolver<F::Key, Con::Subscriber, P>
169where
170 E: Clock + Spawner + Metrics,
171 F: Fetcher + Clone + Send + 'static,
172 F::Value: Clone + Send + 'static,
173 Con: Consumer<Key = F::Key, Value = F::Value>,
174 Con::Subscriber: Ord,
175 P: PublicKey,
176{
177 let (mailbox_tx, mailbox_rx) = mailbox::new(context.child("mailbox"), mailbox_size);
178 Actor::new(
179 context.child("actor"),
180 fetcher,
181 mailbox_rx,
182 consumer,
183 fetch_retry_timeout,
184 )
185 .start();
186 Resolver::new(mailbox_tx)
187}
188
189struct Actor<E, F, Con>
191where
192 E: Clock + Spawner,
193 F: Fetcher,
194 F::Value: Clone + Send + 'static,
195 Con: Consumer<Key = F::Key, Value = F::Value>,
196 Con::Subscriber: Ord,
197{
198 context: ContextCell<E>,
199 fetcher: F,
200 mailbox: mailbox::Receiver<Message<F::Key, Con::Subscriber>>,
201 fetches: AbortablePool<FetchCompletion<F::Key, F::Value>>,
202 deliveries: DeliveryTracker<Con, u64>,
203 requests: BTreeMap<F::Key, Attempt>,
204 subscribers: subscribers::Tracker<F::Key, Con::Subscriber>,
205 retry_schedule: BTreeSet<(SystemTime, F::Key)>,
206 fetch_retry_timeout: Duration,
207 next_id: u64,
208}
209
210enum Attempt {
211 Fetching { id: u64, _aborter: Aborter },
213
214 Delivering { id: u64 },
216
217 Scheduled(SystemTime),
219}
220
221struct FetchCompletion<K, V> {
222 key: K,
223 id: u64,
224 value: Option<V>,
225}
226
227impl<E, F, Con> Actor<E, F, Con>
228where
229 E: Clock + Spawner,
230 F: Fetcher + Clone + Send + 'static,
231 F::Value: Clone + Send + 'static,
232 Con: Consumer<Key = F::Key, Value = F::Value>,
233 Con::Subscriber: Ord,
234{
235 fn new(
236 context: E,
237 fetcher: F,
238 mailbox: mailbox::Receiver<Message<F::Key, Con::Subscriber>>,
239 consumer: Con,
240 fetch_retry_timeout: Duration,
241 ) -> Self {
242 Self {
243 context: ContextCell::new(context),
244 fetcher,
245 mailbox,
246 fetches: AbortablePool::default(),
247 deliveries: DeliveryTracker::new(consumer),
248 requests: BTreeMap::new(),
249 subscribers: subscribers::Tracker::new(),
250 retry_schedule: BTreeSet::new(),
251 fetch_retry_timeout,
252 next_id: 0,
253 }
254 }
255
256 fn start(mut self) -> Handle<()> {
257 spawn_cell!(self.context, self.run())
258 }
259
260 async fn run(mut self) {
261 select_loop! {
262 self.context,
263 on_stopped => {},
264 Ok(result) = self.fetches.next_completed() else continue => {
265 self.handle_fetch_completed(result);
266 },
267 delivery = self.deliveries.next_completion() => {
268 let delivery = match delivery {
269 Ok(delivery) => delivery,
270 Err(_) => continue,
271 };
272 self.handle_delivery_completed(delivery);
273 },
274 _ = match self.retry_schedule.first() {
275 Some((deadline, _)) => Either::Left(self.context.sleep_until(*deadline)),
276 None => Either::Right(future::pending()),
277 } => {
278 self.process_retries();
279 },
280 Some(message) = self.mailbox.recv() else break => {
281 self.handle_message(message);
282 },
283 }
284 }
285
286 fn handle_message(&mut self, message: Message<F::Key, Con::Subscriber>) {
288 match message {
289 Message::Fetch(fetches) => {
290 for fetch in fetches {
291 self.add_fetch(fetch);
292 }
293 }
294 Message::Retain { predicate } => self.retain(predicate),
295 }
296 }
297
298 fn add_fetch(&mut self, fetch: FetchKey<F::Key, Con::Subscriber>) {
300 let FetchKey {
301 key, subscribers, ..
302 } = fetch;
303 let is_new = self.subscribers.insert(key.clone(), subscribers);
304
305 if is_new {
306 assert!(self.deliveries.insert(key.clone()), "delivery entry");
307 self.requests
308 .insert(key.clone(), Attempt::Scheduled(self.context.current()));
309 self.start_fetch(key);
310 }
311 }
312
313 fn retain(&mut self, predicate: ingress::Predicate<F::Key, Con::Subscriber>) {
315 for key in self
316 .subscribers
317 .retain(|key, subscriber| predicate(key, subscriber))
318 {
319 self.deliveries.remove(&key);
320 if let Some(attempt) = self.requests.remove(&key) {
321 match attempt {
322 Attempt::Fetching { .. } | Attempt::Delivering { .. } => {}
323 Attempt::Scheduled(deadline) => {
324 self.retry_schedule.remove(&(deadline, key));
325 }
326 }
327 }
328 }
329 }
330
331 fn start_fetch(&mut self, key: F::Key) {
333 let id = self.next_id;
334 self.next_id = self.next_id.wrapping_add(1);
335 let future = Self::fetch(key.clone(), id, self.fetcher.clone());
336 let aborter = self.fetches.push(future);
337 self.requests.insert(
338 key,
339 Attempt::Fetching {
340 id,
341 _aborter: aborter,
342 },
343 );
344 }
345
346 fn start_delivery(
348 &mut self,
349 key: F::Key,
350 value: F::Value,
351 delivered: NonEmptyVec<(Con::Subscriber, tracing::Span)>,
352 ) {
353 let id = self.next_id;
354 self.next_id = self.next_id.wrapping_add(1);
355 self.deliveries.deliver(
356 Delivery {
357 key: key.clone(),
358 subscribers: delivered,
359 },
360 id,
361 value,
362 );
363 self.requests.insert(key, Attempt::Delivering { id });
364 }
365
366 fn redeliver(&mut self, key: F::Key, delivered: NonEmptyVec<(Con::Subscriber, tracing::Span)>) {
368 self.deliveries.redeliver(Delivery {
369 key,
370 subscribers: delivered,
371 });
372 }
373
374 fn handle_fetch_completed(&mut self, completion: FetchCompletion<F::Key, F::Value>) {
376 let FetchCompletion { key, id, value } = completion;
377 if !self.current_fetch(&key, id) {
378 return;
379 }
380 self.handle_fetched(key, value);
381 }
382
383 fn handle_delivery_completed(
385 &mut self,
386 completion: DeliveryCompletion<F::Key, Con::Subscriber, u64>,
387 ) {
388 let DeliveryCompletion {
389 context: id,
390 delivery,
391 valid,
392 } = completion;
393 let Delivery {
394 key,
395 subscribers: delivered,
396 ..
397 } = delivery;
398 if !self.current_delivery(&key, id) {
399 return;
400 }
401 self.handle_delivered(key, delivered, valid);
402 }
403
404 fn current_fetch(&self, key: &F::Key, id: u64) -> bool {
406 let Some(attempt) = self.requests.get(key) else {
407 trace!(?key, id, "ignoring stale fetch completion");
408 return false;
409 };
410 match attempt {
411 Attempt::Fetching { id: active_id, .. } if *active_id == id => true,
412 Attempt::Fetching { id: active_id, .. } => {
413 trace!(
414 ?key,
415 completed_id = id,
416 active_id,
417 "ignoring replaced fetch completion",
418 );
419 false
420 }
421 Attempt::Delivering { id: active_id } => {
422 trace!(
423 ?key,
424 completed_id = id,
425 active_id,
426 "ignoring fetch completion for delivery attempt",
427 );
428 false
429 }
430 Attempt::Scheduled(deadline) => {
431 trace!(?key, id, ?deadline, "ignoring scheduled fetch completion");
432 false
433 }
434 }
435 }
436
437 fn current_delivery(&self, key: &F::Key, id: u64) -> bool {
439 let Some(attempt) = self.requests.get(key) else {
440 trace!(?key, id, "ignoring stale delivery completion");
441 return false;
442 };
443 match attempt {
444 Attempt::Delivering { id: active_id } if *active_id == id => true,
445 Attempt::Delivering { id: active_id } => {
446 trace!(
447 ?key,
448 completed_id = id,
449 active_id,
450 "ignoring replaced delivery completion",
451 );
452 false
453 }
454 Attempt::Fetching { id: active_id, .. } => {
455 trace!(
456 ?key,
457 completed_id = id,
458 active_id,
459 "ignoring delivery completion for fetch attempt",
460 );
461 false
462 }
463 Attempt::Scheduled(deadline) => {
464 trace!(
465 ?key,
466 id,
467 ?deadline,
468 "ignoring scheduled delivery completion"
469 );
470 false
471 }
472 }
473 }
474
475 fn handle_fetched(&mut self, key: F::Key, value: Option<F::Value>) {
477 match value {
478 None => self.schedule_retry(key),
479 Some(value) => {
480 if let Some(subscribers) = self.subscribers.pending(&key) {
481 self.start_delivery(key, value, subscribers);
482 } else {
483 self.requests.remove(&key);
484 self.subscribers.remove(&key);
485 self.deliveries.remove(&key);
486 }
487 }
488 }
489 }
490
491 fn handle_delivered(
493 &mut self,
494 key: F::Key,
495 delivered: NonEmptyVec<(Con::Subscriber, tracing::Span)>,
496 valid: bool,
497 ) {
498 let accepted = self.deliveries.response_accepted(&key);
499
500 if valid {
501 let remaining = self
502 .subscribers
503 .remove_delivered(&key, delivered.map_into(|(subscriber, _)| subscriber));
504
505 if let Some(subscribers) = remaining {
509 if !accepted {
510 self.deliveries.accept_response(&key);
511 }
512 self.redeliver(key, subscribers);
513 } else {
514 self.requests.remove(&key);
515 self.subscribers.remove(&key);
516 self.deliveries.remove(&key);
517 }
518 return;
519 }
520
521 if accepted {
525 warn!(
526 ?key,
527 "previously accepted resolver response rejected during opaque redelivery"
528 );
529 self.requests.remove(&key);
530 self.subscribers.remove(&key);
531 self.deliveries.remove(&key);
532 return;
533 }
534
535 warn!(?key, "consumer rejected opaque resolver delivery");
536 self.deliveries.discard_response(&key);
537 self.schedule_retry(key);
538 }
539
540 fn schedule_retry(&mut self, key: F::Key) {
542 let deadline = self.context.current() + self.fetch_retry_timeout;
543 let Some(attempt) = self.requests.get_mut(&key) else {
544 return;
545 };
546 *attempt = Attempt::Scheduled(deadline);
547 debug!(?key, ?deadline, "scheduled opaque resolver retry");
548 self.retry_schedule.insert((deadline, key));
549 }
550
551 fn process_retries(&mut self) {
553 let now = self.context.current();
554 while let Some((deadline, key)) = self.retry_schedule.pop_first() {
555 if deadline > now {
556 self.retry_schedule.insert((deadline, key));
557 break;
558 }
559
560 let Some(state) = self.requests.get(&key) else {
561 continue;
562 };
563 match state {
564 Attempt::Scheduled(state_deadline) if *state_deadline == deadline => {
565 debug!(?key, "retrying opaque resolver fetch");
566 self.start_fetch(key);
567 }
568 Attempt::Scheduled(_) | Attempt::Fetching { .. } | Attempt::Delivering { .. } => {}
569 }
570 }
571 }
572
573 async fn fetch(key: F::Key, id: u64, fetcher: F) -> FetchCompletion<F::Key, F::Value> {
575 let value = fetcher.fetch(key.clone()).await;
576 FetchCompletion { key, id, value }
577 }
578}
579
580#[cfg(test)]
581mod tests {
582 use super::*;
583 use crate::Resolver as _;
584 use bytes::Bytes;
585 use commonware_cryptography::{
586 ed25519::{PrivateKey, PublicKey},
587 Signer,
588 };
589 use commonware_runtime::{deterministic, deterministic::Runner, Runner as _, Supervisor as _};
590 use commonware_utils::{channel::oneshot, non_empty_vec, sync::Mutex};
591 use std::{
592 collections::{HashMap, VecDeque},
593 sync::{
594 atomic::{AtomicU32, Ordering},
595 Arc,
596 },
597 };
598
599 const RETRY_TIMEOUT: Duration = Duration::from_millis(100);
600
601 #[derive(Clone, Default)]
602 struct MockFetcher {
603 responses: Arc<Mutex<HashMap<u8, VecDeque<Option<Bytes>>>>>,
604 calls: Arc<AtomicU32>,
605 }
606
607 impl MockFetcher {
608 fn push(&self, key: u8, response: Option<Bytes>) {
609 self.responses
610 .lock()
611 .entry(key)
612 .or_default()
613 .push_back(response);
614 }
615
616 fn calls(&self) -> u32 {
617 self.calls.load(Ordering::Relaxed)
618 }
619 }
620
621 impl Fetcher for MockFetcher {
622 type Key = u8;
623 type Value = Bytes;
624
625 fn fetch(&self, key: Self::Key) -> impl Future<Output = Option<Self::Value>> + Send {
626 let responses = self.responses.clone();
627 let calls = self.calls.clone();
628 async move {
629 calls.fetch_add(1, Ordering::Relaxed);
630 responses
631 .lock()
632 .get_mut(&key)
633 .and_then(VecDeque::pop_front)
634 .flatten()
635 }
636 }
637 }
638
639 #[derive(Clone)]
640 struct BlockingFetcher {
641 started: Arc<Mutex<Option<oneshot::Sender<()>>>>,
642 response: Arc<Mutex<Option<oneshot::Receiver<Option<Bytes>>>>>,
643 }
644
645 impl BlockingFetcher {
646 fn new() -> (Self, oneshot::Receiver<()>, oneshot::Sender<Option<Bytes>>) {
647 let (started_tx, started_rx) = oneshot::channel();
648 let (response_tx, response_rx) = oneshot::channel();
649 (
650 Self {
651 started: Arc::new(Mutex::new(Some(started_tx))),
652 response: Arc::new(Mutex::new(Some(response_rx))),
653 },
654 started_rx,
655 response_tx,
656 )
657 }
658 }
659
660 impl Fetcher for BlockingFetcher {
661 type Key = u8;
662 type Value = Bytes;
663
664 fn fetch(&self, _key: Self::Key) -> impl Future<Output = Option<Self::Value>> + Send {
665 let started = self.started.clone();
666 let response = self.response.clone();
667 async move {
668 if let Some(started) = started.lock().take() {
669 let _ = started.send(());
670 }
671 let response = response.lock().take().expect("missing response");
672 response.await.unwrap_or(None)
673 }
674 }
675 }
676
677 struct CapturedDelivery {
678 delivery: Delivery<u8, u16>,
679 value: Bytes,
680 response: oneshot::Sender<bool>,
681 }
682
683 #[derive(Clone, Default)]
684 struct MockConsumer {
685 deliveries: Arc<Mutex<VecDeque<CapturedDelivery>>>,
686 }
687
688 impl MockConsumer {
689 fn pop(&self) -> Option<CapturedDelivery> {
690 self.deliveries.lock().pop_front()
691 }
692
693 fn len(&self) -> usize {
694 self.deliveries.lock().len()
695 }
696 }
697
698 impl Consumer for MockConsumer {
699 type Key = u8;
700 type Value = Bytes;
701 type Subscriber = u16;
702
703 fn deliver(
704 &mut self,
705 delivery: Delivery<Self::Key, Self::Subscriber>,
706 value: Self::Value,
707 ) -> oneshot::Receiver<bool> {
708 let (response, receiver) = oneshot::channel();
709 self.deliveries.lock().push_back(CapturedDelivery {
710 delivery,
711 value,
712 response,
713 });
714 receiver
715 }
716 }
717
718 fn start_resolver<F>(
719 context: deterministic::Context,
720 fetcher: F,
721 consumer: MockConsumer,
722 ) -> Resolver<u8, u16, PublicKey>
723 where
724 F: Fetcher<Key = u8, Value = Bytes> + Clone + Send + 'static,
725 {
726 init(
727 context,
728 fetcher,
729 consumer,
730 NonZeroUsize::new(16).unwrap(),
731 RETRY_TIMEOUT,
732 )
733 }
734
735 async fn wait_for_delivery(
736 context: &deterministic::Context,
737 consumer: &MockConsumer,
738 ) -> CapturedDelivery {
739 for _ in 0..50 {
740 if let Some(delivery) = consumer.pop() {
741 return delivery;
742 }
743 context.sleep(Duration::from_millis(10)).await;
744 }
745 panic!("timed out waiting for delivery");
746 }
747
748 #[test]
749 fn fetch_during_validation_reuses_response_after_success() {
750 Runner::default().start(|context| async move {
751 let fetcher = MockFetcher::default();
752 fetcher.push(1, Some(Bytes::from_static(b"value")));
753 let consumer = MockConsumer::default();
754 let mut resolver =
755 start_resolver(context.child("resolver"), fetcher.clone(), consumer.clone());
756
757 assert!(resolver
758 .fetch(Fetch {
759 key: 1,
760 subscriber: 10,
761 span: tracing::Span::none(),
762 })
763 .accepted());
764 let first = wait_for_delivery(&context, &consumer).await;
765 assert_eq!(first.value, Bytes::from_static(b"value"));
766
767 assert!(resolver
768 .fetch(Fetch {
769 key: 1,
770 subscriber: 11,
771 span: tracing::Span::none(),
772 })
773 .accepted());
774 context.sleep(Duration::from_millis(10)).await;
775 first.response.send(true).expect("response dropped");
776
777 let second = wait_for_delivery(&context, &consumer).await;
778 assert_eq!(second.value, Bytes::from_static(b"value"));
779 assert_eq!(
780 second
781 .delivery
782 .subscribers
783 .iter()
784 .map(|(subscriber, _)| *subscriber)
785 .collect::<Vec<_>>(),
786 vec![11]
787 );
788 second.response.send(true).expect("response dropped");
789
790 context.sleep(Duration::from_millis(10)).await;
791 assert_eq!(fetcher.calls(), 1);
792 });
793 }
794
795 #[test]
796 fn missing_fetch_retries_until_value_is_available() {
797 Runner::default().start(|context| async move {
798 let fetcher = MockFetcher::default();
799 fetcher.push(1, None);
800 fetcher.push(1, Some(Bytes::from_static(b"value")));
801 let consumer = MockConsumer::default();
802 let mut resolver =
803 start_resolver(context.child("resolver"), fetcher.clone(), consumer.clone());
804
805 assert!(resolver
806 .fetch(Fetch {
807 key: 1,
808 subscriber: 10,
809 span: tracing::Span::none(),
810 })
811 .accepted());
812 context
813 .sleep(RETRY_TIMEOUT + Duration::from_millis(10))
814 .await;
815
816 let delivery = wait_for_delivery(&context, &consumer).await;
817 assert_eq!(delivery.value, Bytes::from_static(b"value"));
818 delivery.response.send(true).expect("response dropped");
819 assert_eq!(fetcher.calls(), 2);
820 });
821 }
822
823 #[test]
824 fn accepted_redelivery_rejection_does_not_refetch() {
825 Runner::default().start(|context| async move {
826 let fetcher = MockFetcher::default();
827 fetcher.push(1, Some(Bytes::from_static(b"value")));
828 let consumer = MockConsumer::default();
829 let mut resolver =
830 start_resolver(context.child("resolver"), fetcher.clone(), consumer.clone());
831
832 assert!(resolver
833 .fetch(Fetch {
834 key: 1,
835 subscriber: 10,
836 span: tracing::Span::none(),
837 })
838 .accepted());
839 let first = wait_for_delivery(&context, &consumer).await;
840
841 assert!(resolver
842 .fetch(Fetch {
843 key: 1,
844 subscriber: 11,
845 span: tracing::Span::none(),
846 })
847 .accepted());
848 context.sleep(Duration::from_millis(10)).await;
849 first.response.send(true).expect("response dropped");
850
851 let second = wait_for_delivery(&context, &consumer).await;
852 second.response.send(false).expect("response dropped");
853
854 context
855 .sleep(RETRY_TIMEOUT + Duration::from_millis(10))
856 .await;
857 assert_eq!(fetcher.calls(), 1);
858 assert_eq!(consumer.len(), 0);
859 });
860 }
861
862 #[test]
863 fn retain_prunes_active_fetch_subscribers() {
864 Runner::default().start(|context| async move {
865 let (fetcher, started, response) = BlockingFetcher::new();
866 let consumer = MockConsumer::default();
867 let mut resolver = start_resolver(context.child("resolver"), fetcher, consumer.clone());
868
869 assert!(resolver
870 .fetch(Fetch {
871 key: 1,
872 subscriber: 10,
873 span: tracing::Span::none(),
874 })
875 .accepted());
876 assert!(resolver
877 .fetch(Fetch {
878 key: 1,
879 subscriber: 11,
880 span: tracing::Span::none(),
881 })
882 .accepted());
883 started.await.expect("fetch did not start");
884 assert!(resolver
885 .retain(|_, subscriber| *subscriber == 11)
886 .accepted());
887 context.sleep(Duration::from_millis(10)).await;
888 response
889 .send(Some(Bytes::from_static(b"value")))
890 .expect("fetcher dropped");
891
892 let delivery = wait_for_delivery(&context, &consumer).await;
893 assert_eq!(
894 delivery
895 .delivery
896 .subscribers
897 .iter()
898 .map(|(subscriber, _)| *subscriber)
899 .collect::<Vec<_>>(),
900 vec![11]
901 );
902 delivery.response.send(true).expect("response dropped");
903 });
904 }
905
906 #[test]
907 fn retain_drops_last_subscriber_aborts_active_fetch() {
908 Runner::default().start(|context| async move {
909 let (fetcher, started, response) = BlockingFetcher::new();
910 let consumer = MockConsumer::default();
911 let mut resolver = start_resolver(context.child("resolver"), fetcher, consumer.clone());
912
913 assert!(resolver
914 .fetch(Fetch {
915 key: 1,
916 subscriber: 10,
917 span: tracing::Span::none(),
918 })
919 .accepted());
920 started.await.expect("fetch did not start");
921 assert!(resolver.retain(|_, _| false).accepted());
922 context.sleep(Duration::from_millis(10)).await;
923
924 assert!(
925 response.send(Some(Bytes::from_static(b"value"))).is_err(),
926 "fetch future should be aborted after its last subscriber is pruned"
927 );
928 context
929 .sleep(RETRY_TIMEOUT + Duration::from_millis(10))
930 .await;
931 assert_eq!(consumer.len(), 0);
932 });
933 }
934
935 #[test]
936 fn retain_drops_last_subscriber_aborts_active_delivery() {
937 Runner::default().start(|context| async move {
938 let fetcher = MockFetcher::default();
939 fetcher.push(1, Some(Bytes::from_static(b"value")));
940 let consumer = MockConsumer::default();
941 let mut resolver =
942 start_resolver(context.child("resolver"), fetcher.clone(), consumer.clone());
943
944 assert!(resolver
945 .fetch(Fetch {
946 key: 1,
947 subscriber: 10,
948 span: tracing::Span::none(),
949 })
950 .accepted());
951 let delivery = wait_for_delivery(&context, &consumer).await;
952 assert!(resolver.retain(|_, _| false).accepted());
953 context.sleep(Duration::from_millis(10)).await;
954
955 assert!(
956 delivery.response.send(false).is_err(),
957 "delivery future should be aborted after its last subscriber is pruned"
958 );
959 context
960 .sleep(RETRY_TIMEOUT + Duration::from_millis(10))
961 .await;
962 assert_eq!(fetcher.calls(), 1);
963 assert_eq!(consumer.len(), 0);
964 });
965 }
966
967 #[test]
968 fn targeted_fetch_uses_same_opaque_fetch_path() {
969 Runner::default().start(|context| async move {
970 let fetcher = MockFetcher::default();
971 fetcher.push(1, Some(Bytes::from_static(b"value")));
972 let consumer = MockConsumer::default();
973 let mut resolver =
974 start_resolver(context.child("resolver"), fetcher.clone(), consumer.clone());
975 let target = PrivateKey::from_seed(0).public_key();
976
977 assert!(resolver
978 .fetch_targeted(
979 Fetch {
980 key: 1,
981 subscriber: 10,
982 span: tracing::Span::none(),
983 },
984 non_empty_vec![target]
985 )
986 .accepted());
987 let delivery = wait_for_delivery(&context, &consumer).await;
988 assert_eq!(delivery.value, Bytes::from_static(b"value"));
989 delivery.response.send(true).expect("response dropped");
990 assert_eq!(fetcher.calls(), 1);
991 });
992 }
993}