zebra_network/peer_set/set/tests/
consumer.rs

1//! Real buffered peer-set and client-channel fixtures for downstream consumer tests.
2
3use std::{net::SocketAddr, time::Duration};
4
5use futures::{stream, StreamExt};
6use tower::{discover::Change, ServiceExt};
7use zebra_chain::{block, chain_tip::mock::MockChainTip, parameters::Network};
8
9use crate::{
10    constants::CURRENT_NETWORK_PROTOCOL_VERSION,
11    peer::{ClientTestHarness, ConnectedAddr, MinimumPeerVersion},
12    peer_set::InventoryChange,
13    types::{InventoryHash, PeerServices},
14    BoxError, NetworkService, PeerSocketAddr, Request, Response,
15};
16
17pub use crate::peer::ReceivedClientRequest;
18
19use super::{PeerSetBuilder, PeerSetGuard};
20
21/// Owns a buffered peer set and the remote endpoints of its real client channels.
22pub struct PeerSetHarness {
23    /// The service to pass to the production sync or inbound constructor.
24    pub network: NetworkService,
25    /// The connected addresses, in the same order as the supplied service flags.
26    pub addrs: Vec<PeerSocketAddr>,
27    peers: Vec<ClientTestHarness>,
28    guard: PeerSetGuard,
29    busy_requests: Vec<super::super::ResponseFuture>,
30}
31
32impl PeerSetHarness {
33    /// Constructs real clients, optionally filling one client's two request-channel slots.
34    ///
35    /// Inventory uses the existing one-entry test channel, so two updates without a poll
36    /// deterministically exercise registry lag. No socket or wall-clock timing is involved.
37    pub async fn new(services: &[PeerServices], busy_peer: Option<usize>) -> Self {
38        let mut guard = PeerSetGuard::new();
39        let inventory_receiver = guard.create_inventory_receiver();
40        let inventory_sender = guard.inventory_sender().as_ref().unwrap().clone();
41        let mut addrs = Vec::new();
42        let mut peers = Vec::new();
43        let mut changes: Vec<Result<_, BoxError>> = Vec::new();
44        for (index, services) in services.iter().enumerate() {
45            let host = u8::try_from(index + 1).expect("test peer count fits in an IPv4 octet");
46            let addr: PeerSocketAddr = SocketAddr::from(([127, 0, 0, host], 8233)).into();
47            let (mut client, handle) = ClientTestHarness::build()
48                .with_version(CURRENT_NETWORK_PROTOCOL_VERSION)
49                .with_services(*services)
50                .with_connected_addr(ConnectedAddr::new_inbound_direct(addr))
51                .finish();
52            client.inv_collector = inventory_sender.clone();
53            changes.push(Ok(Change::Insert(addr, client.into())));
54            addrs.push(addr);
55            peers.push(handle);
56        }
57        let (tip, _tip_sender) = MockChainTip::new();
58        let builder = PeerSetBuilder::new()
59            .with_discover(stream::iter(changes).chain(stream::pending()))
60            .with_minimum_peer_version(MinimumPeerVersion::new(tip, &Network::Mainnet));
61        let (mut peer_set, mut peer_guard) = PeerSetBuilder {
62            inv_stream: Some(inventory_receiver),
63            ..builder
64        }
65        .build();
66        // Keep the sender for the channel actually used by this peer set.
67        *peer_guard.inventory_sender() = guard.inventory_sender().take();
68        peer_set.ready().await.expect("test peers are ready");
69        let mut busy_requests = Vec::new();
70        if let Some(peer) = busy_peer {
71            for byte in [0xF0, 0xF1] {
72                peer_set.ready().await.expect("test peers remain connected");
73                busy_requests.push(peer_set.call_ready_peer(
74                    addrs[peer],
75                    Request::BlocksByHash([block::Hash([byte; 32])].into_iter().collect()),
76                ));
77            }
78            peer_set
79                .ready()
80                .await
81                .expect("an idle fallback peer remains ready");
82            assert!(peer_set.cancel_handles.contains_key(&addrs[peer]));
83        }
84        let (network, _poll_task) = peer_set.into_buffer(3);
85        Self {
86            network,
87            addrs,
88            peers,
89            guard: peer_guard,
90            busy_requests,
91        }
92    }
93
94    /// Receives an already-dispatched client request, without advancing the clock.
95    pub fn take_request(&mut self, peer: usize) -> Option<ReceivedClientRequest> {
96        self.peers[peer].take_request()
97    }
98
99    /// Waits at most one virtual second for the specified client to receive a request.
100    pub async fn next_request(&mut self, peer: usize) -> ReceivedClientRequest {
101        self.next_request_from(&[peer]).await.1
102    }
103
104    /// Receives the next request from any of the supplied peers, with a one-second bound.
105    pub async fn next_request_from(&mut self, peers: &[usize]) -> (usize, ReceivedClientRequest) {
106        tokio::time::timeout(Duration::from_secs(1), async {
107            loop {
108                for &peer in peers {
109                    if let Some(request) = self.take_request(peer) {
110                        return (peer, request);
111                    }
112                }
113                tokio::time::sleep(Duration::from_millis(1)).await;
114            }
115        })
116        .await
117        .expect("a selected peer must receive the consumer request")
118    }
119
120    /// Publishes an inventory event through the real lossy registry channel.
121    pub fn advertise(&mut self, peer: usize, hash: block::Hash) {
122        self.guard
123            .inventory_sender()
124            .as_ref()
125            .unwrap()
126            .send(InventoryChange::new_available(
127                InventoryHash::Block(hash),
128                self.addrs[peer],
129            ))
130            .expect("the real inventory registry is still connected");
131    }
132
133    /// Polls the production peer set through its normal buffered maintenance request.
134    pub async fn flush(&self) {
135        self.network
136            .clone()
137            .oneshot(Request::PollPeerSet)
138            .await
139            .expect("the test peer set is healthy");
140    }
141
142    /// Frees the two pre-filled client slots and completes their unrelated requests.
143    pub async fn release_busy_peer(&mut self, peer: usize) {
144        for byte in [0xF0, 0xF1] {
145            let request = self.next_request(peer).await;
146            assert_eq!(
147                request.request(),
148                &Request::BlocksByHash([block::Hash([byte; 32])].into_iter().collect()),
149            );
150            request.respond(Ok(Response::Nil));
151        }
152        for response in self.busy_requests.drain(..) {
153            response.await.expect("priming requests completed");
154        }
155    }
156}
zebra_network/peer_set/set/tests/consumer.rs

zebra_network/peer_set/set/tests/
consumer.rs
