Skip to content

Commit 6af545f

Browse files
itamar-starkwareitamar-starkware
committed
apollo_gateway: use get_nonce outside extract_state_nonce_and_run_validations
1 parent 2bab9a0 commit 6af545f

File tree

8 files changed

+128
-54
lines changed

8 files changed

+128
-54
lines changed

‎crates/apollo_gateway/src/gateway.rs‎

Lines changed: 62 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,11 @@ use crate::errors::{
4242
GatewayResult,
4343
};
4444
use crate::metrics::{register_metrics, GatewayMetricHandle, GATEWAY_ADD_TX_LATENCY};
45-
use crate::state_reader::StateReaderFactory;
45+
use crate::state_reader::{
46+
GatewayStateReaderWithCompiledClasses,
47+
MempoolStateReader,
48+
StateReaderFactory,
49+
};
4650
use crate::stateful_transaction_validator::{
4751
StatefulTransactionValidatorFactory,
4852
StatefulTransactionValidatorFactoryTrait,
@@ -155,7 +159,7 @@ impl Gateway {
155159
transaction_converter_err_to_deprecated_gw_err(&tx_signature, e)
156160
})?;
157161

158-
let blocking_task =
162+
let mut blocking_task =
159163
ProcessTxBlockingTask::new(self, executable_tx, tokio::runtime::Handle::current())
160164
.await
161165
.map_err(|e| {
@@ -166,10 +170,22 @@ impl Gateway {
166170
metric_counters.record_add_tx_failure(&e);
167171
e
168172
})?;
173+
174+
let state_reader = blocking_task.get_state_reader().await?;
175+
176+
let account_address = blocking_task.executable_tx.contract_address();
177+
let account_nonce = state_reader
178+
.get_account_nonce(account_address)
179+
.await
180+
.map_err(|e| StarknetError::internal_with_logging("Failed to get account nonce", e))?;
181+
182+
let stateful_tx_validator = blocking_task.create_stateful_validator(state_reader).await?;
183+
169184
// Run the blocking task in the current span.
170185
let curr_span = Span::current();
171-
let handle =
172-
tokio::task::spawn_blocking(move || curr_span.in_scope(|| blocking_task.process_tx()));
186+
let handle = tokio::task::spawn_blocking(move || {
187+
curr_span.in_scope(|| blocking_task.process_tx(account_nonce, stateful_tx_validator))
188+
});
173189
let handle_result = handle.await;
174190
let nonce = match handle_result {
175191
Ok(Ok(nonce)) => nonce,
@@ -244,7 +260,8 @@ impl Gateway {
244260
/// CPU-intensive transaction processing, spawned in a blocking thread to avoid blocking other tasks
245261
/// from running.
246262
struct ProcessTxBlockingTask {
247-
stateful_tx_validator: Box<dyn StatefulTransactionValidatorTrait + Send>,
263+
state_reader_factory: Arc<dyn StateReaderFactory>,
264+
stateful_tx_validator_factory: Arc<dyn StatefulTransactionValidatorFactoryTrait>,
248265
mempool_client: SharedMempoolClient,
249266
executable_tx: AccountTransaction,
250267
runtime: tokio::runtime::Handle,
@@ -256,21 +273,54 @@ impl ProcessTxBlockingTask {
256273
executable_tx: AccountTransaction,
257274
runtime: tokio::runtime::Handle,
258275
) -> GatewayResult<Self> {
259-
let stateful_tx_validator = gateway
260-
.stateful_tx_validator_factory
261-
.instantiate_validator(gateway.state_reader_factory.as_ref())
262-
.await?;
263276
Ok(Self {
264-
stateful_tx_validator,
277+
state_reader_factory: gateway.state_reader_factory.clone(),
278+
stateful_tx_validator_factory: gateway.stateful_tx_validator_factory.clone(),
265279
mempool_client: gateway.mempool_client.clone(),
266280
executable_tx,
267281
runtime,
268282
})
269283
}
270284

271-
fn process_tx(mut self) -> GatewayResult<Nonce> {
272-
let nonce = self.stateful_tx_validator.extract_state_nonce_and_run_validations(
285+
pub async fn get_state_reader(
286+
&self,
287+
) -> GatewayResult<Box<dyn GatewayStateReaderWithCompiledClasses>> {
288+
let state_reader = self
289+
.stateful_tx_validator_factory
290+
.get_state_reader_for_validation(self.state_reader_factory.as_ref())
291+
.await
292+
.map_err(|e| {
293+
StarknetError::internal_with_logging(
294+
"Failed to get state reader from latest block",
295+
e,
296+
)
297+
})?;
298+
Ok(state_reader)
299+
}
300+
301+
pub async fn create_stateful_validator(
302+
&mut self,
303+
state_reader: Box<dyn GatewayStateReaderWithCompiledClasses>,
304+
) -> GatewayResult<Box<dyn StatefulTransactionValidatorTrait>> {
305+
self.stateful_tx_validator_factory
306+
.create_validator_from_state_reader(state_reader)
307+
.await
308+
.map_err(|e| {
309+
StarknetError::internal_with_logging(
310+
"Failed to create stateful validator from state reader",
311+
e,
312+
)
313+
})
314+
}
315+
316+
fn process_tx(
317+
self,
318+
account_nonce: Nonce,
319+
mut stateful_tx_validator: Box<dyn StatefulTransactionValidatorTrait>,
320+
) -> GatewayResult<Nonce> {
321+
let nonce = stateful_tx_validator.extract_state_nonce_and_run_validations(
273322
&self.executable_tx,
323+
account_nonce,
274324
self.mempool_client,
275325
self.runtime,
276326
)?;

‎crates/apollo_gateway/src/gateway_test.rs‎

Lines changed: 12 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -98,8 +98,6 @@ use crate::state_reader_test_utils::{local_test_state_reader_factory, TestStateR
9898
use crate::stateful_transaction_validator::{
9999
MockStatefulTransactionValidatorFactoryTrait,
100100
MockStatefulTransactionValidatorTrait,
101-
StatefulTransactionValidatorFactoryTrait,
102-
StatefulTransactionValidatorTrait,
103101
};
104102
use crate::stateless_transaction_validator::MockStatelessTransactionValidatorTrait;
105103

@@ -318,17 +316,12 @@ async fn run_add_tx_and_extract_metrics(
318316
}
319317

320318
async fn process_tx_task(
321-
stateful_transaction_validator_factory: MockStatefulTransactionValidatorFactoryTrait,
319+
stateful_tx_validator_factory: MockStatefulTransactionValidatorFactoryTrait,
322320
) -> ProcessTxBlockingTask {
323321
let state_reader_factory = Arc::new(MockStateReaderFactory::new());
324-
let stateful_tx_validator = StatefulTransactionValidatorFactoryTrait::instantiate_validator(
325-
&stateful_transaction_validator_factory,
326-
state_reader_factory.as_ref(),
327-
)
328-
.await
329-
.expect("instantiate_validator should be mocked in tests");
330322
ProcessTxBlockingTask {
331-
stateful_tx_validator,
323+
state_reader_factory,
324+
stateful_tx_validator_factory: Arc::new(stateful_tx_validator_factory),
332325
mempool_client: Arc::new(MockMempoolClient::new()),
333326
executable_tx: executable_invoke_tx(invoke_args()),
334327
runtime: tokio::runtime::Handle::current(),
@@ -560,7 +553,7 @@ fn test_full_cycle_dump_deserialize_authorized_declarer_accounts(
560553
async fn process_tx_returns_error_when_extract_state_nonce_and_run_validations_fails(
561554
#[case] error_code: StarknetErrorCode,
562555
mut mock_stateful_transaction_validator: MockStatefulTransactionValidatorTrait,
563-
mut mock_stateful_transaction_validator_factory: MockStatefulTransactionValidatorFactoryTrait,
556+
mock_stateful_transaction_validator_factory: MockStatefulTransactionValidatorFactoryTrait,
564557
) {
565558
let expected_error = StarknetError {
566559
code: error_code.clone(),
@@ -569,19 +562,15 @@ async fn process_tx_returns_error_when_extract_state_nonce_and_run_validations_f
569562

570563
mock_stateful_transaction_validator
571564
.expect_extract_state_nonce_and_run_validations()
572-
.return_once(|_, _, _| Err(expected_error));
573-
574-
mock_stateful_transaction_validator_factory.expect_instantiate_validator().return_once(|_| {
575-
Box::pin(async {
576-
Ok::<Box<dyn StatefulTransactionValidatorTrait>, _>(Box::new(
577-
mock_stateful_transaction_validator,
578-
))
579-
})
580-
});
565+
.return_once(|_, _, _, _| Err(expected_error));
581566

582567
let process_tx_task = process_tx_task(mock_stateful_transaction_validator_factory).await;
583-
584-
let result = tokio::task::spawn_blocking(move || process_tx_task.process_tx()).await.unwrap();
568+
let account_nonce = nonce!(0);
569+
let result = tokio::task::spawn_blocking(move || {
570+
process_tx_task.process_tx(account_nonce, Box::new(mock_stateful_transaction_validator))
571+
})
572+
.await
573+
.unwrap();
585574

586575
assert!(result.is_err());
587576
assert_eq!(result.unwrap_err().code, error_code);
@@ -622,7 +611,7 @@ async fn add_tx_returns_error_when_instantiating_validator_fails(
622611
let error_code = StarknetErrorCode::UnknownErrorCode("StarknetErrorCode.InternalError".into());
623612
let expected_error = StarknetError { code: error_code.clone(), message: "placeholder".into() };
624613
mock_stateful_transaction_validator_factory
625-
.expect_instantiate_validator()
614+
.expect_get_state_reader_for_validation()
626615
.return_once(|_| Box::pin(async { Err::<_, _>(expected_error) }));
627616

628617
// Build gateway and inject the failing factory.

‎crates/apollo_gateway/src/rpc_state_reader.rs‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -117,6 +117,19 @@ impl MempoolStateReader for RpcStateReader {
117117
let block_info = block_header.try_into()?;
118118
Ok(block_info)
119119
}
120+
121+
async fn get_account_nonce(&self, contract_address: ContractAddress) -> StateResult<Nonce> {
122+
let get_nonce_params = GetNonceParams { block_id: self.block_id, contract_address };
123+
let result = self.send_rpc_request("starknet_getNonce", get_nonce_params);
124+
match result {
125+
Ok(value) => {
126+
let nonce: Nonce = serde_json::from_value(value).map_err(serde_err_to_state_err)?;
127+
Ok(nonce)
128+
}
129+
Err(RPCStateReaderError::ContractAddressNotFound(_)) => Ok(Nonce::default()),
130+
Err(e) => Err(e)?,
131+
}
132+
}
120133
}
121134

122135
impl BlockifierStateReader for RpcStateReader {

‎crates/apollo_gateway/src/state_reader.rs‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ use starknet_types_core::felt::Felt;
1717
#[async_trait]
1818
pub trait MempoolStateReader: BlockifierStateReader + Send + Sync {
1919
async fn get_block_info(&self) -> StateResult<BlockInfo>;
20+
async fn get_account_nonce(&self, contract_address: ContractAddress) -> StateResult<Nonce>;
2021
}
2122

2223
#[async_trait]
@@ -75,6 +76,10 @@ impl MempoolStateReader for Box<dyn GatewayStateReaderWithCompiledClasses> {
7576
async fn get_block_info(&self) -> StateResult<BlockInfo> {
7677
self.as_ref().get_block_info().await
7778
}
79+
80+
async fn get_account_nonce(&self, contract_address: ContractAddress) -> StateResult<Nonce> {
81+
self.as_ref().get_account_nonce(contract_address).await
82+
}
7883
}
7984

8085
impl GatewayStateReaderWithCompiledClasses for Box<dyn GatewayStateReaderWithCompiledClasses> {}
@@ -86,4 +91,8 @@ impl MempoolStateReader
8691
async fn get_block_info(&self) -> StateResult<BlockInfo> {
8792
self.state_reader.get_block_info().await
8893
}
94+
95+
async fn get_account_nonce(&self, contract_address: ContractAddress) -> StateResult<Nonce> {
96+
self.state_reader.get_account_nonce(contract_address).await
97+
}
8998
}

‎crates/apollo_gateway/src/state_reader_test_utils.rs‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,13 @@ impl MempoolStateReader for TestStateReader {
3434
async fn get_block_info(&self) -> Result<BlockInfo, StateError> {
3535
Ok(self.block_info.clone())
3636
}
37+
38+
async fn get_account_nonce(
39+
&self,
40+
contract_address: ContractAddress,
41+
) -> Result<Nonce, StateError> {
42+
self.blockifier_state_reader.get_nonce_at(contract_address)
43+
}
3744
}
3845

3946
impl FetchCompiledClasses for TestStateReader {

‎crates/apollo_gateway/src/stateful_transaction_validator.rs‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -55,11 +55,12 @@ pub trait StatefulTransactionValidatorFactoryTrait: Send + Sync {
5555
state_reader_factory: &dyn StateReaderFactory,
5656
) -> StatefulTransactionValidatorResult<Box<dyn StatefulTransactionValidatorTrait>>;
5757

58+
// 1) Get a state reader for validation (async, uses the factory).
5859
async fn get_state_reader_for_validation(
5960
&self,
6061
state_reader_factory: &dyn StateReaderFactory,
6162
) -> StatefulTransactionValidatorResult<Box<dyn GatewayStateReaderWithCompiledClasses>>;
62-
63+
// 2) Create the validator from a provided state reader (async, uses ChainInfo/config).
6364
async fn create_validator_from_state_reader(
6465
&self,
6566
state_reader: Box<dyn GatewayStateReaderWithCompiledClasses>,
@@ -145,6 +146,7 @@ pub trait StatefulTransactionValidatorTrait: Send {
145146
fn extract_state_nonce_and_run_validations(
146147
&mut self,
147148
executable_tx: &ExecutableTransaction,
149+
account_nonce: Nonce,
148150
mempool_client: SharedMempoolClient,
149151
runtime: tokio::runtime::Handle,
150152
) -> StatefulTransactionValidatorResult<Nonce>;
@@ -166,11 +168,10 @@ impl<B: BlockifierStatefulValidatorTrait + Send> StatefulTransactionValidatorTra
166168
fn extract_state_nonce_and_run_validations(
167169
&mut self,
168170
executable_tx: &ExecutableTransaction,
171+
account_nonce: Nonce,
169172
mempool_client: SharedMempoolClient,
170173
runtime: tokio::runtime::Handle,
171174
) -> StatefulTransactionValidatorResult<Nonce> {
172-
let address = executable_tx.contract_address();
173-
let account_nonce = self.get_nonce(address)?;
174175
self.run_transaction_validations(executable_tx, account_nonce, mempool_client, runtime)?;
175176
Ok(account_nonce)
176177
}

‎crates/apollo_gateway/src/stateful_transaction_validator_test.rs‎

Lines changed: 9 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -62,32 +62,24 @@ async fn test_get_nonce_fail_on_extract_state_nonce_and_run_validations() {
6262
// The mempool does not have any transactions from the sender.
6363
Ok(false)
6464
});
65-
let mempool_client = Arc::new(mock_mempool_client);
66-
let runtime = tokio::runtime::Handle::current();
65+
// Removed unused mempool client and runtime in this test since nonce is fetched outside now.
6766

6867
let mut stateful_validator = StatefulTransactionValidator {
6968
config: StatefulTransactionValidatorConfig::default(),
7069
blockifier_stateful_tx_validator: mock_blockifier_validator,
7170
};
7271

7372
let executable_tx = create_executable_invoke_tx(CairoVersion::Cairo1(RunnableCairo1::Casm));
74-
let result = tokio::task::spawn_blocking(move || {
75-
stateful_validator.extract_state_nonce_and_run_validations(
76-
&executable_tx,
77-
mempool_client,
78-
runtime,
79-
)
73+
// get_nonce should fail and return an InternalError.
74+
let err = tokio::task::spawn_blocking(move || {
75+
stateful_validator.get_nonce(executable_tx.contract_address())
8076
})
8177
.await
82-
.unwrap();
78+
.unwrap()
79+
.unwrap_err();
8380
assert_eq!(
84-
result,
85-
Err(StarknetError {
86-
code: StarknetErrorCode::UnknownErrorCode(
87-
"StarknetErrorCode.InternalError".to_string()
88-
),
89-
message: "Internal error".to_string(),
90-
})
81+
err.code,
82+
StarknetErrorCode::UnknownErrorCode("StarknetErrorCode.InternalError".to_string())
9183
);
9284
}
9385

@@ -145,6 +137,7 @@ async fn test_extract_state_nonce_and_run_validations(
145137
let result = tokio::task::spawn_blocking(move || {
146138
stateful_validator.extract_state_nonce_and_run_validations(
147139
&executable_tx,
140+
account_nonce,
148141
mempool_client,
149142
runtime,
150143
)

‎crates/apollo_gateway/src/sync_state_reader.rs‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -131,6 +131,18 @@ impl MempoolStateReader for SyncStateReader {
131131

132132
Ok(block_info)
133133
}
134+
135+
async fn get_account_nonce(&self, contract_address: ContractAddress) -> StateResult<Nonce> {
136+
let res = self.state_sync_client.get_nonce_at(self.block_number, contract_address).await;
137+
138+
match res {
139+
Ok(value) => Ok(value),
140+
Err(StateSyncClientError::StateSyncError(StateSyncError::ContractNotFound(_))) => {
141+
Ok(Nonce::default())
142+
}
143+
Err(e) => Err(StateError::StateReadError(e.to_string())),
144+
}
145+
}
134146
}
135147

136148
impl FetchCompiledClasses for SyncStateReader {

0 commit comments

Comments
 (0)