From a0a950761fd8c277d7494d4b2176680bde9d8dd9 Mon Sep 17 00:00:00 2001 From: chen <23054115+cwkang1998@users.noreply.github.com> Date: Wed, 22 May 2024 16:04:35 +0800 Subject: [PATCH 01/10] feat(katana): dedup fork request --- .../provider/src/providers/fork/backend.rs | 108 ++++++++++++------ 1 file changed, 72 insertions(+), 36 deletions(-) diff --git a/crates/katana/storage/provider/src/providers/fork/backend.rs b/crates/katana/storage/provider/src/providers/fork/backend.rs index 265aa502a3..c9b9116ce5 100644 --- a/crates/katana/storage/provider/src/providers/fork/backend.rs +++ b/crates/katana/storage/provider/src/providers/fork/backend.rs @@ -1,4 +1,4 @@ -use std::collections::VecDeque; +use std::collections::{HashMap, HashSet, VecDeque}; use std::pin::Pin; use std::sync::mpsc::{ channel as oneshot, Receiver as OneshotReceiver, RecvError, Sender as OneshotSender, @@ -119,8 +119,10 @@ type BackendRequestFuture = BoxFuture<'static, ()>; pub struct Backend

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

, + // Set that keep track of current requests, for dedup purposes. + request_dedup_set: HashSet, /// Requests that are currently being poll. - pending_requests: Vec, + pending_requests: Vec<(String, BackendRequestFuture)>, /// Requests that are queued to be polled. queued_requests: VecDeque, /// A channel for receiving requests from the [BackendHandle]s. @@ -169,6 +171,7 @@ where block, incoming: rx, provider: Arc::new(provider), + request_dedup_set: HashSet::new(), pending_requests: Vec::new(), queued_requests: VecDeque::new(), }; @@ -182,63 +185,88 @@ 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); + let req_key = format!("nonce_{}", payload); - sender.send(res).expect("failed to send nonce result") - }); + if !self.request_dedup_set.contains(&req_key) { + let fut = Box::pin(async move { + let res = provider + .get_nonce(block, Felt::from(payload)) + .await + .map_err(BackendError::StarknetProvider); - self.pending_requests.push(fut); + sender.send(res).expect("failed to send nonce result") + }); + + self.pending_requests.push((req_key.clone(), fut)); + self.request_dedup_set.insert(req_key); + } } 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); + let req_key = format!("storage_{}_{}", addr, key); - sender.send(res).expect("failed to send storage result") - }); + if !self.request_dedup_set.contains(&req_key) { + let fut = Box::pin(async move { + let res = provider + .get_storage_at(Felt::from(addr), key, block) + .await + .map_err(BackendError::StarknetProvider); - self.pending_requests.push(fut); + sender.send(res).expect("failed to send storage result") + }); + + self.pending_requests.push((req_key.clone(), fut)); + self.request_dedup_set.insert(req_key); + } } 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); + let req_key = format!("classhash_{}", payload); - sender.send(res).expect("failed to send class hash result") - }); + if !self.request_dedup_set.contains(&req_key) { + let fut = Box::pin(async move { + let res = provider + .get_class_hash_at(block, Felt::from(payload)) + .await + .map_err(BackendError::StarknetProvider); - self.pending_requests.push(fut); + sender.send(res).expect("failed to send class hash result") + }); + + self.pending_requests.push((req_key.clone(), fut)); + self.request_dedup_set.insert(req_key); + } } BackendRequest::Class(Request { payload, sender }) => { - let fut = Box::pin(async move { - let res = provider - .get_class(block, payload) - .await - .map_err(BackendError::StarknetProvider); + let req_key = format!("class_{}", payload); + + if !self.request_dedup_set.contains(&req_key) { + 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") - }); + sender.send(res).expect("failed to send class result") + }); - self.pending_requests.push(fut); + self.pending_requests.push((req_key.clone(), fut)); + self.request_dedup_set.insert(req_key); + } } #[cfg(test)] BackendRequest::Stats(sender) => { + // let req_key = "stats"; + // if !self.request_dedup_set.contains(req_key) { let total_ongoing_request = self.pending_requests.len(); sender.send(total_ongoing_request).expect("failed to send backend stats"); + // self.request_dedup_set.insert(req_key.to_string()); + // } } } } @@ -275,11 +303,13 @@ 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); + pin.pending_requests.push((fut_key.clone(), fut)); + } else { + pin.request_dedup_set.remove(&fut_key); } } @@ -736,6 +766,12 @@ mod tests { ); } + #[test] + fn requests_should_be_deduped() { + let backend = create_forked_backend(LOCAL_RPC_URL, 1); + let provider = SharedStateProvider(Arc::new(CacheStateDb::new(backend))); + } + // TODO: unignore this once we have separate the spawning of the backend thread from the backend // creation #[test] From da821efdcf12169c536ee559c0268698ba860f18 Mon Sep 17 00:00:00 2001 From: chen <23054115+cwkang1998@users.noreply.github.com> Date: Sun, 26 May 2024 02:17:55 +0800 Subject: [PATCH 02/10] feat: use enum over string as dedup key --- .../provider/src/providers/fork/backend.rs | 39 +++++++++++-------- 1 file changed, 23 insertions(+), 16 deletions(-) diff --git a/crates/katana/storage/provider/src/providers/fork/backend.rs b/crates/katana/storage/provider/src/providers/fork/backend.rs index c9b9116ce5..7eb43a3c19 100644 --- a/crates/katana/storage/provider/src/providers/fork/backend.rs +++ b/crates/katana/storage/provider/src/providers/fork/backend.rs @@ -1,4 +1,4 @@ -use std::collections::{HashMap, HashSet, VecDeque}; +use std::collections::{HashSet, VecDeque}; use std::pin::Pin; use std::sync::mpsc::{ channel as oneshot, Receiver as OneshotReceiver, RecvError, Sender as OneshotSender, @@ -111,6 +111,17 @@ impl BackendRequest { type BackendRequestFuture = BoxFuture<'static, ()>; + +// 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. /// /// It is responsible for processing [requests](BackendRequest) to fetch data from the remote @@ -120,9 +131,9 @@ pub struct Backend

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

, // Set that keep track of current requests, for dedup purposes. - request_dedup_set: HashSet, + request_dedup_set: HashSet, /// Requests that are currently being poll. - pending_requests: Vec<(String, BackendRequestFuture)>, + pending_requests: Vec<(BackendRequestIdentifier, BackendRequestFuture)>, /// Requests that are queued to be polled. queued_requests: VecDeque, /// A channel for receiving requests from the [BackendHandle]s. @@ -188,7 +199,7 @@ where // Check if there are similar requests in the queue before sending the request match request { BackendRequest::Nonce(Request { payload, sender }) => { - let req_key = format!("nonce_{}", payload); + let req_key = BackendRequestIdentifier::Nonce(payload); if !self.request_dedup_set.contains(&req_key) { let fut = Box::pin(async move { @@ -200,13 +211,13 @@ where sender.send(res).expect("failed to send nonce result") }); - self.pending_requests.push((req_key.clone(), fut)); + self.pending_requests.push((req_key, fut)); self.request_dedup_set.insert(req_key); } } BackendRequest::Storage(Request { payload: (addr, key), sender }) => { - let req_key = format!("storage_{}_{}", addr, key); + let req_key = BackendRequestIdentifier::Storage((addr, key)); if !self.request_dedup_set.contains(&req_key) { let fut = Box::pin(async move { @@ -218,13 +229,13 @@ where sender.send(res).expect("failed to send storage result") }); - self.pending_requests.push((req_key.clone(), fut)); + self.pending_requests.push((req_key, fut)); self.request_dedup_set.insert(req_key); } } BackendRequest::ClassHash(Request { payload, sender }) => { - let req_key = format!("classhash_{}", payload); + let req_key = BackendRequestIdentifier::ClassHash(payload); if !self.request_dedup_set.contains(&req_key) { let fut = Box::pin(async move { @@ -236,13 +247,13 @@ where sender.send(res).expect("failed to send class hash result") }); - self.pending_requests.push((req_key.clone(), fut)); + self.pending_requests.push((req_key, fut)); self.request_dedup_set.insert(req_key); } } BackendRequest::Class(Request { payload, sender }) => { - let req_key = format!("class_{}", payload); + let req_key = BackendRequestIdentifier::Class(payload); if !self.request_dedup_set.contains(&req_key) { let fut = Box::pin(async move { @@ -254,19 +265,15 @@ where sender.send(res).expect("failed to send class result") }); - self.pending_requests.push((req_key.clone(), fut)); + self.pending_requests.push((req_key, fut)); self.request_dedup_set.insert(req_key); } } #[cfg(test)] BackendRequest::Stats(sender) => { - // let req_key = "stats"; - // if !self.request_dedup_set.contains(req_key) { let total_ongoing_request = self.pending_requests.len(); sender.send(total_ongoing_request).expect("failed to send backend stats"); - // self.request_dedup_set.insert(req_key.to_string()); - // } } } } @@ -307,7 +314,7 @@ where // 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_key.clone(), fut)); + pin.pending_requests.push((fut_key, fut)); } else { pin.request_dedup_set.remove(&fut_key); } From 2458c37bf70030c1f990de20a59edb76f00c768c Mon Sep 17 00:00:00 2001 From: chen <23054115+cwkang1998@users.noreply.github.com> Date: Sun, 26 May 2024 04:51:39 +0800 Subject: [PATCH 03/10] tests: add test for dedup --- .../provider/src/providers/fork/backend.rs | 266 +++++++++++++++++- 1 file changed, 260 insertions(+), 6 deletions(-) diff --git a/crates/katana/storage/provider/src/providers/fork/backend.rs b/crates/katana/storage/provider/src/providers/fork/backend.rs index 7eb43a3c19..fe6fffa5af 100644 --- a/crates/katana/storage/provider/src/providers/fork/backend.rs +++ b/crates/katana/storage/provider/src/providers/fork/backend.rs @@ -111,7 +111,6 @@ impl BackendRequest { type BackendRequestFuture = BoxFuture<'static, ()>; - // Identifier for pending requests. // This is used for request deduplication. #[derive(Eq, Hash, PartialEq, Clone, Copy, Debug)] @@ -119,7 +118,7 @@ enum BackendRequestIdentifier { Nonce(ContractAddress), Class(ClassHash), ClassHash(ContractAddress), - Storage((ContractAddress, StorageKey)) + Storage((ContractAddress, StorageKey)), } /// The backend for the forked provider. @@ -657,13 +656,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(); @@ -704,7 +703,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); @@ -723,7 +722,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 || { @@ -742,6 +741,261 @@ 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); + }); + // 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); + }); + // 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); + }); + // 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); + }); + // Different request, should be counted + let h3 = handle.clone(); + thread::spawn(move || { + h3.get_class_at(felt!("0x2")).expect(ERROR_SEND_REQUEST); + }); + // Different request, should be counted + let h4 = handle.clone(); + thread::spawn(move || { + h4.get_compiled_class_hash(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, 3, "Backend should only have 3 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); + }); + // 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); + }); + // 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); + }); + // 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); + }); + + // 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 From 26313448f58f61802cdb04da78d6694a9a1e43bb Mon Sep 17 00:00:00 2001 From: chen <23054115+cwkang1998@users.noreply.github.com> Date: Sun, 1 Dec 2024 08:21:30 +0800 Subject: [PATCH 04/10] feat: dedup should return result correctly --- .../provider/src/providers/fork/backend.rs | 232 ++++++++++++------ 1 file changed, 162 insertions(+), 70 deletions(-) diff --git a/crates/katana/storage/provider/src/providers/fork/backend.rs b/crates/katana/storage/provider/src/providers/fork/backend.rs index fe6fffa5af..ae86debacf 100644 --- a/crates/katana/storage/provider/src/providers/fork/backend.rs +++ b/crates/katana/storage/provider/src/providers/fork/backend.rs @@ -1,16 +1,18 @@ -use std::collections::{HashSet, VecDeque}; +use std::collections::{HashMap, VecDeque}; use std::pin::Pin; +use std::sync::Arc; use std::sync::mpsc::{ - channel as oneshot, Receiver as OneshotReceiver, RecvError, Sender as OneshotSender, + Receiver as OneshotReceiver, RecvError, Sender as OneshotSender, channel as oneshot, }; -use std::sync::Arc; use std::task::{Context, Poll}; use std::{io, thread}; -use futures::channel::mpsc::{channel as async_channel, Receiver, SendError, Sender}; +use anyhow::anyhow; +use futures::channel::mpsc::{Receiver, SendError, Sender, channel as async_channel}; use futures::future::BoxFuture; use futures::stream::Stream; use futures::{Future, FutureExt}; +use katana_primitives::Felt; use katana_primitives::block::BlockHashOrNumber; use katana_primitives::class::{ClassHash, CompiledClass, CompiledClassHash, FlattenedSierraClass}; use katana_primitives::contract::{ContractAddress, Nonce, StorageKey, StorageValue}; @@ -18,44 +20,49 @@ use katana_primitives::conversion::rpc::{ compiled_class_hash_from_flattened_sierra_class, flattened_sierra_to_compiled_class, legacy_rpc_to_compiled_class, }; -use katana_primitives::Felt; use parking_lot::Mutex; use starknet::core::types::{BlockId, ContractClass as RpcContractClass, StarknetError}; use starknet::providers::{Provider, ProviderError as StarknetProviderError}; use tracing::{error, trace}; +use crate::ProviderResult; use crate::error::ProviderError; use crate::providers::in_memory::cache::CacheStateDb; use crate::traits::contract::ContractClassProvider; use crate::traits::state::StateProvider; -use crate::ProviderResult; 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 +70,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 +81,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 +102,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 +114,7 @@ impl BackendRequest { } } -type BackendRequestFuture = BoxFuture<'static, ()>; +type BackendRequestFuture = BoxFuture<'static, BackendResponse>; // Identifier for pending requests. // This is used for request deduplication. @@ -129,8 +134,8 @@ enum BackendRequestIdentifier { pub struct Backend

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

, - // Set that keep track of current requests, for dedup purposes. - request_dedup_set: HashSet, + // HashMap that keep track of current requests, for dedup purposes. + request_dedup_map: HashMap>>, /// Requests that are currently being poll. pending_requests: Vec<(BackendRequestIdentifier, BackendRequestFuture)>, /// Requests that are queued to be polled. @@ -162,7 +167,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."); @@ -181,7 +186,7 @@ where block, incoming: rx, provider: Arc::new(provider), - request_dedup_set: HashSet::new(), + request_dedup_map: HashMap::new(), pending_requests: Vec::new(), queued_requests: VecDeque::new(), }; @@ -200,72 +205,104 @@ where BackendRequest::Nonce(Request { payload, sender }) => { let req_key = BackendRequestIdentifier::Nonce(payload); - if !self.request_dedup_set.contains(&req_key) { + if let std::collections::hash_map::Entry::Vacant(e) = + self.request_dedup_map.entry(req_key) + { let fut = Box::pin(async move { let res = provider .get_nonce(block, Felt::from(payload)) .await - .map_err(BackendError::StarknetProvider); + .map_err(|e| BackendError::StarknetProvider(Arc::new(e))); - sender.send(res).expect("failed to send nonce result") + BackendResponse::Nonce(res) }); self.pending_requests.push((req_key, fut)); - self.request_dedup_set.insert(req_key); + e.insert(vec![sender]); + } else { + let sender_vec = self + .request_dedup_map + .get_mut(&req_key) + .expect("failed to get current request dedup vector"); + sender_vec.push(sender); } } BackendRequest::Storage(Request { payload: (addr, key), sender }) => { let req_key = BackendRequestIdentifier::Storage((addr, key)); - if !self.request_dedup_set.contains(&req_key) { + if let std::collections::hash_map::Entry::Vacant(e) = + self.request_dedup_map.entry(req_key) + { let fut = Box::pin(async move { let res = provider .get_storage_at(Felt::from(addr), key, block) .await - .map_err(BackendError::StarknetProvider); + .map_err(|e| BackendError::StarknetProvider(Arc::new(e))); - sender.send(res).expect("failed to send storage result") + BackendResponse::Storage(res) }); self.pending_requests.push((req_key, fut)); - self.request_dedup_set.insert(req_key); + e.insert(vec![sender]); + } else { + let sender_vec = self + .request_dedup_map + .get_mut(&req_key) + .expect("failed to get current request dedup vector"); + sender_vec.push(sender); } } BackendRequest::ClassHash(Request { payload, sender }) => { let req_key = BackendRequestIdentifier::ClassHash(payload); - if !self.request_dedup_set.contains(&req_key) { + if let std::collections::hash_map::Entry::Vacant(e) = + self.request_dedup_map.entry(req_key) + { let fut = Box::pin(async move { let res = provider .get_class_hash_at(block, Felt::from(payload)) .await - .map_err(BackendError::StarknetProvider); + .map_err(|e| BackendError::StarknetProvider(Arc::new(e))); - sender.send(res).expect("failed to send class hash result") + BackendResponse::ClassHashAt(res) }); self.pending_requests.push((req_key, fut)); - self.request_dedup_set.insert(req_key); + e.insert(vec![sender]); + } else { + let sender_vec = self + .request_dedup_map + .get_mut(&req_key) + .expect("failed to get current request dedup vector"); + sender_vec.push(sender); } } BackendRequest::Class(Request { payload, sender }) => { let req_key = BackendRequestIdentifier::Class(payload); - if !self.request_dedup_set.contains(&req_key) { + if let std::collections::hash_map::Entry::Vacant(e) = + self.request_dedup_map.entry(req_key) + { let fut = Box::pin(async move { let res = provider .get_class(block, payload) .await - .map_err(BackendError::StarknetProvider); + .map_err(|e| BackendError::StarknetProvider(Arc::new(e))); - sender.send(res).expect("failed to send class result") + BackendResponse::ClassAt(res) }); self.pending_requests.push((req_key, fut)); - self.request_dedup_set.insert(req_key); + e.insert(vec![sender]); + } else { + let sender_vec = self + .request_dedup_map + .get_mut(&req_key) + .expect("failed to get current request dedup vector"); + sender_vec.push(sender); } } @@ -312,10 +349,25 @@ where 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_key, fut)); - } else { - pin.request_dedup_set.remove(&fut_key); + 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()).expect( + format!("failed to send result of request {:?}", fut_key).as_str(), + ); + }); + + pin.request_dedup_map.remove(&fut_key); + } } } @@ -345,7 +397,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( @@ -356,21 +413,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( @@ -385,7 +457,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))) } } } @@ -628,9 +700,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), } @@ -642,8 +717,8 @@ pub(crate) mod test_utils { use std::sync::mpsc::sync_channel; use katana_primitives::block::BlockNumber; - use starknet::providers::jsonrpc::HttpTransport; use starknet::providers::JsonRpcClient; + use starknet::providers::jsonrpc::HttpTransport; use tokio::net::TcpListener; use url::Url; @@ -676,6 +751,29 @@ pub(crate) mod test_utils { rx.recv().unwrap(); } + + // Starts a mocked starknet rpc server that never close the connection. + // The server will only accept certain function calls. + pub fn start_mock_starknet_rpc_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(addr).await.unwrap(); + let mut connections = Vec::new(); + + tx.send(()).unwrap(); + + loop { + let (socket, _) = listener.accept().await.unwrap(); + connections.push(socket); + } + }); + }); + + rx.recv().unwrap(); + } } #[cfg(test)] @@ -1009,10 +1107,10 @@ mod tests { .or_default() .insert(STORAGE_KEY, ADDR_1_STORAGE_VALUE); - state_db.contract_state.write().insert( - ADDR_1, - GenericContractInfo { nonce: ADDR_1_NONCE, class_hash: ADDR_1_CLASS_HASH }, - ); + state_db.contract_state.write().insert(ADDR_1, GenericContractInfo { + nonce: ADDR_1_NONCE, + class_hash: ADDR_1_CLASS_HASH, + }); let provider = SharedStateProvider(Arc::new(state_db)); @@ -1027,12 +1125,6 @@ mod tests { ); } - #[test] - fn requests_should_be_deduped() { - let backend = create_forked_backend(LOCAL_RPC_URL, 1); - let provider = SharedStateProvider(Arc::new(CacheStateDb::new(backend))); - } - // TODO: unignore this once we have separate the spawning of the backend thread from the backend // creation #[test] From c2157358a2c2f404af594a9ded423bd90d24c426 Mon Sep 17 00:00:00 2001 From: chen <23054115+cwkang1998@users.noreply.github.com> Date: Tue, 10 Dec 2024 05:11:49 +0800 Subject: [PATCH 05/10] chore: add tests for deduped response --- .../provider/src/providers/fork/backend.rs | 111 ++++++++++++++---- 1 file changed, 88 insertions(+), 23 deletions(-) diff --git a/crates/katana/storage/provider/src/providers/fork/backend.rs b/crates/katana/storage/provider/src/providers/fork/backend.rs index ae86debacf..9b95d55371 100644 --- a/crates/katana/storage/provider/src/providers/fork/backend.rs +++ b/crates/katana/storage/provider/src/providers/fork/backend.rs @@ -714,11 +714,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::{SyncSender, sync_channel}; use katana_primitives::block::BlockNumber; use starknet::providers::JsonRpcClient; use starknet::providers::jsonrpc::HttpTransport; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::TcpListener; use url::Url; @@ -752,33 +753,50 @@ pub(crate) mod test_utils { rx.recv().unwrap(); } - // Starts a mocked starknet rpc server that never close the connection. - // The server will only accept certain function calls. - pub fn start_mock_starknet_rpc_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(addr).await.unwrap(); - let mut connections = Vec::new(); - - tx.send(()).unwrap(); - - loop { - let (socket, _) = listener.accept().await.unwrap(); - connections.push(socket); - } - }); + // 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\n\ + content-length: {}\r\n\ + content-type: application/json\r\n\ + \r\n\ + {}", + response.len(), + response + ); + + socket.write_all(http_response.as_bytes()).await.unwrap(); + socket.flush().await.unwrap(); + } }); - - rx.recv().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; @@ -1125,6 +1143,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] From 471abb8566d625a5dcbecbc39f93ec1f07277e73 Mon Sep 17 00:00:00 2001 From: chen <23054115+cwkang1998@users.noreply.github.com> Date: Tue, 10 Dec 2024 05:22:06 +0800 Subject: [PATCH 06/10] chore: add error handling instead of expect --- .../provider/src/providers/fork/backend.rs | 64 ++++++++++++------- 1 file changed, 41 insertions(+), 23 deletions(-) diff --git a/crates/katana/storage/provider/src/providers/fork/backend.rs b/crates/katana/storage/provider/src/providers/fork/backend.rs index 9b95d55371..850107525b 100644 --- a/crates/katana/storage/provider/src/providers/fork/backend.rs +++ b/crates/katana/storage/provider/src/providers/fork/backend.rs @@ -220,11 +220,16 @@ where self.pending_requests.push((req_key, fut)); e.insert(vec![sender]); } else { - let sender_vec = self - .request_dedup_map - .get_mut(&req_key) - .expect("failed to get current request dedup vector"); - sender_vec.push(sender); + 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"); + } + } } } @@ -246,11 +251,16 @@ where self.pending_requests.push((req_key, fut)); e.insert(vec![sender]); } else { - let sender_vec = self - .request_dedup_map - .get_mut(&req_key) - .expect("failed to get current request dedup vector"); - sender_vec.push(sender); + 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"); + } + } } } @@ -272,11 +282,16 @@ where self.pending_requests.push((req_key, fut)); e.insert(vec![sender]); } else { - let sender_vec = self - .request_dedup_map - .get_mut(&req_key) - .expect("failed to get current request dedup vector"); - sender_vec.push(sender); + 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"); + } + } } } @@ -298,11 +313,16 @@ where self.pending_requests.push((req_key, fut)); e.insert(vec![sender]); } else { - let sender_vec = self - .request_dedup_map - .get_mut(&req_key) - .expect("failed to get current request dedup vector"); - sender_vec.push(sender); + 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"); + } + } } } @@ -361,9 +381,7 @@ where // Send the response to all the senders waiting on the same request sender_vec.iter().for_each(|sender| { - sender.send(res.clone()).expect( - format!("failed to send result of request {:?}", fut_key).as_str(), - ); + sender.send(res.clone()).unwrap_or_else(|_| error!(target: LOG_TARGET, "failed to send result of request {:?} to sender {:?}", fut_key, sender)); }); pin.request_dedup_map.remove(&fut_key); From fb4b08d22b6c9e9acfb3c609ce27c7dd50c242c7 Mon Sep 17 00:00:00 2001 From: chen <23054115+cwkang1998@users.noreply.github.com> Date: Tue, 10 Dec 2024 05:46:20 +0800 Subject: [PATCH 07/10] chore: refactor dedup logic into fn --- .../provider/src/providers/fork/backend.rs | 124 ++++++------------ 1 file changed, 43 insertions(+), 81 deletions(-) diff --git a/crates/katana/storage/provider/src/providers/fork/backend.rs b/crates/katana/storage/provider/src/providers/fork/backend.rs index 850107525b..d1f3fb9a4b 100644 --- a/crates/katana/storage/provider/src/providers/fork/backend.rs +++ b/crates/katana/storage/provider/src/providers/fork/backend.rs @@ -205,125 +205,60 @@ where BackendRequest::Nonce(Request { payload, sender }) => { let req_key = BackendRequestIdentifier::Nonce(payload); - if let std::collections::hash_map::Entry::Vacant(e) = - self.request_dedup_map.entry(req_key) - { - let fut = Box::pin(async move { + self.dedup_request(req_key, sender, move || { + 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) - }); - - self.pending_requests.push((req_key, fut)); - 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"); - } - } - } + }) + }); } BackendRequest::Storage(Request { payload: (addr, key), sender }) => { let req_key = BackendRequestIdentifier::Storage((addr, key)); - if let std::collections::hash_map::Entry::Vacant(e) = - self.request_dedup_map.entry(req_key) - { - let fut = Box::pin(async move { + self.dedup_request(req_key, sender, move || { + 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) - }); - - self.pending_requests.push((req_key, fut)); - 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"); - } - } - } + }) + }); } BackendRequest::ClassHash(Request { payload, sender }) => { let req_key = BackendRequestIdentifier::ClassHash(payload); - if let std::collections::hash_map::Entry::Vacant(e) = - self.request_dedup_map.entry(req_key) - { - let fut = Box::pin(async move { + self.dedup_request(req_key, sender, move || { + 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) - }); - - self.pending_requests.push((req_key, fut)); - 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"); - } - } - } + }) + }); } BackendRequest::Class(Request { payload, sender }) => { let req_key = BackendRequestIdentifier::Class(payload); - if let std::collections::hash_map::Entry::Vacant(e) = - self.request_dedup_map.entry(req_key) - { - let fut = Box::pin(async move { + self.dedup_request(req_key, sender, move || { + Box::pin(async move { let res = provider .get_class(block, payload) .await .map_err(|e| BackendError::StarknetProvider(Arc::new(e))); BackendResponse::ClassAt(res) - }); - - self.pending_requests.push((req_key, fut)); - 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"); - } - } - } + }) + }); } #[cfg(test)] @@ -333,6 +268,33 @@ where } } } + + fn dedup_request( + &mut self, + req_key: BackendRequestIdentifier, + sender: OneshotSender, + rpc_call_future: F, + ) where + F: FnOnce() -> BoxFuture<'static, BackendResponse>, + { + if let std::collections::hash_map::Entry::Vacant(e) = self.request_dedup_map.entry(req_key) + { + let fut = rpc_call_future(); + self.pending_requests.push((req_key, fut)); + 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

From d6da703fdb05a8ddb7b2d356df58989ecae9d0f7 Mon Sep 17 00:00:00 2001 From: chen <23054115+cwkang1998@users.noreply.github.com> Date: Thu, 19 Dec 2024 03:51:33 +0800 Subject: [PATCH 08/10] chore: refactor tests & apply fmting --- .../provider/src/providers/fork/backend.rs | 126 ++++++++++++------ 1 file changed, 82 insertions(+), 44 deletions(-) diff --git a/crates/katana/storage/provider/src/providers/fork/backend.rs b/crates/katana/storage/provider/src/providers/fork/backend.rs index d1f3fb9a4b..b6adba5095 100644 --- a/crates/katana/storage/provider/src/providers/fork/backend.rs +++ b/crates/katana/storage/provider/src/providers/fork/backend.rs @@ -1,18 +1,18 @@ +use std::collections::hash_map::Entry; use std::collections::{HashMap, VecDeque}; use std::pin::Pin; -use std::sync::Arc; use std::sync::mpsc::{ - Receiver as OneshotReceiver, RecvError, Sender as OneshotSender, channel as oneshot, + channel as oneshot, Receiver as OneshotReceiver, RecvError, Sender as OneshotSender, }; +use std::sync::Arc; use std::task::{Context, Poll}; use std::{io, thread}; use anyhow::anyhow; -use futures::channel::mpsc::{Receiver, SendError, Sender, channel as async_channel}; +use futures::channel::mpsc::{channel as async_channel, Receiver, SendError, Sender}; use futures::future::BoxFuture; use futures::stream::Stream; use futures::{Future, FutureExt}; -use katana_primitives::Felt; use katana_primitives::block::BlockHashOrNumber; use katana_primitives::class::{ClassHash, CompiledClass, CompiledClassHash, FlattenedSierraClass}; use katana_primitives::contract::{ContractAddress, Nonce, StorageKey, StorageValue}; @@ -20,16 +20,17 @@ use katana_primitives::conversion::rpc::{ compiled_class_hash_from_flattened_sierra_class, flattened_sierra_to_compiled_class, legacy_rpc_to_compiled_class, }; +use katana_primitives::Felt; use parking_lot::Mutex; use starknet::core::types::{BlockId, ContractClass as RpcContractClass, StarknetError}; use starknet::providers::{Provider, ProviderError as StarknetProviderError}; use tracing::{error, trace}; -use crate::ProviderResult; use crate::error::ProviderError; use crate::providers::in_memory::cache::CacheStateDb; use crate::traits::contract::ContractClassProvider; use crate::traits::state::StateProvider; +use crate::ProviderResult; const LOG_TARGET: &str = "forking::backend"; @@ -205,21 +206,25 @@ where BackendRequest::Nonce(Request { payload, sender }) => { let req_key = BackendRequestIdentifier::Nonce(payload); - self.dedup_request(req_key, sender, move || { + 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 req_key = BackendRequestIdentifier::Storage((addr, key)); - self.dedup_request(req_key, sender, move || { + self.dedup_request( + req_key, + sender, Box::pin(async move { let res = provider .get_storage_at(Felt::from(addr), key, block) @@ -227,14 +232,16 @@ where .map_err(|e| BackendError::StarknetProvider(Arc::new(e))); BackendResponse::Storage(res) - }) - }); + }), + ); } BackendRequest::ClassHash(Request { payload, sender }) => { let req_key = BackendRequestIdentifier::ClassHash(payload); - self.dedup_request(req_key, sender, move || { + self.dedup_request( + req_key, + sender, Box::pin(async move { let res = provider .get_class_hash_at(block, Felt::from(payload)) @@ -242,14 +249,16 @@ where .map_err(|e| BackendError::StarknetProvider(Arc::new(e))); BackendResponse::ClassHashAt(res) - }) - }); + }), + ); } BackendRequest::Class(Request { payload, sender }) => { let req_key = BackendRequestIdentifier::Class(payload); - self.dedup_request(req_key, sender, move || { + self.dedup_request( + req_key, + sender, Box::pin(async move { let res = provider .get_class(block, payload) @@ -257,8 +266,8 @@ where .map_err(|e| BackendError::StarknetProvider(Arc::new(e))); BackendResponse::ClassAt(res) - }) - }); + }), + ); } #[cfg(test)] @@ -269,18 +278,14 @@ where } } - fn dedup_request( + fn dedup_request( &mut self, req_key: BackendRequestIdentifier, sender: OneshotSender, - rpc_call_future: F, - ) where - F: FnOnce() -> BoxFuture<'static, BackendResponse>, - { - if let std::collections::hash_map::Entry::Vacant(e) = self.request_dedup_map.entry(req_key) - { - let fut = rpc_call_future(); - self.pending_requests.push((req_key, fut)); + 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) { @@ -343,7 +348,9 @@ where // 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!(target: LOG_TARGET, "failed to send result of request {:?} to sender {:?}", fut_key, 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); @@ -694,11 +701,11 @@ fn handle_not_found_err(result: Result) -> Result, #[cfg(test)] pub(crate) mod test_utils { - use std::sync::mpsc::{SyncSender, sync_channel}; + use std::sync::mpsc::{sync_channel, SyncSender}; use katana_primitives::block::BlockNumber; - use starknet::providers::JsonRpcClient; use starknet::providers::jsonrpc::HttpTransport; + use starknet::providers::JsonRpcClient; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::TcpListener; use url::Url; @@ -754,11 +761,8 @@ pub(crate) mod test_utils { // After reading, we send the pre-determined response let http_response = format!( - "HTTP/1.1 200 OK\r\n\ - content-length: {}\r\n\ - content-type: application/json\r\n\ - \r\n\ - {}", + "HTTP/1.1 200 OK\r\ncontent-length: {}\r\ncontent-type: \ + application/json\r\n\r\n{}", response.len(), response ); @@ -857,6 +861,11 @@ mod tests { thread::spawn(move || { h2.get_nonce(felt!("0x1").into()).expect(ERROR_SEND_REQUEST); }); + + // 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 || { @@ -891,6 +900,11 @@ mod tests { thread::spawn(move || { h2.get_class_at(felt!("0x1")).expect(ERROR_SEND_REQUEST); }); + + // check current request count + let stats = handle.stats().expect(ERROR_STATS); + assert_eq!(stats, 2, "Backend should have 1 ongoing requests."); + // Different request, should be counted let h3 = handle.clone(); thread::spawn(move || { @@ -925,6 +939,11 @@ mod tests { thread::spawn(move || { h2.get_compiled_class_hash(felt!("0x1")).expect(ERROR_SEND_REQUEST); }); + + // 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 || { @@ -960,23 +979,23 @@ mod tests { thread::spawn(move || { h2.get_compiled_class_hash(felt!("0x1")).expect(ERROR_SEND_REQUEST); }); + + // 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); }); - // Different request, should be counted - let h4 = handle.clone(); - thread::spawn(move || { - h4.get_compiled_class_hash(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, 3, "Backend should only have 3 ongoing requests.") + assert_eq!(stats, 1, "Backend should only have 2 ongoing requests.") } #[test] @@ -999,6 +1018,11 @@ mod tests { thread::spawn(move || { h2.get_class_hash_at(felt!("0x1").into()).expect(ERROR_SEND_REQUEST); }); + + // 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 || { @@ -1033,6 +1057,11 @@ mod tests { thread::spawn(move || { h2.get_storage(felt!("0x1").into(), felt!("0x1")).expect(ERROR_SEND_REQUEST); }); + + // 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 || { @@ -1067,6 +1096,11 @@ mod tests { thread::spawn(move || { h2.get_storage(felt!("0x1").into(), felt!("0x1")).expect(ERROR_SEND_REQUEST); }); + + // 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 || { @@ -1078,6 +1112,10 @@ mod tests { h4.get_storage(felt!("0x1").into(), felt!("0x6")).expect(ERROR_SEND_REQUEST); }); + // 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 || { @@ -1105,10 +1143,10 @@ mod tests { .or_default() .insert(STORAGE_KEY, ADDR_1_STORAGE_VALUE); - state_db.contract_state.write().insert(ADDR_1, GenericContractInfo { - nonce: ADDR_1_NONCE, - class_hash: ADDR_1_CLASS_HASH, - }); + state_db.contract_state.write().insert( + ADDR_1, + GenericContractInfo { nonce: ADDR_1_NONCE, class_hash: ADDR_1_CLASS_HASH }, + ); let provider = SharedStateProvider(Arc::new(state_db)); From d2c624c8e631aa7bbe52f39d8539f6990ba0e847 Mon Sep 17 00:00:00 2001 From: Ammar Arif Date: Thu, 19 Dec 2024 10:08:46 -0500 Subject: [PATCH 09/10] fix tests --- .../provider/src/providers/fork/backend.rs | 27 ++++++++++++++++--- 1 file changed, 24 insertions(+), 3 deletions(-) diff --git a/crates/katana/storage/provider/src/providers/fork/backend.rs b/crates/katana/storage/provider/src/providers/fork/backend.rs index b6adba5095..02ef6fa603 100644 --- a/crates/katana/storage/provider/src/providers/fork/backend.rs +++ b/crates/katana/storage/provider/src/providers/fork/backend.rs @@ -862,6 +862,9 @@ mod tests { 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."); @@ -901,9 +904,12 @@ mod tests { 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, 2, "Backend should have 1 ongoing requests."); + assert_eq!(stats, 1, "Backend should have 1 ongoing requests."); // Different request, should be counted let h3 = handle.clone(); @@ -940,6 +946,9 @@ mod tests { 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."); @@ -980,6 +989,9 @@ mod tests { 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."); @@ -995,7 +1007,7 @@ mod tests { // check request are handled let stats = handle.stats().expect(ERROR_STATS); - assert_eq!(stats, 1, "Backend should only have 2 ongoing requests.") + assert_eq!(stats, 2, "Backend should only have 2 ongoing requests.") } #[test] @@ -1019,6 +1031,9 @@ mod tests { 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."); @@ -1058,6 +1073,9 @@ mod tests { 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."); @@ -1097,6 +1115,9 @@ mod tests { 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."); @@ -1114,7 +1135,7 @@ mod tests { // check current request count let stats = handle.stats().expect(ERROR_STATS); - assert_eq!(stats, 3, "Backend should have 3 ongoing requests."); + assert_eq!(stats, 1, "Backend should have 1 ongoing requests."); // Same request as the last one, shouldn't be counted let h5 = handle.clone(); From ec163c5706f8f0496c6a6ed2ee2183327301bb3a Mon Sep 17 00:00:00 2001 From: Ammar Arif Date: Thu, 19 Dec 2024 10:16:45 -0500 Subject: [PATCH 10/10] fix test --- crates/katana/storage/provider/src/providers/fork/backend.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/crates/katana/storage/provider/src/providers/fork/backend.rs b/crates/katana/storage/provider/src/providers/fork/backend.rs index 02ef6fa603..64e921ecdc 100644 --- a/crates/katana/storage/provider/src/providers/fork/backend.rs +++ b/crates/katana/storage/provider/src/providers/fork/backend.rs @@ -1133,9 +1133,12 @@ mod tests { 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, 1, "Backend should have 1 ongoing requests."); + assert_eq!(stats, 3, "Backend should have 3 ongoing requests."); // Same request as the last one, shouldn't be counted let h5 = handle.clone();