Add PeerSet readiness and request future cancel-safety tests (#3252)

* add a PeerSet drop test

* replace all `now_or_never()` in test

* add a drop of the ready future to test

* make sure requests always go to client

* fix imports and runtime

* add a peer sent hang test

* replace `std::mem::discriminant` calls

* replace `unreachable` calls

Co-authored-by: Janito Vaqueiro Ferreira Filho <janito.vff@gmail.com>

* move comments

Co-authored-by: Janito Vaqueiro Ferreira Filho <janito.vff@gmail.com>

* add `yield_now` call

Co-authored-by: Janito Vaqueiro Ferreira Filho <janito.vff@gmail.com>

* increase timeout but make the test fast pausing the runtime

Co-authored-by: Janito Vaqueiro Ferreira Filho <janito.vff@gmail.com>

* apply last fixes

Co-authored-by: Janito Vaqueiro Ferreira Filho <janito.vff@gmail.com>

Co-authored-by: Janito Vaqueiro Ferreira Filho <janito.vff@gmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

zebra-network/src/peer_set/set/tests.rs

@@ -31,6 +31,9 @@ use crate::{
 #[cfg(test)]
 mod prop;
 
+#[cfg(test)]
+mod vectors;
+
 /// The maximum number of arbitrary peers to generate in [`PeerVersions`].
 ///
 /// This affects the maximum number of peer connections added to the [`PeerSet`] during the tests.

zebra-network/src/peer_set/set/tests/vectors.rs

@@ -0,0 +1,166 @@
+use std::time::Duration;
+use tokio::time::timeout;
+use tower::{Service, ServiceExt};
+
+use zebra_chain::parameters::{Network, NetworkUpgrade};
+
+use super::{PeerSetBuilder, PeerVersions};
+use crate::{
+    peer::{ClientRequest, MinimumPeerVersion},
+    protocol::external::types::Version,
+    Request,
+};
+
+#[test]
+fn peer_set_ready_single_connection() {
+    // We are going to use just one peer version in this test
+    let peer_versions = PeerVersions {
+        peer_versions: vec![Version::min_specified_for_upgrade(
+            Network::Mainnet,
+            NetworkUpgrade::Canopy,
+        )],
+    };
+
+    // Start the runtime
+    let runtime = zebra_test::init_async();
+    let _guard = runtime.enter();
+
+    // Get peers and client handles of them
+    let (discovered_peers, handles) = peer_versions.mock_peer_discovery();
+    let (minimum_peer_version, _best_tip_height) =
+        MinimumPeerVersion::with_mock_chain_tip(Network::Mainnet);
+
+    // We will just use the first peer handle
+    let mut client_handle = handles
+        .into_iter()
+        .next()
+        .expect("we always have at least one client");
+
+    // Client did not received anything yet
+    assert!(client_handle
+        .try_to_receive_outbound_client_request()
+        .is_empty());
+
+    runtime.block_on(async move {
+        // Build a peerset
+        let (mut peer_set, _peer_set_guard) = PeerSetBuilder::new()
+            .with_discover(discovered_peers)
+            .with_minimum_peer_version(minimum_peer_version.clone())
+            .build();
+
+        // Get a ready future
+        let peer_ready_future = peer_set.ready();
+
+        // Drop the future
+        std::mem::drop(peer_ready_future);
+
+        // Peer set will remain ready for requests
+        let peer_ready1 = peer_set
+            .ready()
+            .await
+            .expect("peer set service is always ready");
+
+        // Make sure the client did not received anything yet
+        assert!(client_handle
+            .try_to_receive_outbound_client_request()
+            .is_empty());
+
+        // Make a call to the peer set that returns a future
+        let fut = peer_ready1.call(Request::Peers);
+
+        // Client received the request
+        assert!(matches!(
+            client_handle
+                .try_to_receive_outbound_client_request()
+                .request(),
+            Some(ClientRequest {
+                request: Request::Peers,
+                ..
+            })
+        ));
+
+        // Drop the future
+        std::mem::drop(fut);
+
+        // Peer set will remain ready for requests
+        let peer_ready2 = peer_set
+            .ready()
+            .await
+            .expect("peer set service is always ready");
+
+        // Get a new future calling a different request than before
+        let _fut = peer_ready2.call(Request::MempoolTransactionIds);
+
+        // Client received the request
+        assert!(matches!(
+            client_handle
+                .try_to_receive_outbound_client_request()
+                .request(),
+            Some(ClientRequest {
+                request: Request::MempoolTransactionIds,
+                ..
+            })
+        ));
+    });
+}
+
+#[test]
+fn peer_set_ready_multiple_connections() {
+    // Use three peers with the same version
+    let peer_version = Version::min_specified_for_upgrade(Network::Mainnet, NetworkUpgrade::Canopy);
+    let peer_versions = PeerVersions {
+        peer_versions: vec![peer_version, peer_version, peer_version],
+    };
+
+    // Start the runtime
+    let runtime = zebra_test::init_async();
+    let _guard = runtime.enter();
+
+    // Pause the runtime's timer so that it advances automatically.
+    //
+    // CORRECTNESS: This test does not depend on external resources that could really timeout, like
+    // real network connections.
+    tokio::time::pause();
+
+    // Get peers and client handles of them
+    let (discovered_peers, handles) = peer_versions.mock_peer_discovery();
+    let (minimum_peer_version, _best_tip_height) =
+        MinimumPeerVersion::with_mock_chain_tip(Network::Mainnet);
+
+    // Make sure we have the right number of peers
+    assert_eq!(handles.len(), 3);
+
+    runtime.block_on(async move {
+        // Build a peerset
+        let (mut peer_set, _peer_set_guard) = PeerSetBuilder::new()
+            .with_discover(discovered_peers)
+            .with_minimum_peer_version(minimum_peer_version.clone())
+            .build();
+
+        // Get peerset ready
+        let peer_ready = peer_set
+            .ready()
+            .await
+            .expect("peer set service is always ready");
+
+        // Check we have the right amount of ready services
+        assert_eq!(peer_ready.ready_services.len(), 3);
+
+        // Stop some peer connections but not all
+        handles[0].stop_connection_task().await;
+        handles[1].stop_connection_task().await;
+
+        // We can still make the peer set ready
+        peer_set
+            .ready()
+            .await
+            .expect("peer set service is always ready");
+
+        // Stop the connection of the last peer
+        handles[2].stop_connection_task().await;
+
+        // Peer set hangs when no more connections are present
+        let peer_ready = peer_set.ready();
+        assert!(timeout(Duration::from_secs(10), peer_ready).await.is_err());
+    });
+}