diff --git a/crates/katana/storage/provider/src/providers/fork/backend.rs b/crates/katana/storage/provider/src/providers/fork/backend.rs index 265aa502a3..64e921ecdc 100644 --- a/crates/katana/storage/provider/src/providers/fork/backend.rs +++ b/crates/katana/storage/provider/src/providers/fork/backend.rs @@ -1,4 +1,5 @@ -use std::collections::VecDeque; +use std::collections::hash_map::Entry; +use std::collections::{HashMap, VecDeque}; use std::pin::Pin; use std::sync::mpsc::{ channel as oneshot, Receiver as OneshotReceiver, RecvError, Sender as OneshotSender, @@ -7,6 +8,7 @@ use std::sync::Arc; use std::task::{Context, Poll}; use std::{io, thread}; +use anyhow::anyhow; use futures::channel::mpsc::{channel as async_channel, Receiver, SendError, Sender}; use futures::future::BoxFuture; use futures::stream::Stream; @@ -34,28 +36,34 @@ const LOG_TARGET: &str = "forking::backend"; type BackendResult = Result; -type GetNonceResult = BackendResult; -type GetStorageResult = BackendResult; -type GetClassHashAtResult = BackendResult; -type GetClassAtResult = BackendResult; +/// The types of response from [`Backend`]. +#[derive(Debug, Clone)] +enum BackendResponse { + Nonce(BackendResult), + Storage(BackendResult), + ClassHashAt(BackendResult), + ClassAt(BackendResult), +} -#[derive(Debug, thiserror::Error)] +#[derive(Debug, thiserror::Error, Clone)] pub enum BackendError { #[error("failed to send request to backend: {0}")] FailedSendRequest(#[from] SendError), #[error("failed to receive result from backend: {0}")] FailedReceiveResult(#[from] RecvError), #[error("compute class hash error: {0}")] - ComputeClassHashError(anyhow::Error), + ComputeClassHashError(Arc), #[error("failed to spawn backend thread: {0}")] - BackendThreadInit(#[from] io::Error), + BackendThreadInit(#[from] Arc), #[error("rpc provider error: {0}")] - StarknetProvider(#[from] starknet::providers::ProviderError), + StarknetProvider(#[from] Arc), + #[error("unexpected received result: {0}")] + UnexpectedReceiveResult(Arc), } -struct Request { +struct Request

{ payload: P, - sender: OneshotSender>, + sender: OneshotSender, } /// The types of request that can be sent to [`Backend`]. @@ -63,10 +71,10 @@ struct Request { /// Each request consists of a payload and the sender half of a oneshot channel that will be used /// to send the result back to the backend handle. enum BackendRequest { - Nonce(Request), - Class(Request), - ClassHash(Request), - Storage(Request<(ContractAddress, StorageKey), StorageValue>), + Nonce(Request), + Class(Request), + ClassHash(Request), + Storage(Request<(ContractAddress, StorageKey)>), // Test-only request kind for requesting the backend stats #[cfg(test)] Stats(OneshotSender), @@ -74,21 +82,19 @@ enum BackendRequest { impl BackendRequest { /// Create a new request for fetching the nonce of a contract. - fn nonce(address: ContractAddress) -> (BackendRequest, OneshotReceiver) { + fn nonce(address: ContractAddress) -> (BackendRequest, OneshotReceiver) { let (sender, receiver) = oneshot(); (BackendRequest::Nonce(Request { payload: address, sender }), receiver) } /// Create a new request for fetching the class definitions of a contract. - fn class(hash: ClassHash) -> (BackendRequest, OneshotReceiver) { + fn class(hash: ClassHash) -> (BackendRequest, OneshotReceiver) { let (sender, receiver) = oneshot(); (BackendRequest::Class(Request { payload: hash, sender }), receiver) } /// Create a new request for fetching the class hash of a contract. - fn class_hash( - address: ContractAddress, - ) -> (BackendRequest, OneshotReceiver) { + fn class_hash(address: ContractAddress) -> (BackendRequest, OneshotReceiver) { let (sender, receiver) = oneshot(); (BackendRequest::ClassHash(Request { payload: address, sender }), receiver) } @@ -97,7 +103,7 @@ impl BackendRequest { fn storage( address: ContractAddress, key: StorageKey, - ) -> (BackendRequest, OneshotReceiver) { + ) -> (BackendRequest, OneshotReceiver) { let (sender, receiver) = oneshot(); (BackendRequest::Storage(Request { payload: (address, key), sender }), receiver) } @@ -109,7 +115,17 @@ impl BackendRequest { } } -type BackendRequestFuture = BoxFuture<'static, ()>; +type BackendRequestFuture = BoxFuture<'static, BackendResponse>; + +// Identifier for pending requests. +// This is used for request deduplication. +#[derive(Eq, Hash, PartialEq, Clone, Copy, Debug)] +enum BackendRequestIdentifier { + Nonce(ContractAddress), + Class(ClassHash), + ClassHash(ContractAddress), + Storage((ContractAddress, StorageKey)), +} /// The backend for the forked provider. /// @@ -119,8 +135,10 @@ type BackendRequestFuture = BoxFuture<'static, ()>; pub struct Backend

{ /// The Starknet RPC provider that will be used to fetch data from. provider: Arc

, + // HashMap that keep track of current requests, for dedup purposes. + request_dedup_map: HashMap>>, /// Requests that are currently being poll. - pending_requests: Vec, + pending_requests: Vec<(BackendRequestIdentifier, BackendRequestFuture)>, /// Requests that are queued to be polled. queued_requests: VecDeque, /// A channel for receiving requests from the [BackendHandle]s. @@ -150,7 +168,7 @@ where .expect("failed to create tokio runtime") .block_on(backend); }) - .map_err(BackendError::BackendThreadInit)?; + .map_err(|e| BackendError::BackendThreadInit(Arc::new(e)))?; trace!(target: LOG_TARGET, "Forking backend started."); @@ -169,6 +187,7 @@ where block, incoming: rx, provider: Arc::new(provider), + request_dedup_map: HashMap::new(), pending_requests: Vec::new(), queued_requests: VecDeque::new(), }; @@ -182,57 +201,73 @@ where let block = self.block; let provider = self.provider.clone(); + // Check if there are similar requests in the queue before sending the request match request { BackendRequest::Nonce(Request { payload, sender }) => { - let fut = Box::pin(async move { - let res = provider - .get_nonce(block, Felt::from(payload)) - .await - .map_err(BackendError::StarknetProvider); - - sender.send(res).expect("failed to send nonce result") - }); - - self.pending_requests.push(fut); + let req_key = BackendRequestIdentifier::Nonce(payload); + + self.dedup_request( + req_key, + sender, + Box::pin(async move { + let res = provider + .get_nonce(block, Felt::from(payload)) + .await + .map_err(|e| BackendError::StarknetProvider(Arc::new(e))); + BackendResponse::Nonce(res) + }), + ); } BackendRequest::Storage(Request { payload: (addr, key), sender }) => { - let fut = Box::pin(async move { - let res = provider - .get_storage_at(Felt::from(addr), key, block) - .await - .map_err(BackendError::StarknetProvider); - - sender.send(res).expect("failed to send storage result") - }); - - self.pending_requests.push(fut); + let req_key = BackendRequestIdentifier::Storage((addr, key)); + + self.dedup_request( + req_key, + sender, + Box::pin(async move { + let res = provider + .get_storage_at(Felt::from(addr), key, block) + .await + .map_err(|e| BackendError::StarknetProvider(Arc::new(e))); + + BackendResponse::Storage(res) + }), + ); } BackendRequest::ClassHash(Request { payload, sender }) => { - let fut = Box::pin(async move { - let res = provider - .get_class_hash_at(block, Felt::from(payload)) - .await - .map_err(BackendError::StarknetProvider); - - sender.send(res).expect("failed to send class hash result") - }); - - self.pending_requests.push(fut); + let req_key = BackendRequestIdentifier::ClassHash(payload); + + self.dedup_request( + req_key, + sender, + Box::pin(async move { + let res = provider + .get_class_hash_at(block, Felt::from(payload)) + .await + .map_err(|e| BackendError::StarknetProvider(Arc::new(e))); + + BackendResponse::ClassHashAt(res) + }), + ); } BackendRequest::Class(Request { payload, sender }) => { - let fut = Box::pin(async move { - let res = provider - .get_class(block, payload) - .await - .map_err(BackendError::StarknetProvider); - - sender.send(res).expect("failed to send class result") - }); - - self.pending_requests.push(fut); + let req_key = BackendRequestIdentifier::Class(payload); + + self.dedup_request( + req_key, + sender, + Box::pin(async move { + let res = provider + .get_class(block, payload) + .await + .map_err(|e| BackendError::StarknetProvider(Arc::new(e))); + + BackendResponse::ClassAt(res) + }), + ); } #[cfg(test)] @@ -242,6 +277,29 @@ where } } } + + fn dedup_request( + &mut self, + req_key: BackendRequestIdentifier, + sender: OneshotSender, + rpc_call_future: BoxFuture<'static, BackendResponse>, + ) { + if let Entry::Vacant(e) = self.request_dedup_map.entry(req_key) { + self.pending_requests.push((req_key, rpc_call_future)); + e.insert(vec![sender]); + } else { + match self.request_dedup_map.get_mut(&req_key) { + Some(sender_vec) => { + sender_vec.push(sender); + } + None => { + // Log this and do nothing here, as this should never happen. + // If this does happen it is an unexpected bug. + error!(target: LOG_TARGET, "failed to get current request dedup vector"); + } + } + } + } } impl

Future for Backend

@@ -275,11 +333,28 @@ where // poll all pending requests for n in (0..pin.pending_requests.len()).rev() { - let mut fut = pin.pending_requests.swap_remove(n); + let (fut_key, mut fut) = pin.pending_requests.swap_remove(n); // poll the future and if the future is still pending, push it back to the // pending requests so that it will be polled again - if fut.poll_unpin(cx).is_pending() { - pin.pending_requests.push(fut); + match fut.poll_unpin(cx) { + Poll::Pending => { + pin.pending_requests.push((fut_key, fut)); + } + Poll::Ready(res) => { + let sender_vec = pin + .request_dedup_map + .get(&fut_key) + .expect("failed to get sender vector"); + + // Send the response to all the senders waiting on the same request + sender_vec.iter().for_each(|sender| { + sender.send(res.clone()).unwrap_or_else(|error| { + error!(target: LOG_TARGET, key = ?fut_key, %error, "Failed to send result.") + }); + }); + + pin.request_dedup_map.remove(&fut_key); + } } } @@ -309,7 +384,12 @@ impl BackendHandle { trace!(target: LOG_TARGET, %address, "Requesting contract nonce."); let (req, rx) = BackendRequest::nonce(address); self.request(req)?; - rx.recv()? + match rx.recv()? { + BackendResponse::Nonce(res) => res, + response => { + Err(BackendError::UnexpectedReceiveResult(Arc::new(anyhow!("{:?}", response)))) + } + } } pub fn get_storage( @@ -320,21 +400,36 @@ impl BackendHandle { trace!(target: LOG_TARGET, %address, key = %format!("{key:#x}"), "Requesting contract storage."); let (req, rx) = BackendRequest::storage(address, key); self.request(req)?; - rx.recv()? + match rx.recv()? { + BackendResponse::Storage(res) => res, + response => { + Err(BackendError::UnexpectedReceiveResult(Arc::new(anyhow!("{:?}", response)))) + } + } } pub fn get_class_hash_at(&self, address: ContractAddress) -> Result { trace!(target: LOG_TARGET, %address, "Requesting contract class hash."); let (req, rx) = BackendRequest::class_hash(address); self.request(req)?; - rx.recv()? + match rx.recv()? { + BackendResponse::ClassHashAt(res) => res, + response => { + Err(BackendError::UnexpectedReceiveResult(Arc::new(anyhow!("{:?}", response)))) + } + } } pub fn get_class_at(&self, class_hash: ClassHash) -> Result { trace!(target: LOG_TARGET, class_hash = %format!("{class_hash:#x}"), "Requesting class."); let (req, rx) = BackendRequest::class(class_hash); self.request(req)?; - rx.recv()? + match rx.recv()? { + BackendResponse::ClassAt(res) => res, + response => { + Err(BackendError::UnexpectedReceiveResult(Arc::new(anyhow!("{:?}", response)))) + } + } } pub fn get_compiled_class_hash( @@ -349,7 +444,7 @@ impl BackendHandle { RpcContractClass::Legacy(_) => Ok(class_hash), RpcContractClass::Sierra(sierra_class) => { compiled_class_hash_from_flattened_sierra_class(&sierra_class) - .map_err(BackendError::ComputeClassHashError) + .map_err(|e| BackendError::ComputeClassHashError(Arc::new(e))) } } } @@ -592,9 +687,12 @@ fn handle_not_found_err(result: Result) -> Result, match result { Ok(value) => Ok(Some(value)), - Err(BackendError::StarknetProvider(StarknetProviderError::StarknetError( - StarknetError::ContractNotFound | StarknetError::ClassHashNotFound, - ))) => Ok(None), + Err(BackendError::StarknetProvider(err_in_arc)) => match err_in_arc.as_ref() { + StarknetProviderError::StarknetError( + StarknetError::ContractNotFound | StarknetError::ClassHashNotFound, + ) => Ok(None), + _ => Err(BackendError::StarknetProvider(err_in_arc)), + }, Err(e) => Err(e), } @@ -603,11 +701,12 @@ fn handle_not_found_err(result: Result) -> Result, #[cfg(test)] pub(crate) mod test_utils { - use std::sync::mpsc::sync_channel; + use std::sync::mpsc::{sync_channel, SyncSender}; use katana_primitives::block::BlockNumber; use starknet::providers::jsonrpc::HttpTransport; use starknet::providers::JsonRpcClient; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::TcpListener; use url::Url; @@ -620,13 +719,13 @@ pub(crate) mod test_utils { } // Starts a TCP server that never close the connection. - pub fn start_tcp_server() { + pub fn start_tcp_server(addr: String) { use tokio::runtime::Builder; let (tx, rx) = sync_channel::<()>(1); thread::spawn(move || { Builder::new_current_thread().enable_all().build().unwrap().block_on(async move { - let listener = TcpListener::bind("127.0.0.1:8080").await.unwrap(); + let listener = TcpListener::bind(addr).await.unwrap(); let mut connections = Vec::new(); tx.send(()).unwrap(); @@ -640,11 +739,48 @@ pub(crate) mod test_utils { rx.recv().unwrap(); } + + // Helper function to start a TCP server that returns predefined JSON-RPC responses + pub fn start_mock_rpc_server(addr: String, response: String) -> SyncSender<()> { + use tokio::runtime::Builder; + let (tx, rx) = sync_channel::<()>(1); + + thread::spawn(move || { + Builder::new_current_thread().enable_all().build().unwrap().block_on(async move { + let listener = TcpListener::bind(addr).await.unwrap(); + + loop { + let (mut socket, _) = listener.accept().await.unwrap(); + + // Read the request, so hyper would not close the connection + let mut buffer = [0; 1024]; + let _ = socket.read(&mut buffer).await.unwrap(); + + // Wait for a signal to return the response. + rx.recv().unwrap(); + + // After reading, we send the pre-determined response + let http_response = format!( + "HTTP/1.1 200 OK\r\ncontent-length: {}\r\ncontent-type: \ + application/json\r\n\r\n{}", + response.len(), + response + ); + + socket.write_all(http_response.as_bytes()).await.unwrap(); + socket.flush().await.unwrap(); + } + }); + }); + + // Returning the sender to allow controlling the response timing. + tx + } } #[cfg(test)] mod tests { - + use std::sync::Mutex; use std::time::Duration; use katana_primitives::contract::GenericContractInfo; @@ -667,7 +803,7 @@ mod tests { #[test] fn handle_incoming_requests() { // start a mock remote network - start_tcp_server(); + start_tcp_server("127.0.0.1:8080".to_string()); let handle = create_forked_backend("http://127.0.0.1:8080", 1); @@ -686,7 +822,7 @@ mod tests { }); let h3 = handle.clone(); thread::spawn(move || { - h3.get_compiled_class_hash(felt!("0x1")).expect(ERROR_SEND_REQUEST); + h3.get_compiled_class_hash(felt!("0x2")).expect(ERROR_SEND_REQUEST); }); let h4 = handle.clone(); thread::spawn(move || { @@ -705,6 +841,319 @@ mod tests { assert_eq!(stats, 5, "Backend should have 5 ongoing requests.") } + #[test] + fn get_nonce_request_should_be_deduplicated() { + // start a mock remote network + start_tcp_server("127.0.0.1:8081".to_string()); + + let handle = create_forked_backend("http://127.0.0.1:8081", 1); + + // check no pending requests + let stats = handle.stats().expect(ERROR_STATS); + assert_eq!(stats, 0, "Backend should not have any ongoing requests."); + + // send requests to the backend + let h1 = handle.clone(); + thread::spawn(move || { + h1.get_nonce(felt!("0x1").into()).expect(ERROR_SEND_REQUEST); + }); + let h2 = handle.clone(); + thread::spawn(move || { + h2.get_nonce(felt!("0x1").into()).expect(ERROR_SEND_REQUEST); + }); + + // wait for the requests to be handled + thread::sleep(Duration::from_secs(1)); + + // check current request count + let stats = handle.stats().expect(ERROR_STATS); + assert_eq!(stats, 1, "Backend should have 1 ongoing requests."); + + // Different request, should be counted + let h3 = handle.clone(); + thread::spawn(move || { + h3.get_nonce(felt!("0x2").into()).expect(ERROR_SEND_REQUEST); + }); + + // wait for the requests to be handled + thread::sleep(Duration::from_secs(1)); + + // check request are handled + let stats = handle.stats().expect(ERROR_STATS); + assert_eq!(stats, 2, "Backend should only have 2 ongoing requests.") + } + + #[test] + fn get_class_at_request_should_be_deduplicated() { + // start a mock remote network + start_tcp_server("127.0.0.1:8082".to_string()); + + let handle = create_forked_backend("http://127.0.0.1:8082", 1); + + // check no pending requests + let stats = handle.stats().expect(ERROR_STATS); + assert_eq!(stats, 0, "Backend should not have any ongoing requests."); + + // send requests to the backend + let h1 = handle.clone(); + thread::spawn(move || { + h1.get_class_at(felt!("0x1")).expect(ERROR_SEND_REQUEST); + }); + let h2 = handle.clone(); + thread::spawn(move || { + h2.get_class_at(felt!("0x1")).expect(ERROR_SEND_REQUEST); + }); + + // wait for the requests to be handled + thread::sleep(Duration::from_secs(1)); + + // check current request count + let stats = handle.stats().expect(ERROR_STATS); + assert_eq!(stats, 1, "Backend should have 1 ongoing requests."); + + // Different request, should be counted + let h3 = handle.clone(); + thread::spawn(move || { + h3.get_class_at(felt!("0x2")).expect(ERROR_SEND_REQUEST); + }); + + // wait for the requests to be handled + thread::sleep(Duration::from_secs(1)); + + // check request are handled + let stats = handle.stats().expect(ERROR_STATS); + assert_eq!(stats, 2, "Backend should only have 2 ongoing requests.") + } + + #[test] + fn get_compiled_class_hash_request_should_be_deduplicated() { + // start a mock remote network + start_tcp_server("127.0.0.1:8083".to_string()); + + let handle = create_forked_backend("http://127.0.0.1:8083", 1); + + // check no pending requests + let stats = handle.stats().expect(ERROR_STATS); + assert_eq!(stats, 0, "Backend should not have any ongoing requests."); + + // send requests to the backend + let h1 = handle.clone(); + thread::spawn(move || { + h1.get_compiled_class_hash(felt!("0x1")).expect(ERROR_SEND_REQUEST); + }); + let h2 = handle.clone(); + thread::spawn(move || { + h2.get_compiled_class_hash(felt!("0x1")).expect(ERROR_SEND_REQUEST); + }); + + // wait for the requests to be handled + thread::sleep(Duration::from_secs(1)); + + // check current request count + let stats = handle.stats().expect(ERROR_STATS); + assert_eq!(stats, 1, "Backend should have 1 ongoing requests."); + + // Different request, should be counted + let h3 = handle.clone(); + thread::spawn(move || { + h3.get_compiled_class_hash(felt!("0x2")).expect(ERROR_SEND_REQUEST); + }); + + // wait for the requests to be handled + thread::sleep(Duration::from_secs(1)); + + // check request are handled + let stats = handle.stats().expect(ERROR_STATS); + assert_eq!(stats, 2, "Backend should only have 2 ongoing requests.") + } + + #[test] + fn get_class_at_and_get_compiled_class_hash_request_should_be_deduplicated() { + // start a mock remote network + start_tcp_server("127.0.0.1:8084".to_string()); + + let handle = create_forked_backend("http://127.0.0.1:8084", 1); + + // check no pending requests + let stats = handle.stats().expect(ERROR_STATS); + assert_eq!(stats, 0, "Backend should not have any ongoing requests."); + + // send requests to the backend + let h1 = handle.clone(); + thread::spawn(move || { + h1.get_class_at(felt!("0x1")).expect(ERROR_SEND_REQUEST); + }); + // Since this also calls to the same request as the previous one, it should be deduped + let h2 = handle.clone(); + thread::spawn(move || { + h2.get_compiled_class_hash(felt!("0x1")).expect(ERROR_SEND_REQUEST); + }); + + // wait for the requests to be handled + thread::sleep(Duration::from_secs(1)); + + // check current request count + let stats = handle.stats().expect(ERROR_STATS); + assert_eq!(stats, 1, "Backend should have 1 ongoing requests."); + + // Different request, should be counted + let h3 = handle.clone(); + thread::spawn(move || { + h3.get_class_at(felt!("0x2")).expect(ERROR_SEND_REQUEST); + }); + + // wait for the requests to be handled + thread::sleep(Duration::from_secs(1)); + + // check request are handled + let stats = handle.stats().expect(ERROR_STATS); + assert_eq!(stats, 2, "Backend should only have 2 ongoing requests.") + } + + #[test] + fn get_class_hash_at_request_should_be_deduplicated() { + // start a mock remote network + start_tcp_server("127.0.0.1:8085".to_string()); + + let handle = create_forked_backend("http://127.0.0.1:8085", 1); + + // check no pending requests + let stats = handle.stats().expect(ERROR_STATS); + assert_eq!(stats, 0, "Backend should not have any ongoing requests."); + + // send requests to the backend + let h1 = handle.clone(); + thread::spawn(move || { + h1.get_class_hash_at(felt!("0x1").into()).expect(ERROR_SEND_REQUEST); + }); + let h2 = handle.clone(); + thread::spawn(move || { + h2.get_class_hash_at(felt!("0x1").into()).expect(ERROR_SEND_REQUEST); + }); + + // wait for the requests to be handled + thread::sleep(Duration::from_secs(1)); + + // check current request count + let stats = handle.stats().expect(ERROR_STATS); + assert_eq!(stats, 1, "Backend should have 1 ongoing requests."); + + // Different request, should be counted + let h3 = handle.clone(); + thread::spawn(move || { + h3.get_class_hash_at(felt!("0x2").into()).expect(ERROR_SEND_REQUEST); + }); + + // wait for the requests to be handled + thread::sleep(Duration::from_secs(1)); + + // check request are handled + let stats = handle.stats().expect(ERROR_STATS); + assert_eq!(stats, 2, "Backend should only have 2 ongoing requests.") + } + + #[test] + fn get_storage_request_should_be_deduplicated() { + // start a mock remote network + start_tcp_server("127.0.0.1:8086".to_string()); + + let handle = create_forked_backend("http://127.0.0.1:8086", 1); + + // check no pending requests + let stats = handle.stats().expect(ERROR_STATS); + assert_eq!(stats, 0, "Backend should not have any ongoing requests."); + + // send requests to the backend + let h1 = handle.clone(); + thread::spawn(move || { + h1.get_storage(felt!("0x1").into(), felt!("0x1")).expect(ERROR_SEND_REQUEST); + }); + let h2 = handle.clone(); + thread::spawn(move || { + h2.get_storage(felt!("0x1").into(), felt!("0x1")).expect(ERROR_SEND_REQUEST); + }); + + // wait for the requests to be handled + thread::sleep(Duration::from_secs(1)); + + // check current request count + let stats = handle.stats().expect(ERROR_STATS); + assert_eq!(stats, 1, "Backend should have 1 ongoing requests."); + + // Different request, should be counted + let h3 = handle.clone(); + thread::spawn(move || { + h3.get_storage(felt!("0x2").into(), felt!("0x3")).expect(ERROR_SEND_REQUEST); + }); + + // wait for the requests to be handled + thread::sleep(Duration::from_secs(1)); + + // check request are handled + let stats = handle.stats().expect(ERROR_STATS); + assert_eq!(stats, 2, "Backend should only have 2 ongoing requests.") + } + + #[test] + fn get_storage_request_on_same_address_with_different_key_should_be_deduplicated() { + // start a mock remote network + start_tcp_server("127.0.0.1:8087".to_string()); + + let handle = create_forked_backend("http://127.0.0.1:8087", 1); + + // check no pending requests + let stats = handle.stats().expect(ERROR_STATS); + assert_eq!(stats, 0, "Backend should not have any ongoing requests."); + + // send requests to the backend + let h1 = handle.clone(); + thread::spawn(move || { + h1.get_storage(felt!("0x1").into(), felt!("0x1")).expect(ERROR_SEND_REQUEST); + }); + let h2 = handle.clone(); + thread::spawn(move || { + h2.get_storage(felt!("0x1").into(), felt!("0x1")).expect(ERROR_SEND_REQUEST); + }); + + // wait for the requests to be handled + thread::sleep(Duration::from_secs(1)); + + // check current request count + let stats = handle.stats().expect(ERROR_STATS); + assert_eq!(stats, 1, "Backend should have 1 ongoing requests."); + + // Different request, should be counted + let h3 = handle.clone(); + thread::spawn(move || { + h3.get_storage(felt!("0x1").into(), felt!("0x3")).expect(ERROR_SEND_REQUEST); + }); + // Different request, should be counted + let h4 = handle.clone(); + thread::spawn(move || { + h4.get_storage(felt!("0x1").into(), felt!("0x6")).expect(ERROR_SEND_REQUEST); + }); + + // wait for the requests to be handled + thread::sleep(Duration::from_secs(1)); + + // check current request count + let stats = handle.stats().expect(ERROR_STATS); + assert_eq!(stats, 3, "Backend should have 3 ongoing requests."); + + // Same request as the last one, shouldn't be counted + let h5 = handle.clone(); + thread::spawn(move || { + h5.get_storage(felt!("0x1").into(), felt!("0x6")).expect(ERROR_SEND_REQUEST); + }); + + // wait for the requests to be handled + thread::sleep(Duration::from_secs(1)); + + // check request are handled + let stats = handle.stats().expect(ERROR_STATS); + assert_eq!(stats, 3, "Backend should only have 3 ongoing requests.") + } + #[test] fn get_from_cache_if_exist() { // setup @@ -736,6 +1185,53 @@ mod tests { ); } + #[test] + fn test_deduplicated_request_should_return_similar_results() { + // Start mock server with a predefined nonce response + let response = r#"{"jsonrpc":"2.0","result":"0x123","id":1}"#; + let sender = start_mock_rpc_server("127.0.0.1:8090".to_string(), response.to_string()); + + let handle = create_forked_backend("http://127.0.0.1:8090", 1); + let addr = ContractAddress(felt!("0x1")); + + // Collect results from multiple identical nonce requests + let results: Arc>>> = + Arc::new(Mutex::new(Vec::new())); + let handles: Vec<_> = (0..5) + .map(|_| { + let h = handle.clone(); + let results = results.clone(); + thread::spawn(move || { + let res = h.get_nonce(addr); + results.lock().unwrap().push(res); + }) + }) + .collect(); + + // wait for the requests to be sent to the rpc server + thread::sleep(Duration::from_secs(1)); + + // Check that there's only one request, meaning it is deduplicated. + let stats = handle.stats().expect(ERROR_STATS); + assert_eq!(stats, 1, "Backend should only have 1 ongoing requests."); + + // Send the signal to tell the mock rpc server to return the response + sender.send(()).unwrap(); + + // Join all request threads + handles.into_iter().for_each(|h| h.join().unwrap()); + + // Verify all results are identical + let results = results.lock().unwrap(); + for result in results.iter() { + assert_eq!( + "0x123", + format!("{:#x}", result.as_ref().unwrap()), + "All deduplicated nonce requests should return the same result" + ); + } + } + // TODO: unignore this once we have separate the spawning of the backend thread from the backend // creation #[test]