rust_ipfs/p2p/
bitswap.rs

1mod handler;
2mod message;
3mod pb;
4mod prefix;
5mod session;
6mod wantlist;
7
8use std::{
9    collections::{HashMap, HashSet, VecDeque, hash_map::Entry},
10    fmt::Debug,
11    task::{Context, Poll, Waker},
12    time::{Duration, Instant},
13};
14
15use connexa::prelude::{
16    Multiaddr, PeerId,
17    swarm::{
18        ConnectionClosed, ConnectionDenied, ConnectionId, FromSwarm, NetworkBehaviour,
19        NotifyHandler, THandler, THandlerInEvent, THandlerOutEvent, ToSwarm,
20        behaviour::ConnectionEstablished, dial_opts::DialOpts,
21    },
22    transport::Endpoint,
23    transport::transport::PortUse,
24};
25use futures::{FutureExt, StreamExt};
26use futures_timer::Delay;
27use ipld_core::cid::Cid;
28
29use pollable_map::stream::StreamMap;
30
31mod bitswap_pb {
32    pub use super::pb::bitswap_pb::Message;
33    pub mod message {
34        use super::super::pb::bitswap_pb::mod_Message as message;
35        pub use message::Wantlist;
36        pub use message::mod_Wantlist as wantlist;
37        pub use message::{Block, BlockPresence, BlockPresenceType};
38    }
39}
40
41use self::{
42    message::RequestType,
43    session::{PeerSession, PeerSessionEvent, ServeBudget},
44    wantlist::{Wantlist, WantlistDriver, WantlistEvent},
45};
46use crate::repo::DefaultStorage;
47use crate::repo::Repo;
48
49const CAP_THRESHOLD: usize = 100;
50const DEFAULT_PRIORITY: i32 = 1;
51const BLOCK_REQUEST_TIMEOUT: Duration = Duration::from_secs(5);
52const BLOCK_TIMER_TICK: Duration = Duration::from_secs(1);
53
54#[derive(Debug)]
55pub enum Event {
56    NeedBlock { cid: Cid },
57    BlockRetrieved { cid: Cid },
58    CancelBlock { cid: Cid },
59}
60
61#[derive(Debug, Default, Clone, Copy)]
62pub struct BitswapStats {
63    pub blocks_sent: u64,
64    pub bytes_sent: u64,
65    pub blocks_received: u64,
66    pub bytes_received: u64,
67}
68
69#[derive(Debug, Clone, Copy)]
70pub struct Config {
71    pub global_serve_limit: usize,
72}
73
74impl Default for Config {
75    fn default() -> Self {
76        Self {
77            global_serve_limit: 1024,
78        }
79    }
80}
81
82pub struct Behaviour {
83    events: VecDeque<ToSwarm<Event, THandlerInEvent<Self>>>,
84    connections: HashMap<PeerId, HashSet<ConnectionId>>,
85    unsupported: HashSet<PeerId>,
86    store: Repo<DefaultStorage>,
87    wantlist: Wantlist,
88    driver: WantlistDriver,
89    sessions: StreamMap<PeerId, PeerSession>,
90    budget: ServeBudget,
91    candidates: HashMap<Cid, VecDeque<PeerId>>,
92    block_inflight: HashMap<Cid, (PeerId, Instant)>,
93    block_timer: Delay,
94    waker: Option<Waker>,
95}
96
97impl Behaviour {
98    pub fn new(store: &Repo<DefaultStorage>, config: Config) -> Self {
99        let wantlist = Wantlist::default();
100        Self {
101            events: VecDeque::new(),
102            connections: HashMap::new(),
103            unsupported: HashSet::new(),
104            store: store.clone(),
105            driver: WantlistDriver::new(wantlist.clone()),
106            wantlist,
107            sessions: StreamMap::new(),
108            budget: ServeBudget::new(config.global_serve_limit),
109            candidates: HashMap::new(),
110            block_inflight: HashMap::new(),
111            block_timer: Delay::new(BLOCK_TIMER_TICK),
112            waker: None,
113        }
114    }
115
116    pub fn get(&mut self, cid: &Cid, providers: &[PeerId], timeout: Option<Duration>) {
117        self.gets(vec![*cid], providers, timeout)
118    }
119
120    pub fn gets(&mut self, cids: Vec<Cid>, providers: &[PeerId], timeout: Option<Duration>) {
121        for cid in &cids {
122            self.wantlist
123                .want(*cid, RequestType::Have, DEFAULT_PRIORITY, timeout);
124        }
125
126        // A target is any peer we can broadcast to right now. Explicit providers that are not yet
127        // connected are dialed but do not count until the connection is up.
128        let mut have_target = false;
129        if providers.is_empty() {
130            have_target = self
131                .connections
132                .keys()
133                .any(|peer| !self.unsupported.contains(peer));
134        } else {
135            for peer in providers {
136                if self.unsupported.contains(peer) {
137                    continue;
138                }
139                if self.connections.contains_key(peer) {
140                    have_target = true;
141                } else {
142                    self.events.push_back(ToSwarm::Dial {
143                        opts: DialOpts::peer_id(*peer).build(),
144                    });
145                }
146            }
147        }
148
149        for session in self.sessions.values_mut() {
150            session.sync();
151        }
152
153        if !have_target {
154            for cid in cids {
155                self.events
156                    .push_back(ToSwarm::GenerateEvent(Event::NeedBlock { cid }));
157            }
158        }
159
160        self.wake();
161    }
162
163    pub fn local_wantlist(&self) -> Vec<Cid> {
164        self.wantlist.cids()
165    }
166
167    pub fn peer_wantlist(&self, peer_id: PeerId) -> Vec<Cid> {
168        self.sessions
169            .get(&peer_id)
170            .map(|session| session.peer_wantlist())
171            .unwrap_or_default()
172    }
173
174    // Note: This is called specifically to cancel the request and not just emitting a request
175    //       after receiving a request.
176    pub fn cancel(&mut self, cid: Cid) {
177        if !self.wantlist.cancel(&cid) {
178            return;
179        }
180
181        self.candidates.remove(&cid);
182        self.block_inflight.remove(&cid);
183        for session in self.sessions.values_mut() {
184            session.sync();
185        }
186
187        self.events
188            .push_back(ToSwarm::GenerateEvent(Event::CancelBlock { cid }));
189        self.wake();
190    }
191
192    // Notify connected peers that asked for these blocks that we now have them.
193    pub fn notify_new_blocks(&mut self, cid: impl IntoIterator<Item = Cid>) {
194        let blocks = cid.into_iter().collect::<Vec<_>>();
195        if blocks.is_empty() {
196            return;
197        }
198
199        for session in self.sessions.values_mut() {
200            session.serve_wanted(&blocks);
201        }
202
203        self.wake();
204    }
205
206    pub fn peers(&self) -> Vec<PeerId> {
207        self.sessions.keys().copied().collect()
208    }
209
210    pub fn stats(&self) -> BitswapStats {
211        let mut stats = BitswapStats::default();
212        for session in self.sessions.values() {
213            let ledger = session.ledger();
214            stats.blocks_sent += ledger.blocks_sent;
215            stats.bytes_sent += ledger.bytes_sent;
216            stats.blocks_received += ledger.blocks_recv;
217            stats.bytes_received += ledger.bytes_recv;
218        }
219        stats
220    }
221
222    fn on_have(&mut self, peer: PeerId, cid: Cid) {
223        if !self.wantlist.contains(&cid) {
224            return;
225        }
226        if self.block_inflight.get(&cid).map(|(p, _)| *p) == Some(peer) {
227            return;
228        }
229        let queue = self.candidates.entry(cid).or_default();
230        if !queue.contains(&peer) {
231            queue.push_back(peer);
232        }
233        if !self.block_inflight.contains_key(&cid) {
234            self.choose_block_source(cid);
235        }
236    }
237
238    fn on_dont_have(&mut self, peer: PeerId, cid: Cid) {
239        if let Some(queue) = self.candidates.get_mut(&cid) {
240            queue.retain(|p| *p != peer);
241        }
242        if self.block_inflight.get(&cid).map(|(p, _)| *p) == Some(peer) {
243            self.block_inflight.remove(&cid);
244            self.choose_block_source(cid);
245        }
246    }
247
248    fn choose_block_source(&mut self, cid: Cid) {
249        while let Some(peer) = self.candidates.get_mut(&cid).and_then(|q| q.pop_front()) {
250            let Some(session) = self.sessions.get_mut(&peer) else {
251                continue;
252            };
253            let Some(message) = session.request_block(cid) else {
254                continue;
255            };
256            self.block_inflight.insert(cid, (peer, Instant::now()));
257            self.events.push_back(ToSwarm::NotifyHandler {
258                peer_id: peer,
259                handler: NotifyHandler::Any,
260                event: handler::FromBehaviour::Send(message),
261            });
262            self.wake();
263            return;
264        }
265        self.candidates.remove(&cid);
266        self.block_inflight.remove(&cid);
267        if self.wantlist.contains(&cid) {
268            self.wantlist.rearm_discovery(&cid);
269        }
270    }
271
272    fn failover_stale(&mut self) {
273        let now = Instant::now();
274        let stale: Vec<(Cid, PeerId)> = self
275            .block_inflight
276            .iter()
277            .filter(|(_, (_, requested))| now.duration_since(*requested) >= BLOCK_REQUEST_TIMEOUT)
278            .map(|(cid, (peer, _))| (*cid, *peer))
279            .collect();
280        for (cid, peer) in stale {
281            self.block_inflight.remove(&cid);
282            if let Some(session) = self.sessions.get_mut(&peer) {
283                session.reset_block(cid);
284            }
285            self.choose_block_source(cid);
286        }
287    }
288
289    fn on_peer_gone(&mut self, peer: PeerId) {
290        for queue in self.candidates.values_mut() {
291            queue.retain(|p| *p != peer);
292        }
293        let affected: Vec<Cid> = self
294            .block_inflight
295            .iter()
296            .filter(|(_, (p, _))| *p == peer)
297            .map(|(cid, _)| *cid)
298            .collect();
299        for cid in affected {
300            self.block_inflight.remove(&cid);
301            self.choose_block_source(cid);
302        }
303    }
304
305    fn wake(&mut self) {
306        if let Some(waker) = self.waker.take() {
307            waker.wake();
308        }
309    }
310
311    fn ensure_session(&mut self, peer_id: PeerId) {
312        if self.unsupported.contains(&peer_id) || self.sessions.contains_key(&peer_id) {
313            return;
314        }
315        let session =
316            PeerSession::new(self.wantlist.clone(), self.store.clone(), self.budget.clone());
317        self.sessions.insert(peer_id, session);
318        self.wake();
319    }
320
321    fn on_connection_established(
322        &mut self,
323        ConnectionEstablished {
324            connection_id,
325            peer_id,
326            ..
327        }: ConnectionEstablished,
328    ) {
329        self.connections
330            .entry(peer_id)
331            .or_default()
332            .insert(connection_id);
333    }
334
335    fn on_connection_close(
336        &mut self,
337        ConnectionClosed {
338            connection_id,
339            peer_id,
340            remaining_established,
341            ..
342        }: ConnectionClosed,
343    ) {
344        if let Entry::Occupied(mut entry) = self.connections.entry(peer_id) {
345            let list = entry.get_mut();
346            list.remove(&connection_id);
347            if list.is_empty() {
348                entry.remove();
349            }
350        }
351
352        if remaining_established == 0 {
353            self.sessions.remove(&peer_id);
354            self.unsupported.remove(&peer_id);
355            self.on_peer_gone(peer_id);
356        }
357    }
358}
359
360impl NetworkBehaviour for Behaviour {
361    type ConnectionHandler = handler::Handler;
362    type ToSwarm = Event;
363
364    fn handle_pending_inbound_connection(
365        &mut self,
366        _: ConnectionId,
367        _: &Multiaddr,
368        _: &Multiaddr,
369    ) -> Result<(), ConnectionDenied> {
370        Ok(())
371    }
372
373    fn handle_pending_outbound_connection(
374        &mut self,
375        _: ConnectionId,
376        _: Option<PeerId>,
377        _: &[Multiaddr],
378        _: Endpoint,
379    ) -> Result<Vec<Multiaddr>, ConnectionDenied> {
380        Ok(vec![])
381    }
382
383    fn handle_established_inbound_connection(
384        &mut self,
385        _: ConnectionId,
386        _: PeerId,
387        _: &Multiaddr,
388        _: &Multiaddr,
389    ) -> Result<THandler<Self>, ConnectionDenied> {
390        Ok(handler::Handler::default())
391    }
392
393    fn handle_established_outbound_connection(
394        &mut self,
395        _: ConnectionId,
396        _: PeerId,
397        _: &Multiaddr,
398        _: Endpoint,
399        _: PortUse,
400    ) -> Result<THandler<Self>, ConnectionDenied> {
401        Ok(handler::Handler::default())
402    }
403
404    fn on_connection_handler_event(
405        &mut self,
406        peer_id: PeerId,
407        _connection_id: ConnectionId,
408        event: THandlerOutEvent<Self>,
409    ) {
410        match event {
411            handler::ToBehaviour::Ready => self.ensure_session(peer_id),
412            handler::ToBehaviour::Message(message) => {
413                self.ensure_session(peer_id);
414                if let Some(session) = self.sessions.get_mut(&peer_id) {
415                    session.on_message(message);
416                }
417                self.wake();
418            }
419            handler::ToBehaviour::Unsupported => {
420                self.unsupported.insert(peer_id);
421                self.sessions.remove(&peer_id);
422                self.on_peer_gone(peer_id);
423            }
424            handler::ToBehaviour::Failed => {
425                self.sessions.remove(&peer_id);
426                self.on_peer_gone(peer_id);
427            }
428        }
429    }
430
431    fn on_swarm_event(&mut self, event: FromSwarm) {
432        match event {
433            FromSwarm::ConnectionEstablished(event) => self.on_connection_established(event),
434            FromSwarm::ConnectionClosed(event) => self.on_connection_close(event),
435            _ => {}
436        }
437    }
438
439    fn poll(&mut self, ctx: &mut Context) -> Poll<ToSwarm<Self::ToSwarm, THandlerInEvent<Self>>> {
440        if let Some(event) = self.events.pop_front() {
441            return Poll::Ready(event);
442        } else if self.events.capacity() > CAP_THRESHOLD {
443            self.events.shrink_to_fit();
444        }
445
446        while let Poll::Ready(Some(event)) = self.driver.poll_next_unpin(ctx) {
447            match event {
448                WantlistEvent::NeedProviders(cid) => {
449                    return Poll::Ready(ToSwarm::GenerateEvent(Event::NeedBlock { cid }));
450                }
451                WantlistEvent::Expired(cid) => {
452                    self.candidates.remove(&cid);
453                    self.block_inflight.remove(&cid);
454                    for session in self.sessions.values_mut() {
455                        session.sync();
456                    }
457                    return Poll::Ready(ToSwarm::GenerateEvent(Event::CancelBlock { cid }));
458                }
459                WantlistEvent::Rebroadcast => {
460                    for session in self.sessions.values_mut() {
461                        session.sync();
462                    }
463                }
464            }
465        }
466
467        while let Poll::Ready(Some((peer_id, event))) = self.sessions.poll_next_unpin(ctx) {
468            match event {
469                PeerSessionEvent::Send(message) => {
470                    return Poll::Ready(ToSwarm::NotifyHandler {
471                        peer_id,
472                        handler: NotifyHandler::Any,
473                        event: handler::FromBehaviour::Send(message),
474                    });
475                }
476                PeerSessionEvent::Have(cid) => self.on_have(peer_id, cid),
477                PeerSessionEvent::DontHave(cid) => self.on_dont_have(peer_id, cid),
478                PeerSessionEvent::Stored(cid) => {
479                    self.candidates.remove(&cid);
480                    self.block_inflight.remove(&cid);
481                    self.wantlist.cancel(&cid);
482                    for session in self.sessions.values_mut() {
483                        session.sync();
484                    }
485                    return Poll::Ready(ToSwarm::GenerateEvent(Event::BlockRetrieved { cid }));
486                }
487            }
488        }
489
490        while let Poll::Ready(()) = self.block_timer.poll_unpin(ctx) {
491            self.block_timer.reset(BLOCK_TIMER_TICK);
492            self.failover_stale();
493        }
494
495        if let Some(event) = self.events.pop_front() {
496            return Poll::Ready(event);
497        }
498
499        self.waker = Some(ctx.waker().clone());
500
501        Poll::Pending
502    }
503}
504
505#[cfg(test)]
506mod test {
507    use std::time::Duration;
508
509    use crate::{block::BlockCodec, repo::DefaultStorage};
510    use connexa::prelude::{
511        Multiaddr, PeerId,
512        swarm::{NetworkBehaviour, Swarm, SwarmBuilder, SwarmEvent, dial_opts::DialOpts},
513        transport::{
514            noise,
515            transport::{MemoryTransport, Transport},
516            upgrade::Version,
517            yamux,
518        },
519    };
520    use futures::StreamExt;
521    use ipld_core::cid::Cid;
522    use multihash_codetable::{Code, MultihashDigest};
523
524    use crate::{Block, repo::Repo};
525
526    fn create_block() -> Block {
527        let data = b"hello block\n".to_vec();
528        let cid = Cid::new_v1(BlockCodec::Raw.into(), Code::Sha2_256.digest(&data));
529
530        Block::new_unchecked(cid, data)
531    }
532
533    async fn wait_on_connection(
534        swarm1: &mut Swarm<Behaviour>,
535        swarm2: &mut Swarm<Behaviour>,
536        peer_id: PeerId,
537    ) {
538        loop {
539            futures::select! {
540                event = swarm1.select_next_some() => {
541                    if let SwarmEvent::ConnectionEstablished { peer_id: peer, .. } = event {
542                        assert_eq!(peer, peer_id);
543                        break;
544                    }
545                }
546                _ = swarm2.next() => {}
547            }
548        }
549    }
550
551    #[tokio::test]
552    async fn exchange_blocks() -> anyhow::Result<()> {
553        let (_, _, mut swarm1, repo) = build_swarm().await;
554        let (peer2, addr2, mut swarm2, repo2) = build_swarm().await;
555
556        let block = create_block();
557
558        let cid = *block.cid();
559
560        repo.put_block(&block).await?;
561
562        let opt = DialOpts::peer_id(peer2)
563            .addresses(vec![addr2.clone()])
564            .build();
565
566        swarm1.dial(opt)?;
567
568        wait_on_connection(&mut swarm1, &mut swarm2, peer2).await;
569
570        swarm2.behaviour_mut().bitswap.get(&cid, &[], None);
571
572        loop {
573            tokio::select! {
574                _ = swarm1.next() => {}
575                e = swarm2.select_next_some() => {
576                    match e {
577                        SwarmEvent::Behaviour(BehaviourEvent::Bitswap(super::Event::BlockRetrieved { cid: inner_cid })) => {
578                            assert_eq!(inner_cid, cid);
579                        }
580                        SwarmEvent::Behaviour(BehaviourEvent::Bitswap(super::Event::CancelBlock { cid: inner_cid })) => {
581                            assert_eq!(inner_cid, cid);
582                            unreachable!("exchange should not timeout");
583                        }
584                        _ => {}
585                    }
586                },
587                Ok(true) = repo2.contains(&cid) => {
588                    break;
589                }
590            }
591        }
592
593        let b = repo2
594            .get_block_now(&cid)
595            .await
596            .unwrap()
597            .expect("block exist");
598
599        assert_eq!(b, block);
600
601        Ok(())
602    }
603
604    #[tokio::test]
605    async fn notify_swarm() -> anyhow::Result<()> {
606        let (_, _, mut swarm1, _) = build_swarm().await;
607
608        let block = create_block();
609
610        let cid = *block.cid();
611
612        swarm1
613            .behaviour_mut()
614            .bitswap
615            .get(&cid, &[], Some(Duration::from_millis(500)));
616
617        let mut notified_counter = 0;
618
619        loop {
620            tokio::select! {
621                e = swarm1.select_next_some() => {
622                    match e {
623                        SwarmEvent::Behaviour(BehaviourEvent::Bitswap(super::Event::NeedBlock { cid: inner_cid })) => {
624                            assert_eq!(inner_cid, cid);
625                            notified_counter += 1;
626                        }
627                        SwarmEvent::Behaviour(BehaviourEvent::Bitswap(super::Event::CancelBlock { cid: inner_cid })) => {
628                            assert_eq!(inner_cid, cid);
629                            unreachable!()
630                        }
631                        _ => {}
632                    }
633                },
634            }
635
636            if notified_counter == 2 {
637                break;
638            }
639        }
640
641        Ok(())
642    }
643
644    #[tokio::test]
645    async fn bitswap_timeout() -> anyhow::Result<()> {
646        let (_, _, mut swarm1, _) = build_swarm().await;
647        let (peer2, addr2, mut swarm2, _) = build_swarm().await;
648
649        let block = create_block();
650
651        let cid = *block.cid();
652
653        let opt = DialOpts::peer_id(peer2)
654            .addresses(vec![addr2.clone()])
655            .build();
656
657        swarm1.dial(opt)?;
658
659        wait_on_connection(&mut swarm1, &mut swarm2, peer2).await;
660
661        swarm2
662            .behaviour_mut()
663            .bitswap
664            .get(&cid, &[], Some(Duration::from_millis(150)));
665
666        loop {
667            tokio::select! {
668                _ = swarm1.next() => {}
669                e = swarm2.select_next_some() => {
670                    if let SwarmEvent::Behaviour(BehaviourEvent::Bitswap(super::Event::CancelBlock { cid: inner_cid })) = e {
671                        assert_eq!(inner_cid, cid);
672                        break;
673                    }
674                },
675            }
676        }
677
678        Ok(())
679    }
680
681    #[tokio::test]
682    async fn exchange_blocks_with_explicit_peer() -> anyhow::Result<()> {
683        let (peer1, _, mut swarm1, repo) = build_swarm().await;
684        let (peer2, addr2, mut swarm2, repo2) = build_swarm().await;
685
686        let block = create_block();
687
688        let cid = *block.cid();
689
690        repo.put_block(&block).await?;
691
692        let opt = DialOpts::peer_id(peer2)
693            .addresses(vec![addr2.clone()])
694            .build();
695
696        swarm1.dial(opt)?;
697
698        wait_on_connection(&mut swarm1, &mut swarm2, peer2).await;
699
700        swarm2.behaviour_mut().bitswap.get(&cid, &[peer1], None);
701
702        loop {
703            tokio::select! {
704                _ = swarm1.next() => {}
705                e = swarm2.select_next_some() => {
706                    if let SwarmEvent::Behaviour(BehaviourEvent::Bitswap(super::Event::BlockRetrieved { cid: inner_cid })) = e {
707                        assert_eq!(inner_cid, cid);
708                    }
709                },
710                Ok(true) = repo2.contains(&cid) => {
711                    break;
712                }
713            }
714        }
715
716        let b = repo2
717            .get_block_now(&cid)
718            .await
719            .unwrap()
720            .expect("block exist");
721
722        assert_eq!(b, block);
723
724        Ok(())
725    }
726
727    #[tokio::test]
728    async fn notify_after_block_exchange() -> anyhow::Result<()> {
729        let (peer1, _, mut swarm1, repo) = build_swarm().await;
730        let (peer2, addr2, mut swarm2, _) = build_swarm().await;
731        let (peer3, addr3, mut swarm3, repo3) = build_swarm().await;
732
733        let block = create_block();
734
735        let cid = *block.cid();
736
737        repo.put_block(&block).await?;
738
739        let opt = DialOpts::peer_id(peer2)
740            .addresses(vec![addr2.clone()])
741            .build();
742        swarm1.dial(opt)?;
743
744        let opt = DialOpts::peer_id(peer3)
745            .addresses(vec![addr3.clone()])
746            .build();
747
748        swarm2.dial(opt)?;
749        let mut peer_1_connected = false;
750        let mut peer_2_connected = false;
751        let mut peer_3_connected = false;
752
753        loop {
754            futures::select! {
755                event = swarm1.select_next_some() => {
756                    if let SwarmEvent::ConnectionEstablished { .. } = event {
757                        peer_1_connected = true;
758                    }
759                }
760                event = swarm2.select_next_some() => {
761                    if let SwarmEvent::ConnectionEstablished { peer_id, .. } = event {
762                        if peer_id == peer1 {
763                            peer_2_connected = true;
764                        }
765                    }
766                }
767
768                event = swarm3.select_next_some() => {
769                    if let SwarmEvent::ConnectionEstablished { peer_id, .. } = event {
770                        assert_eq!(peer_id, peer2);
771                        peer_3_connected = true;
772                    }
773                }
774            }
775            if peer_1_connected && peer_2_connected && peer_3_connected {
776                break;
777            }
778        }
779        swarm2.behaviour_mut().bitswap.get(&cid, &[peer1], None);
780        swarm3.behaviour_mut().bitswap.get(&cid, &[], None);
781
782        loop {
783            tokio::select! {
784                _ = swarm1.next() => {}
785                e = swarm2.select_next_some() => {
786                    if let SwarmEvent::Behaviour(BehaviourEvent::Bitswap(super::Event::BlockRetrieved { cid: inner_cid })) = e {
787                        assert_eq!(inner_cid, cid);
788                        swarm2.behaviour_mut().bitswap.notify_new_blocks(std::iter::once(cid));
789                    }
790                },
791                e = swarm3.select_next_some() => {
792                    if let SwarmEvent::Behaviour(BehaviourEvent::Bitswap(super::Event::BlockRetrieved { cid: inner_cid })) = e {
793                        assert_eq!(inner_cid, cid);
794                        break;
795                    }
796                },
797            }
798        }
799
800        let b = repo3
801            .get_block_now(&cid)
802            .await
803            .unwrap()
804            .expect("block exist");
805
806        assert_eq!(b, block);
807
808        Ok(())
809    }
810
811    #[tokio::test]
812    async fn cancel_block_exchange() -> anyhow::Result<()> {
813        let (_, _, mut swarm1, _) = build_swarm().await;
814
815        let block = create_block();
816
817        let cid = *block.cid();
818
819        swarm1.behaviour_mut().bitswap.get(&cid, &[], None);
820        swarm1.behaviour_mut().bitswap.cancel(cid);
821
822        loop {
823            tokio::select! {
824                e = swarm1.select_next_some() => {
825                    if let SwarmEvent::Behaviour(BehaviourEvent::Bitswap(super::Event::CancelBlock { cid: inner_cid })) = e {
826                        assert_eq!(inner_cid, cid);
827                        break;
828                    }
829                },
830            }
831        }
832
833        Ok(())
834    }
835
836    #[tokio::test]
837    async fn local_wantlist() -> anyhow::Result<()> {
838        let (_, _, mut swarm1, _) = build_swarm().await;
839
840        let block = create_block();
841
842        let cid = *block.cid();
843
844        swarm1.behaviour_mut().bitswap.get(&cid, &[], None);
845
846        let list = swarm1.behaviour().bitswap.local_wantlist();
847
848        assert_eq!(list[0], cid);
849
850        Ok(())
851    }
852
853    #[tokio::test]
854    async fn peer_wantlist() -> anyhow::Result<()> {
855        let (peer1, _, mut swarm1, _) = build_swarm().await;
856        let (peer2, addr2, mut swarm2, _) = build_swarm().await;
857
858        let block = create_block();
859
860        let cid = *block.cid();
861
862        let opt = DialOpts::peer_id(peer2)
863            .addresses(vec![addr2.clone()])
864            .build();
865        swarm1.dial(opt)?;
866
867        let mut peer_1_connected = false;
868        let mut peer_2_connected = false;
869
870        loop {
871            futures::select! {
872                event = swarm1.select_next_some() => {
873                    if let SwarmEvent::ConnectionEstablished { .. } = event {
874                        peer_1_connected = true;
875                    }
876                }
877                event = swarm2.select_next_some() => {
878                    if let SwarmEvent::ConnectionEstablished { peer_id, .. } = event {
879                        if peer_id == peer1 {
880                            peer_2_connected = true;
881                        }
882                    }
883                }
884            }
885            if peer_1_connected && peer_2_connected {
886                break;
887            }
888        }
889        swarm2.behaviour_mut().bitswap.get(&cid, &[peer1], None);
890
891        loop {
892            tokio::select! {
893                _ = swarm1.next() => {}
894                e = swarm2.select_next_some() => {
895                    if let SwarmEvent::Behaviour(BehaviourEvent::Bitswap(super::Event::NeedBlock { cid: inner_cid })) = e {
896                        assert_eq!(inner_cid, cid);
897                        break;
898                    }
899                },
900            }
901        }
902
903        let list = swarm1.behaviour().bitswap.peer_wantlist(peer2);
904        assert_eq!(list[0], cid);
905
906        Ok(())
907    }
908
909    #[tokio::test]
910    async fn single_block_fetch_across_seeders() -> anyhow::Result<()> {
911        let (peer_a, addr_a, mut swarm_a, repo_a) = build_swarm().await;
912        let (peer_b, addr_b, mut swarm_b, repo_b) = build_swarm().await;
913        let (_, _, mut swarm_f, repo_f) = build_swarm().await;
914
915        let block = create_block();
916        let cid = *block.cid();
917        repo_a.put_block(&block).await?;
918        repo_b.put_block(&block).await?;
919
920        swarm_f.dial(DialOpts::peer_id(peer_a).addresses(vec![addr_a]).build())?;
921        swarm_f.dial(DialOpts::peer_id(peer_b).addresses(vec![addr_b]).build())?;
922
923        let mut connected = std::collections::HashSet::new();
924        while connected.len() < 2 {
925            tokio::select! {
926                _ = swarm_a.next() => {}
927                _ = swarm_b.next() => {}
928                e = swarm_f.select_next_some() => {
929                    if let SwarmEvent::ConnectionEstablished { peer_id, .. } = e {
930                        connected.insert(peer_id);
931                    }
932                }
933            }
934        }
935
936        swarm_f.behaviour_mut().bitswap.get(&cid, &[], None);
937
938        loop {
939            tokio::select! {
940                _ = swarm_a.next() => {}
941                _ = swarm_b.next() => {}
942                _ = swarm_f.next() => {}
943                Ok(true) = repo_f.contains(&cid) => break,
944            }
945        }
946
947        let settle = tokio::time::sleep(Duration::from_millis(300));
948        tokio::pin!(settle);
949        loop {
950            tokio::select! {
951                _ = swarm_a.next() => {}
952                _ = swarm_b.next() => {}
953                _ = swarm_f.next() => {}
954                _ = &mut settle => break,
955            }
956        }
957
958        let sent = swarm_a.behaviour().bitswap.stats().blocks_sent
959            + swarm_b.behaviour().bitswap.stats().blocks_sent;
960        assert_eq!(sent, 1, "exactly one seeder should send the block");
961        assert_eq!(swarm_f.behaviour().bitswap.stats().blocks_received, 1);
962
963        Ok(())
964    }
965
966    #[tokio::test]
967    async fn ledger_counts_transfers() -> anyhow::Result<()> {
968        let (peer_a, addr_a, mut swarm_a, repo_a) = build_swarm().await;
969        let (_, _, mut swarm_f, repo_f) = build_swarm().await;
970
971        let block = create_block();
972        let cid = *block.cid();
973        let len = block.data().len() as u64;
974        repo_a.put_block(&block).await?;
975
976        swarm_f.dial(DialOpts::peer_id(peer_a).addresses(vec![addr_a]).build())?;
977        wait_on_connection(&mut swarm_f, &mut swarm_a, peer_a).await;
978
979        swarm_f.behaviour_mut().bitswap.get(&cid, &[], None);
980
981        loop {
982            tokio::select! {
983                _ = swarm_a.next() => {}
984                _ = swarm_f.next() => {}
985                Ok(true) = repo_f.contains(&cid) => break,
986            }
987        }
988
989        let seeder = swarm_a.behaviour().bitswap.stats();
990        assert_eq!(seeder.blocks_sent, 1);
991        assert_eq!(seeder.bytes_sent, len);
992
993        let fetcher = swarm_f.behaviour().bitswap.stats();
994        assert_eq!(fetcher.blocks_received, 1);
995        assert_eq!(fetcher.bytes_received, len);
996        assert!(!swarm_f.behaviour().bitswap.peers().is_empty());
997
998        Ok(())
999    }
1000
1001    #[tokio::test]
1002    async fn fetch_succeeds_when_one_peer_lacks_block() -> anyhow::Result<()> {
1003        let (peer_a, addr_a, mut swarm_a, _repo_a) = build_swarm().await;
1004        let (peer_b, addr_b, mut swarm_b, repo_b) = build_swarm().await;
1005        let (_, _, mut swarm_f, repo_f) = build_swarm().await;
1006
1007        let block = create_block();
1008        let cid = *block.cid();
1009        repo_b.put_block(&block).await?;
1010
1011        swarm_f.dial(DialOpts::peer_id(peer_a).addresses(vec![addr_a]).build())?;
1012        swarm_f.dial(DialOpts::peer_id(peer_b).addresses(vec![addr_b]).build())?;
1013
1014        let mut connected = std::collections::HashSet::new();
1015        while connected.len() < 2 {
1016            tokio::select! {
1017                _ = swarm_a.next() => {}
1018                _ = swarm_b.next() => {}
1019                e = swarm_f.select_next_some() => {
1020                    if let SwarmEvent::ConnectionEstablished { peer_id, .. } = e {
1021                        connected.insert(peer_id);
1022                    }
1023                }
1024            }
1025        }
1026
1027        swarm_f.behaviour_mut().bitswap.get(&cid, &[], None);
1028
1029        loop {
1030            tokio::select! {
1031                _ = swarm_a.next() => {}
1032                _ = swarm_b.next() => {}
1033                _ = swarm_f.next() => {}
1034                Ok(true) = repo_f.contains(&cid) => break,
1035            }
1036        }
1037
1038        assert_eq!(swarm_b.behaviour().bitswap.stats().blocks_sent, 1);
1039
1040        Ok(())
1041    }
1042
1043    async fn build_swarm() -> (PeerId, Multiaddr, Swarm<Behaviour>, Repo<DefaultStorage>) {
1044        let repo = Repo::new_memory();
1045
1046        let mut swarm = SwarmBuilder::with_new_identity()
1047            .with_tokio()
1048            .with_other_transport(|kp| {
1049                MemoryTransport::default()
1050                    .upgrade(Version::V1)
1051                    .authenticate(noise::Config::new(kp).expect("valid config"))
1052                    .multiplex(yamux::Config::default())
1053                    .timeout(Duration::from_secs(20))
1054                    .boxed()
1055            })
1056            .expect("")
1057            .with_behaviour(|_| Behaviour {
1058                bitswap: super::Behaviour::new(&repo, super::Config::default()),
1059                address_book: crate::p2p::addressbook::Behaviour::with_config(
1060                    crate::p2p::addressbook::Config {
1061                        store_on_connection: true,
1062                        ..Default::default()
1063                    },
1064                ),
1065            })
1066            .expect("")
1067            .with_swarm_config(|c| c.with_idle_connection_timeout(Duration::from_secs(30)))
1068            .build();
1069
1070        Swarm::listen_on(&mut swarm, "/memory/0".parse().unwrap()).unwrap();
1071
1072        if let Some(SwarmEvent::NewListenAddr { address, .. }) = swarm.next().await {
1073            let peer_id = swarm.local_peer_id();
1074            return (*peer_id, address, swarm, repo);
1075        }
1076
1077        unreachable!()
1078    }
1079
1080    #[derive(NetworkBehaviour)]
1081    #[behaviour(prelude = "connexa::prelude::swarm::derive_prelude")]
1082    struct Behaviour {
1083        bitswap: super::Behaviour,
1084        address_book: crate::p2p::addressbook::Behaviour,
1085    }
1086}
rust_ipfs/p2p/bitswap.rs

rust_ipfs/p2p/
bitswap.rs
