From 1f756fcc81ed7d6343d8bce4052e373f65c1c6ae Mon Sep 17 00:00:00 2001 From: Janito Vaqueiro Ferreira Filho Date: Wed, 8 Dec 2021 21:18:17 -0300 Subject: [PATCH] Add `zebra_test::init_async` helper function (#3169) * Use a single-thread shared Tokio runtime This allows it to pause the time and more closely resembles the environment that's set by default for asynchronous tests. * Add a `zebra_test::init_async` helper function Calls `zebra_test::init` but also constructs a single-thread Tokio runtime and returns it. This makes it simpler to initialize asynchronous tests that can't use the `#[tokio::test]` attribute. * Replace usages of `Runtime::new` in tests Use the new `zebra_test::init_async()` helper function instead. * Replace `runtime::Builder::new_current_thread()` Use the new `zebra_test::init_async()` helper function instead. * Replace `runtime::Builder::new_multi_thread()` Use the new `zebra_test::init_async()` helper function instead. The test with the change doesn't necessarily have to use a multi-thread runtime. --- zebra-network/src/meta_addr/tests/prop.rs | 20 ++++---------- .../src/peer_set/candidate_set/tests/prop.rs | 13 +++------ .../peer_set/candidate_set/tests/vectors.rs | 17 +++--------- zebra-network/src/protocol/external/codec.rs | 27 +++++-------------- zebra-test/src/lib.rs | 20 ++++++++++++-- .../src/components/mempool/crawler/tests.rs | 20 +++----------- zebrad/src/components/mempool/tests/prop.rs | 18 +++---------- zebrad/src/components/sync/status/tests.rs | 5 +--- zebrad/src/components/sync/tests/timing.rs | 11 +++----- 9 files changed, 47 insertions(+), 104 deletions(-) diff --git a/zebra-network/src/meta_addr/tests/prop.rs b/zebra-network/src/meta_addr/tests/prop.rs index a7f70896..88179c73 100644 --- a/zebra-network/src/meta_addr/tests/prop.rs +++ b/zebra-network/src/meta_addr/tests/prop.rs @@ -12,7 +12,7 @@ use std::{ use chrono::Utc; use proptest::{collection::vec, prelude::*}; -use tokio::{runtime, time::Instant}; +use tokio::time::Instant; use tower::service_fn; use tracing::Span; @@ -199,7 +199,8 @@ proptest! { fn individual_peer_retry_limit_candidate_set( (addr, changes) in MetaAddrChange::addr_changes_strategy(MAX_ADDR_CHANGE) ) { - zebra_test::init(); + let runtime = zebra_test::init_async(); + let _guard = runtime.enter(); // Run the test for this many simulated live peer durations const LIVE_PEER_INTERVALS: u32 = 3; @@ -214,12 +215,6 @@ proptest! { "there are enough changes for good test coverage", ); - let runtime = runtime::Builder::new_current_thread() - .enable_all() - .build() - .expect("Failed to create Tokio runtime"); - let _guard = runtime.enter(); - // Only put valid addresses in the address book. // This means some tests will start with an empty address book. let addrs = if addr.last_known_info_is_valid_for_outbound() { @@ -302,7 +297,8 @@ proptest! { 2..MAX_ADDR_CHANGE ), ) { - zebra_test::init(); + let runtime = zebra_test::init_async(); + let _guard = runtime.enter(); let instant_now = std::time::Instant::now(); let chrono_now = Utc::now(); @@ -320,12 +316,6 @@ proptest! { "there are enough changes for good test coverage", ); - let runtime = runtime::Builder::new_current_thread() - .enable_all() - .build() - .expect("Failed to create Tokio runtime"); - let _guard = runtime.enter(); - let attempt_counts = runtime.block_on(async move { tokio::time::pause(); diff --git a/zebra-network/src/peer_set/candidate_set/tests/prop.rs b/zebra-network/src/peer_set/candidate_set/tests/prop.rs index e2b6f995..2edc92d2 100644 --- a/zebra-network/src/peer_set/candidate_set/tests/prop.rs +++ b/zebra-network/src/peer_set/candidate_set/tests/prop.rs @@ -8,10 +8,7 @@ use std::{ use futures::FutureExt; use proptest::{collection::vec, prelude::*}; -use tokio::{ - runtime::Runtime, - time::{sleep, timeout}, -}; +use tokio::time::{sleep, timeout}; use tracing::Span; use zebra_chain::serialization::DateTime32; @@ -60,9 +57,7 @@ proptest! { /// using a "not ready for attempt" peer generation strategy #[test] fn skipping_outbound_peer_connection_skips_rate_limit(next_peer_attempts in 0..TEST_ADDRESSES) { - zebra_test::init(); - - let runtime = Runtime::new().expect("Failed to create Tokio runtime"); + let runtime = zebra_test::init_async(); let _guard = runtime.enter(); let peer_service = tower::service_fn(|_| async { @@ -97,9 +92,7 @@ proptest! { initial_candidates in 0..MAX_TEST_CANDIDATES, extra_candidates in 0..MAX_TEST_CANDIDATES, ) { - zebra_test::init(); - - let runtime = Runtime::new().expect("Failed to create Tokio runtime"); + let runtime = zebra_test::init_async(); let _guard = runtime.enter(); let peer_service = tower::service_fn(|_| async { diff --git a/zebra-network/src/peer_set/candidate_set/tests/vectors.rs b/zebra-network/src/peer_set/candidate_set/tests/vectors.rs index 1daee008..a9a99960 100644 --- a/zebra-network/src/peer_set/candidate_set/tests/vectors.rs +++ b/zebra-network/src/peer_set/candidate_set/tests/vectors.rs @@ -7,10 +7,7 @@ use std::{ }; use chrono::{DateTime, Duration, Utc}; -use tokio::{ - runtime, - time::{self, Instant}, -}; +use tokio::time::{self, Instant}; use tracing::Span; use zebra_chain::serialization::DateTime32; @@ -136,12 +133,7 @@ fn candidate_set_updates_are_rate_limited() { // How many times should `update` be called in each rate limit interval const POLL_FREQUENCY_FACTOR: u32 = 3; - zebra_test::init(); - - let runtime = runtime::Builder::new_current_thread() - .enable_all() - .build() - .expect("Failed to create Tokio runtime"); + let runtime = zebra_test::init_async(); let _guard = runtime.enter(); let address_book = AddressBook::new(SocketAddr::from_str("0.0.0.0:0").unwrap(), Span::none()); @@ -182,10 +174,7 @@ fn candidate_set_updates_are_rate_limited() { /// rate limited. #[test] fn candidate_set_update_after_update_initial_is_rate_limited() { - let runtime = runtime::Builder::new_current_thread() - .enable_all() - .build() - .expect("Failed to create Tokio runtime"); + let runtime = zebra_test::init_async(); let _guard = runtime.enter(); let address_book = AddressBook::new(SocketAddr::from_str("0.0.0.0:0").unwrap(), Span::none()); diff --git a/zebra-network/src/protocol/external/codec.rs b/zebra-network/src/protocol/external/codec.rs index af6438ee..f1c53829 100644 --- a/zebra-network/src/protocol/external/codec.rs +++ b/zebra-network/src/protocol/external/codec.rs @@ -678,7 +678,6 @@ mod tests { use futures::prelude::*; use lazy_static::lazy_static; use std::net::{IpAddr, Ipv4Addr, SocketAddr}; - use tokio::runtime::Runtime; lazy_static! { static ref VERSION_TEST_VECTOR: Message = { @@ -707,8 +706,7 @@ mod tests { /// Check that the version test vector serializes and deserializes correctly #[test] fn version_message_round_trip() { - zebra_test::init(); - let rt = Runtime::new().unwrap(); + let rt = zebra_test::init_async(); let v = &*VERSION_TEST_VECTOR; @@ -761,8 +759,7 @@ mod tests { /// Deserialize a `Version` message containing `time`, and return the result. fn deserialize_version_with_time(time: i64) -> Result { - zebra_test::init(); - let rt = Runtime::new().unwrap(); + let rt = zebra_test::init_async(); let v = &*VERSION_TEST_VECTOR; @@ -803,9 +800,7 @@ mod tests { #[test] fn filterload_message_round_trip() { - zebra_test::init(); - - let rt = Runtime::new().unwrap(); + let rt = zebra_test::init_async(); let v = Message::FilterLoad { filter: Filter(vec![0; 35999]), @@ -839,9 +834,7 @@ mod tests { #[test] fn reject_message_no_extra_data_round_trip() { - zebra_test::init(); - - let rt = Runtime::new().unwrap(); + let rt = zebra_test::init_async(); let v = Message::Reject { message: "experimental".to_string(), @@ -875,9 +868,7 @@ mod tests { #[test] fn reject_message_extra_data_round_trip() { - zebra_test::init(); - - let rt = Runtime::new().unwrap(); + let rt = zebra_test::init_async(); let v = Message::Reject { message: "block".to_string(), @@ -911,9 +902,7 @@ mod tests { #[test] fn filterload_message_too_large_round_trip() { - zebra_test::init(); - - let rt = Runtime::new().unwrap(); + let rt = zebra_test::init_async(); let v = Message::FilterLoad { filter: Filter(vec![0; 40000]), @@ -947,9 +936,7 @@ mod tests { fn max_msg_size_round_trip() { use zebra_chain::serialization::ZcashDeserializeInto; - zebra_test::init(); - - let rt = Runtime::new().unwrap(); + let rt = zebra_test::init_async(); // make tests with a Tx message let tx: Transaction = zebra_test::vectors::DUMMY_TX1 diff --git a/zebra-test/src/lib.rs b/zebra-test/src/lib.rs index b9f841bb..8ccb24cf 100644 --- a/zebra-test/src/lib.rs +++ b/zebra-test/src/lib.rs @@ -30,7 +30,7 @@ pub mod zip0143; pub mod zip0243; pub mod zip0244; -/// A multi-threaded Tokio runtime that can be shared between tests. +/// A single-threaded Tokio runtime that can be shared between tests. /// /// This shared runtime should be used in tests that use shared background tasks. An example is /// with shared global `Lazy` types, because they spawn a background task when they @@ -46,7 +46,7 @@ pub mod zip0244; /// for example) and that means that the next test will already start with an incorrect timer /// state. pub static RUNTIME: Lazy = Lazy::new(|| { - tokio::runtime::Builder::new_multi_thread() + tokio::runtime::Builder::new_current_thread() .enable_all() .build() .expect("Failed to create Tokio runtime") @@ -121,6 +121,22 @@ pub fn init() { }) } +/// Initialize globals for tests that need a separate Tokio runtime instance. +/// +/// This is generally used in proptests, which don't support the `#[tokio::test]` attribute. +/// +/// If a runtime needs to be shared between tests, use the [`RUNTIME`] instance instead. +/// +/// See also the [`init`] function, which is called by this function. +pub fn init_async() -> tokio::runtime::Runtime { + init(); + + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("Failed to create Tokio runtime") +} + struct SkipTestReturnedErrPanicMessages; impl PanicMessage for SkipTestReturnedErrPanicMessages { diff --git a/zebrad/src/components/mempool/crawler/tests.rs b/zebrad/src/components/mempool/crawler/tests.rs index 819b22eb..652a3a25 100644 --- a/zebrad/src/components/mempool/crawler/tests.rs +++ b/zebrad/src/components/mempool/crawler/tests.rs @@ -50,11 +50,7 @@ proptest! { /// enabled, i.e., if the block synchronizer is likely close to the chain tip. #[test] fn crawler_requests_for_transaction_ids(mut sync_lengths in any::>()) { - let runtime = tokio::runtime::Builder::new_current_thread() - .enable_all() - .build() - .expect("Failed to create Tokio runtime"); - let _guard = runtime.enter(); + let runtime = zebra_test::init_async(); // Add a dummy last element, so that all of the original values are used. sync_lengths.push(0); @@ -101,11 +97,7 @@ proptest! { fn crawled_transactions_are_forwarded_to_downloader( transaction_ids in hash_set(any::(), 1..MAX_CRAWLED_TX), ) { - let runtime = tokio::runtime::Builder::new_current_thread() - .enable_all() - .build() - .expect("Failed to create Tokio runtime"); - let _guard = runtime.enter(); + let runtime = zebra_test::init_async(); let transaction_id_count = transaction_ids.len(); @@ -144,16 +136,12 @@ proptest! { vec(any::<(UnminedTxId, Result<(), MempoolError>)>(), 1..MAX_CRAWLED_TX), transaction_ids_for_return_to_normal in hash_set(any::(), 1..MAX_CRAWLED_TX), ) { + let runtime = zebra_test::init_async(); + // Make transaction_ids_and_responses unique let unique_transaction_ids_and_responses: HashSet = transaction_ids_and_responses.iter().map(|(id, _result)| id).copied().collect(); let transaction_ids_and_responses: Vec<(UnminedTxId, Result<(), MempoolError>)> = unique_transaction_ids_and_responses.iter().map(|unique_id| transaction_ids_and_responses.iter().find(|(id, _result)| id == unique_id).unwrap()).cloned().collect(); - let runtime = tokio::runtime::Builder::new_current_thread() - .enable_all() - .build() - .expect("Failed to create Tokio runtime"); - let _guard = runtime.enter(); - runtime.block_on(async move { let (mut peer_set, mut mempool, _sync_status, mut recent_sync_lengths, _chain_tip_sender) = setup_crawler(); diff --git a/zebrad/src/components/mempool/tests/prop.rs b/zebrad/src/components/mempool/tests/prop.rs index b0c69bce..998c7cfe 100644 --- a/zebrad/src/components/mempool/tests/prop.rs +++ b/zebrad/src/components/mempool/tests/prop.rs @@ -37,11 +37,7 @@ proptest! { transaction in any::(), chain_tip in any::(), ) { - let runtime = tokio::runtime::Builder::new_current_thread() - .enable_all() - .build() - .expect("Failed to create Tokio runtime"); - let _guard = runtime.enter(); + let runtime = zebra_test::init_async(); runtime.block_on(async move { let ( @@ -92,11 +88,7 @@ proptest! { mut transactions in vec(any::(), 0..CHAIN_LENGTH), fake_chain_tips in vec(any::(), 0..CHAIN_LENGTH), ) { - let runtime = tokio::runtime::Builder::new_current_thread() - .enable_all() - .build() - .expect("Failed to create Tokio runtime"); - let _guard = runtime.enter(); + let runtime = zebra_test::init_async(); runtime.block_on(async move { let ( @@ -178,11 +170,7 @@ proptest! { network in any::(), transaction in any::(), ) { - let runtime = tokio::runtime::Builder::new_current_thread() - .enable_all() - .build() - .expect("Failed to create Tokio runtime"); - let _guard = runtime.enter(); + let runtime = zebra_test::init_async(); runtime.block_on(async move { let ( diff --git a/zebrad/src/components/sync/status/tests.rs b/zebrad/src/components/sync/status/tests.rs index 1e962eaf..e501f5f9 100644 --- a/zebrad/src/components/sync/status/tests.rs +++ b/zebrad/src/components/sync/status/tests.rs @@ -33,10 +33,7 @@ proptest! { /// length updates and verifies if the other task was awakened by the update. #[test] fn waits_until_close_to_tip(sync_lengths in any::>()) { - let runtime = tokio::runtime::Builder::new_multi_thread() - .enable_all() - .build() - .expect("Failed to create Tokio runtime"); + let runtime = zebra_test::init_async(); let _guard = runtime.enter(); runtime.block_on(timeout(MAX_TEST_EXECUTION, root_task(sync_lengths)))??; diff --git a/zebrad/src/components/sync/tests/timing.rs b/zebrad/src/components/sync/tests/timing.rs index 276b0b8d..84ec1d38 100644 --- a/zebrad/src/components/sync/tests/timing.rs +++ b/zebrad/src/components/sync/tests/timing.rs @@ -3,10 +3,7 @@ use std::sync::{ atomic::{AtomicU8, Ordering}, Arc, }; -use tokio::{ - runtime::Runtime, - time::{timeout, Duration}, -}; +use tokio::time::{timeout, Duration}; use super::super::*; use crate::config::ZebradConfig; @@ -70,7 +67,8 @@ fn ensure_timeouts_consistent() { /// Test that calls to [`ChainSync::request_genesis`] are rate limited. #[test] fn request_genesis_is_rate_limited() { - zebra_test::init(); + let runtime = zebra_test::init_async(); + let _guard = runtime.enter(); // The number of calls to `request_genesis()` we are going to be testing for const RETRIES_TO_RUN: u8 = 3; @@ -81,9 +79,6 @@ fn request_genesis_is_rate_limited() { let state_requests_counter = Arc::new(AtomicU8::new(0)); let state_requests_counter_in_service = Arc::clone(&state_requests_counter); - let runtime = Runtime::new().expect("Failed to create Tokio runtime"); - let _guard = runtime.enter(); - // create a fake peer service that respond with `Error` to `BlocksByHash` or // panic in any other type of request. let peer_service = tower::service_fn(move |request| {