From ec406afa7a5b5e5e326eec56e2006ef3026276bb Mon Sep 17 00:00:00 2001 From: grandizzy Date: Thu, 18 Jul 2024 11:31:22 +0300 Subject: [PATCH 1/2] fix(fork): preserve block number and timestamp when changing fork --- crates/evm/core/src/backend/mod.rs | 12 ++- crates/evm/core/src/fork/multi.rs | 124 ++++++++++++++---------- crates/forge/tests/it/repros.rs | 3 + testdata/default/repros/Issue8168.t.sol | 39 ++++++++ 4 files changed, 127 insertions(+), 51 deletions(-) create mode 100644 testdata/default/repros/Issue8168.t.sol diff --git a/crates/evm/core/src/backend/mod.rs b/crates/evm/core/src/backend/mod.rs index 7839bb262e93b..23c54e6e21545 100644 --- a/crates/evm/core/src/backend/mod.rs +++ b/crates/evm/core/src/backend/mod.rs @@ -1024,6 +1024,16 @@ impl DatabaseExt for Backend { return Ok(()); } + // Update block number and timestamp of active fork (if any) with current env values, + // in order to preserve values changed by using `roll` and `warp` cheatcodes. + if let Some(active_fork_id) = self.active_fork_id() { + self.forks.update_block( + self.ensure_fork_id(active_fork_id).cloned()?, + env.block.number, + env.block.timestamp, + )?; + } + let fork_id = self.ensure_fork_id(id).cloned()?; let idx = self.inner.ensure_fork_index(&fork_id)?; let fork_env = self @@ -1100,7 +1110,7 @@ impl DatabaseExt for Backend { } self.active_fork_ids = Some((id, idx)); - // update the environment accordingly + // Update current environment with environment of newly selected fork. update_current_env_with_fork_env(env, fork_env); Ok(()) diff --git a/crates/evm/core/src/fork/multi.rs b/crates/evm/core/src/fork/multi.rs index 372794ed1eb71..001c654699685 100644 --- a/crates/evm/core/src/fork/multi.rs +++ b/crates/evm/core/src/fork/multi.rs @@ -1,9 +1,10 @@ -//! Support for running multiple fork backends +//! Support for running multiple fork backends. //! //! The design is similar to the single `SharedBackend`, `BackendHandler` but supports multiple //! concurrently active pairs at once. use super::CreateFork; +use alloy_primitives::U256; use alloy_transport::layers::RetryBackoffService; use foundry_common::provider::{ runtime_transport::RuntimeTransport, ProviderBuilder, RetryProvider, @@ -65,13 +66,13 @@ impl> From for ForkId { } /// The Sender half of multi fork pair. -/// Can send requests to the `MultiForkHandler` to create forks +/// Can send requests to the `MultiForkHandler` to create forks. #[derive(Clone, Debug)] #[must_use] pub struct MultiFork { - /// Channel to send `Request`s to the handler + /// Channel to send `Request`s to the handler. handler: Sender, - /// Ensures that all rpc resources get flushed properly + /// Ensures that all rpc resources get flushed properly. _shutdown: Arc, } @@ -81,8 +82,8 @@ impl MultiFork { trace!(target: "fork::multi", "spawning multifork"); let (fork, mut handler) = Self::new(); - // spawn a light-weight thread with a thread-local async runtime just for - // sending and receiving data from the remote client(s) + // Spawn a light-weight thread with a thread-local async runtime just for + // sending and receiving data from the remote client(s). std::thread::Builder::new() .name("multi-fork-backend".into()) .spawn(move || { @@ -92,10 +93,10 @@ impl MultiFork { .expect("failed to build tokio runtime"); rt.block_on(async move { - // flush cache every 60s, this ensures that long-running fork tests get their - // cache flushed from time to time + // Flush cache every 60s, this ensures that long-running fork tests get their + // cache flushed from time to time. // NOTE: we install the interval here because the `tokio::timer::Interval` - // requires a rt + // requires a rt. handler.set_flush_cache_interval(Duration::from_secs(60)); handler.await }); @@ -115,9 +116,9 @@ impl MultiFork { (Self { handler, _shutdown }, MultiForkHandler::new(handler_rx)) } - /// Returns a fork backend + /// Returns a fork backend. /// - /// If no matching fork backend exists it will be created + /// If no matching fork backend exists it will be created. pub fn create_fork(&self, fork: CreateFork) -> eyre::Result<(ForkId, SharedBackend, Env)> { trace!("Creating new fork, url={}, block={:?}", fork.url, fork.evm_opts.fork_block_number); let (sender, rx) = oneshot_channel(); @@ -126,9 +127,9 @@ impl MultiFork { rx.recv()? } - /// Rolls the block of the fork + /// Rolls the block of the fork. /// - /// If no matching fork backend exists it will be created + /// If no matching fork backend exists it will be created. pub fn roll_fork( &self, fork: ForkId, @@ -141,7 +142,7 @@ impl MultiFork { rx.recv()? } - /// Returns the `Env` of the given fork, if any + /// Returns the `Env` of the given fork, if any. pub fn get_env(&self, fork: ForkId) -> eyre::Result> { trace!(?fork, "getting env config"); let (sender, rx) = oneshot_channel(); @@ -150,7 +151,16 @@ impl MultiFork { Ok(rx.recv()?) } - /// Returns the corresponding fork if it exists + /// Updates block number and timestamp of given fork with new values. + pub fn update_block(&self, fork: ForkId, number: U256, timestamp: U256) -> eyre::Result<()> { + trace!(?fork, ?number, ?timestamp, "update fork block"); + self.handler + .clone() + .try_send(Request::UpdateBlock(fork, number, timestamp)) + .map_err(|e| eyre::eyre!("{:?}", e)) + } + + /// Returns the corresponding fork if it exists. /// /// Returns `None` if no matching fork backend is available. pub fn get_fork(&self, id: impl Into) -> eyre::Result> { @@ -162,7 +172,7 @@ impl MultiFork { Ok(rx.recv()?) } - /// Returns the corresponding fork url if it exists + /// Returns the corresponding fork url if it exists. /// /// Returns `None` if no matching fork is available. pub fn get_fork_url(&self, id: impl Into) -> eyre::Result> { @@ -180,37 +190,39 @@ type CreateFuture = type CreateSender = OneshotSender>; type GetEnvSender = OneshotSender>; -/// Request that's send to the handler +/// Request that's send to the handler. #[derive(Debug)] enum Request { - /// Creates a new ForkBackend + /// Creates a new ForkBackend. CreateFork(Box, CreateSender), - /// Returns the Fork backend for the `ForkId` if it exists + /// Returns the Fork backend for the `ForkId` if it exists. GetFork(ForkId, OneshotSender>), - /// Adjusts the block that's being forked, by creating a new fork at the new block + /// Adjusts the block that's being forked, by creating a new fork at the new block. RollFork(ForkId, u64, CreateSender), - /// Returns the environment of the fork + /// Returns the environment of the fork. GetEnv(ForkId, GetEnvSender), + /// Updates the block number and timestamp of the fork. + UpdateBlock(ForkId, U256, U256), /// Shutdowns the entire `MultiForkHandler`, see `ShutDownMultiFork` ShutDown(OneshotSender<()>), - /// Returns the Fork Url for the `ForkId` if it exists + /// Returns the Fork Url for the `ForkId` if it exists. GetForkUrl(ForkId, OneshotSender>), } enum ForkTask { - /// Contains the future that will establish a new fork + /// Contains the future that will establish a new fork. Create(CreateFuture, ForkId, CreateSender, Vec), } -/// The type that manages connections in the background +/// The type that manages connections in the background. #[must_use = "futures do nothing unless polled"] pub struct MultiForkHandler { /// Incoming requests from the `MultiFork`. incoming: Fuse>, - /// All active handlers + /// All active handlers. /// - /// It's expected that this list will be rather small (<10) + /// It's expected that this list will be rather small (<10). handlers: Vec<(ForkId, Handler)>, // tasks currently in progress @@ -222,7 +234,7 @@ pub struct MultiForkHandler { /// block number. forks: HashMap, - /// Optional periodic interval to flush rpc cache + /// Optional periodic interval to flush rpc cache. flush_cache_interval: Option, } @@ -237,7 +249,7 @@ impl MultiForkHandler { } } - /// Sets the interval after which all rpc caches should be flushed periodically + /// Sets the interval after which all rpc caches should be flushed periodically. pub fn set_flush_cache_interval(&mut self, period: Duration) -> &mut Self { self.flush_cache_interval = Some(tokio::time::interval_at(tokio::time::Instant::now() + period, period)); @@ -261,13 +273,13 @@ impl MultiForkHandler { let fork_id = ForkId::new(&fork.url, fork.evm_opts.fork_block_number); trace!(?fork_id, "created new forkId"); - // there could already be a task for the requested fork in progress + // There could already be a task for the requested fork in progress. if let Some(in_progress) = self.find_in_progress_task(&fork_id) { in_progress.push(sender); return; } - // need to create a new fork + // Need to create a new fork. let task = Box::pin(create_fork(fork)); self.pending_tasks.push(ForkTask::Create(task, fork_id, sender, Vec::new())); } @@ -282,7 +294,7 @@ impl MultiForkHandler { self.forks.insert(fork_id.clone(), fork.clone()); let _ = sender.send(Ok((fork_id.clone(), fork.backend.clone(), fork.opts.env.clone()))); - // notify all additional senders and track unique forkIds + // Notify all additional senders and track unique forkIds. for sender in additional_senders { let next_fork_id = fork.inc_senders(fork_id.clone()); self.forks.insert(next_fork_id.clone(), fork.clone()); @@ -290,6 +302,15 @@ impl MultiForkHandler { } } + /// Update fork block number and timestamp. Used to preserve values set by `roll` and `warp` + /// cheatcodes when new fork selected. + fn update_block(&mut self, fork_id: ForkId, block_number: U256, block_timestamp: U256) { + if let Some(fork) = self.forks.get_mut(&fork_id) { + fork.opts.env.block.number = block_number; + fork.opts.env.block.timestamp = block_timestamp; + } + } + fn on_request(&mut self, req: Request) { match req { Request::CreateFork(fork, sender) => self.create_fork(*fork, sender), @@ -310,9 +331,12 @@ impl MultiForkHandler { Request::GetEnv(fork_id, sender) => { let _ = sender.send(self.forks.get(&fork_id).map(|fork| fork.opts.env.clone())); } + Request::UpdateBlock(fork_id, block_number, block_timestamp) => { + self.update_block(fork_id, block_number, block_timestamp); + } Request::ShutDown(sender) => { trace!(target: "fork::multi", "received shutdown signal"); - // we're emptying all fork backends, this way we ensure all caches get flushed + // We're emptying all fork backends, this way we ensure all caches get flushed. self.forks.clear(); self.handlers.clear(); let _ = sender.send(()); @@ -325,22 +349,22 @@ impl MultiForkHandler { } } -// Drives all handler to completion -// This future will finish once all underlying BackendHandler are completed +// Drives all handler to completion. +// This future will finish once all underlying BackendHandler are completed. impl Future for MultiForkHandler { type Output = (); fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let pin = self.get_mut(); - // receive new requests + // Receive new requests. loop { match Pin::new(&mut pin.incoming).poll_next(cx) { Poll::Ready(Some(req)) => { pin.on_request(req); } Poll::Ready(None) => { - // channel closed, but we still need to drive the fork handlers to completion + // Channel closed, but we still need to drive the fork handlers to completion. trace!(target: "fork::multi", "request channel closed"); break; } @@ -348,7 +372,7 @@ impl Future for MultiForkHandler { } } - // advance all tasks + // Advance all tasks. for n in (0..pin.pending_tasks.len()).rev() { let task = pin.pending_tasks.swap_remove(n); match task { @@ -387,7 +411,7 @@ impl Future for MultiForkHandler { } } - // advance all handlers + // Advance all handlers. for n in (0..pin.handlers.len()).rev() { let (id, mut handler) = pin.handlers.swap_remove(n); match handler.poll_unpin(cx) { @@ -405,7 +429,7 @@ impl Future for MultiForkHandler { return Poll::Ready(()); } - // periodically flush cached RPC state + // Periodically flush cached RPC state. if pin .flush_cache_interval .as_mut() @@ -415,7 +439,7 @@ impl Future for MultiForkHandler { { trace!(target: "fork::multi", "tick flushing caches"); let forks = pin.forks.values().map(|f| f.backend.clone()).collect::>(); - // flush this on new thread to not block here + // Flush this on new thread to not block here. std::thread::Builder::new() .name("flusher".into()) .spawn(move || { @@ -431,12 +455,12 @@ impl Future for MultiForkHandler { /// Tracks the created Fork #[derive(Debug, Clone)] struct CreatedFork { - /// How the fork was initially created + /// How the fork was initially created. opts: CreateFork, - /// Copy of the sender + /// Copy of the sender. backend: SharedBackend, /// How many consumers there are, since a `SharedBacked` can be used by multiple - /// consumers + /// consumers. num_senders: Arc, } @@ -445,7 +469,7 @@ impl CreatedFork { Self { opts, backend, num_senders: Arc::new(AtomicUsize::new(1)) } } - /// Increment senders and return unique identifier of the fork + /// Increment senders and return unique identifier of the fork. fn inc_senders(&self, fork_id: ForkId) -> ForkId { format!( "{}-{}", @@ -458,7 +482,7 @@ impl CreatedFork { /// A type that's used to signaling the `MultiForkHandler` when it's time to shut down. /// -/// This is essentially a sync on drop, so that the `MultiForkHandler` can flush all rpc cashes +/// This is essentially a sync on drop, so that the `MultiForkHandler` can flush all rpc cashes. /// /// This type intentionally does not implement `Clone` since it's intended that there's only once /// instance. @@ -481,9 +505,9 @@ impl Drop for ShutDownMultiFork { } } -/// Creates a new fork +/// Creates a new fork. /// -/// This will establish a new `Provider` to the endpoint and return the Fork Backend +/// This will establish a new `Provider` to the endpoint and return the Fork Backend. async fn create_fork(mut fork: CreateFork) -> eyre::Result<(ForkId, CreatedFork, Handler)> { let provider = Arc::new( ProviderBuilder::new(fork.url.as_str()) @@ -493,16 +517,16 @@ async fn create_fork(mut fork: CreateFork) -> eyre::Result<(ForkId, CreatedFork, .build()?, ); - // initialise the fork environment + // Initialise the fork environment. let (env, block) = fork.evm_opts.fork_evm_env(&fork.url).await?; fork.env = env; let meta = BlockchainDbMeta::new(fork.env.clone(), fork.url.clone()); - // we need to use the block number from the block because the env's number can be different on + // We need to use the block number from the block because the env's number can be different on // some L2s (e.g. Arbitrum). let number = block.header.number.unwrap_or(meta.block_env.number.to()); - // determine the cache path if caching is enabled + // Determine the cache path if caching is enabled. let cache_path = if fork.enable_caching { Config::foundry_block_cache_dir(meta.cfg_env.chain_id, number) } else { diff --git a/crates/forge/tests/it/repros.rs b/crates/forge/tests/it/repros.rs index f1b1c8a96ef4b..9216d275488d3 100644 --- a/crates/forge/tests/it/repros.rs +++ b/crates/forge/tests/it/repros.rs @@ -358,3 +358,6 @@ test_repro!(8277); // https://github.com/foundry-rs/foundry/issues/8287 test_repro!(8287); + +// https://github.com/foundry-rs/foundry/issues/8168 +test_repro!(8168); diff --git a/testdata/default/repros/Issue8168.t.sol b/testdata/default/repros/Issue8168.t.sol new file mode 100644 index 0000000000000..dab69c4568190 --- /dev/null +++ b/testdata/default/repros/Issue8168.t.sol @@ -0,0 +1,39 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 +pragma solidity 0.8.18; + +import "ds-test/test.sol"; +import "cheats/Vm.sol"; + +// https://github.com/foundry-rs/foundry/issues/8168 +contract Issue8168Test is DSTest { + Vm constant vm = Vm(HEVM_ADDRESS); + + function testForkWarpRollPreserved() public { + uint256 fork1 = vm.createFork("rpcAlias"); + uint256 fork2 = vm.createFork("rpcAlias"); + + vm.selectFork(fork1); + uint256 initial_fork1_number = block.number; + uint256 initial_fork1_ts = block.timestamp; + vm.warp(block.timestamp + 1000); + vm.roll(block.number + 100); + assertEq(block.timestamp, initial_fork1_ts + 1000); + assertEq(block.number, initial_fork1_number + 100); + + vm.selectFork(fork2); + uint256 initial_fork2_number = block.number; + uint256 initial_fork2_ts = block.timestamp; + vm.warp(block.timestamp + 2000); + vm.roll(block.number + 200); + assertEq(block.timestamp, initial_fork1_ts + 2000); + assertEq(block.number, initial_fork1_number + 200); + + vm.selectFork(fork1); + assertEq(block.timestamp, initial_fork1_ts + 1000); + assertEq(block.number, initial_fork1_number + 100); + + vm.selectFork(fork2); + assertEq(block.timestamp, initial_fork2_ts + 2000); + assertEq(block.number, initial_fork2_number + 200); + } +} From c3f675d53264d941ca04a08a94a4b70abd4fc134 Mon Sep 17 00:00:00 2001 From: grandizzy Date: Thu, 18 Jul 2024 18:00:00 +0300 Subject: [PATCH 2/2] Minor test update, could have been failing if forks created at different blocks --- testdata/default/repros/Issue8168.t.sol | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/testdata/default/repros/Issue8168.t.sol b/testdata/default/repros/Issue8168.t.sol index dab69c4568190..b9bd5757abc70 100644 --- a/testdata/default/repros/Issue8168.t.sol +++ b/testdata/default/repros/Issue8168.t.sol @@ -25,8 +25,8 @@ contract Issue8168Test is DSTest { uint256 initial_fork2_ts = block.timestamp; vm.warp(block.timestamp + 2000); vm.roll(block.number + 200); - assertEq(block.timestamp, initial_fork1_ts + 2000); - assertEq(block.number, initial_fork1_number + 200); + assertEq(block.timestamp, initial_fork2_ts + 2000); + assertEq(block.number, initial_fork2_number + 200); vm.selectFork(fork1); assertEq(block.timestamp, initial_fork1_ts + 1000);