From 114b91afe6641966710cc22787630d43692eea21 Mon Sep 17 00:00:00 2001 From: Michael Vines Date: Wed, 11 Nov 2020 17:08:32 -0800 Subject: [PATCH] program-test now generates new blockhashes for test usage --- Cargo.lock | 2 ++ program-test/Cargo.toml | 2 ++ program-test/src/lib.rs | 74 ++++++++++++++++++++++++++++++++++++----- 3 files changed, 69 insertions(+), 9 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index bc4af20edf..6a04f7b280 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4580,6 +4580,7 @@ dependencies = [ name = "solana-program-test" version = "1.5.0" dependencies = [ + "async-trait", "base64 0.12.3", "chrono", "chrono-humanize", @@ -4591,6 +4592,7 @@ dependencies = [ "solana-program 1.5.0", "solana-runtime", "solana-sdk", + "tokio 0.3.2", ] [[package]] diff --git a/program-test/Cargo.toml b/program-test/Cargo.toml index 0daebbf32e..e128dd48d4 100644 --- a/program-test/Cargo.toml +++ b/program-test/Cargo.toml @@ -8,6 +8,7 @@ repository = "https://github.com/solana-labs/solana" version = "1.5.0" [dependencies] +async-trait = "0.1.36" base64 = "0.12.3" chrono = "0.4.19" chrono-humanize = "0.1.1" @@ -19,3 +20,4 @@ solana-logger = { path = "../logger", version = "1.5.0" } solana-program = { path = "../sdk/program", version = "1.5.0" } solana-runtime = { path = "../runtime", version = "1.5.0" } solana-sdk = { path = "../sdk", version = "1.5.0" } +tokio = { version = "0.3", features = ["full"] } diff --git a/program-test/src/lib.rs b/program-test/src/lib.rs index 4a5c45b434..fe64a14067 100644 --- a/program-test/src/lib.rs +++ b/program-test/src/lib.rs @@ -1,13 +1,15 @@ //! The solana-program-test provides a BanksClient-based test framework BPF programs +use async_trait::async_trait; use chrono_humanize::{Accuracy, HumanTime, Tense}; use log::*; use solana_banks_client::start_client; use solana_banks_server::banks_server::start_local_server; use solana_program::{ - account_info::AccountInfo, entrypoint::ProgramResult, hash::Hash, instruction::Instruction, - instruction::InstructionError, message::Message, native_token::sol_to_lamports, - program_error::ProgramError, program_stubs, pubkey::Pubkey, rent::Rent, + account_info::AccountInfo, entrypoint::ProgramResult, fee_calculator::FeeCalculator, + hash::Hash, instruction::Instruction, instruction::InstructionError, message::Message, + native_token::sol_to_lamports, program_error::ProgramError, program_stubs, pubkey::Pubkey, + rent::Rent, }; use solana_runtime::{ bank::{Bank, Builtin}, @@ -26,10 +28,11 @@ use std::{ collections::HashMap, convert::TryFrom, fs::File, - io::Read, + io::{self, Read}, path::{Path, PathBuf}, rc::Rc, sync::{Arc, RwLock}, + time::{Duration, Instant}, }; // Export types so test clients can limit their solana crate dependencies @@ -579,6 +582,7 @@ impl ProgramTest { let payer = gci.mint_keypair; debug!("Payer address: {}", payer.pubkey()); debug!("Genesis config: {}", genesis_config); + let target_tick_duration = genesis_config.poh_config.target_tick_duration; let mut bank = Bank::new(&genesis_config); @@ -628,10 +632,12 @@ impl ProgramTest { // non-zero fee calculator let last_blockhash = bank.last_blockhash(); while last_blockhash == bank.last_blockhash() { - bank.register_tick(&Hash::default()); + bank.register_tick(&Hash::new_unique()); } + let last_blockhash = bank.last_blockhash(); + // Make sure the new last_blockhash now requires a fee assert_ne!( - bank.get_fee_calculator(&bank.last_blockhash()) + bank.get_fee_calculator(&last_blockhash) .expect("fee_calculator") .lamports_per_signature, 0 @@ -639,11 +645,61 @@ impl ProgramTest { let bank_forks = Arc::new(RwLock::new(BankForks::new(bank))); let transport = start_local_server(&bank_forks).await; - let mut banks_client = start_client(transport) + let banks_client = start_client(transport) .await .unwrap_or_else(|err| panic!("Failed to start banks client: {}", err)); - let recent_blockhash = banks_client.get_recent_blockhash().await.unwrap(); - (banks_client, payer, recent_blockhash) + // Run a simulated PohService to provide the client with new blockhashes. New blockhashes + // are required when sending multiple otherwise identical transactions in series from a + // test + tokio::spawn(async move { + loop { + bank_forks + .read() + .unwrap() + .working_bank() + .register_tick(&Hash::new_unique()); + tokio::time::sleep(target_tick_duration).await; + } + }); + + (banks_client, payer, last_blockhash) + } +} + +#[async_trait] +pub trait ProgramTestBanksClientExt { + async fn get_new_blockhash(&mut self, blockhash: &Hash) -> io::Result<(Hash, FeeCalculator)>; +} + +#[async_trait] +impl ProgramTestBanksClientExt for BanksClient { + /// Get a new blockhash, similar in spirit to RpcClient::get_new_blockhash() + /// + /// This probably should eventually be moved into BanksClient proper in some form + async fn get_new_blockhash(&mut self, blockhash: &Hash) -> io::Result<(Hash, FeeCalculator)> { + let mut num_retries = 0; + let start = Instant::now(); + while start.elapsed().as_secs() < 5 { + if let Ok((fee_calculator, new_blockhash, _slot)) = self.get_fees().await { + if new_blockhash != *blockhash { + return Ok((new_blockhash, fee_calculator)); + } + } + debug!("Got same blockhash ({:?}), will retry...", blockhash); + + tokio::time::sleep(Duration::from_millis(200)).await; + num_retries += 1; + } + + Err(io::Error::new( + io::ErrorKind::Other, + format!( + "Unable to get new blockhash after {}ms (retried {} times), stuck at {}", + start.elapsed().as_millis(), + num_retries, + blockhash + ), + )) } }