492 update block timestamp

crates/katana-core/src/block_context.rs

@@ -60,3 +60,9 @@ pub fn block_context_from_config(config: &StarknetConfig) -> BlockContext {
         invoke_tx_max_n_steps: 1_000_000,
     }
 }
+
+#[derive(Default)]
+pub struct BlockContextGenerator {
+    pub block_timestamp_offset: i64,
+    pub next_block_start_time: u64,
+}

crates/katana-core/src/sequencer.rs

@@ -10,7 +10,7 @@ use blockifier::transaction::account_transaction::AccountTransaction;
 use blockifier::transaction::transaction_execution::Transaction;
 use blockifier::transaction::transactions::ExecutableTransaction;
 use starknet::core::types::{BlockId, BlockTag, FeeEstimate, StateUpdate, TransactionStatus};
-use starknet_api::block::{BlockHash, BlockNumber};
+use starknet_api::block::{BlockHash, BlockNumber, BlockTimestamp};
 use starknet_api::core::{calculate_contract_address, ChainId, ClassHash, ContractAddress, Nonce};
 use starknet_api::hash::StarkFelt;
 use starknet_api::stark_felt;
@@ -321,6 +321,10 @@ impl Sequencer for KatanaSequencer {
         self.starknet.read().await.block_context.block_number
     }
 
+    async fn next_block_timestamp(&self) -> BlockTimestamp {
+        self.starknet.read().await.block_context.block_timestamp
+    }
+
     async fn block(&self, block_id: BlockId) -> Option<StarknetBlock> {
         match block_id {
             BlockId::Tag(BlockTag::Pending) => {
@@ -512,6 +516,14 @@ impl Sequencer for KatanaSequencer {
         self.starknet.write().await.generate_pending_block();
     }
 
+    async fn set_next_block_timestamp(&self, timestamp: u64) -> SequencerResult<()> {
+        self.starknet.write().await.set_next_block_timestamp(timestamp)
+    }
+
+    async fn increase_next_block_timestamp(&self, timestamp: u64) -> SequencerResult<()> {
+        self.starknet.write().await.increase_next_block_timestamp(timestamp)
+    }
+
     async fn predeployed_accounts(&self) -> Vec<Account> {
         self.starknet.read().await.predeployed_accounts.accounts.clone()
     }
@@ -540,6 +552,8 @@ pub trait Sequencer {
 
     async fn block_number(&self) -> BlockNumber;
 
+    async fn next_block_timestamp(&self) -> BlockTimestamp;
+
     async fn block(&self, block_id: BlockId) -> Option<StarknetBlock>;
 
     async fn transaction(
@@ -601,4 +615,8 @@ pub trait Sequencer {
     ) -> SequencerResult<Vec<EmittedEvent>>;
 
     async fn state_update(&self, block_id: BlockId) -> SequencerResult<StateUpdate>;
+
+    async fn set_next_block_timestamp(&self, timestamp: u64) -> SequencerResult<()>;
+
+    async fn increase_next_block_timestamp(&self, timestamp: u64) -> SequencerResult<()>;
 }

crates/katana-core/src/sequencer_error.rs

@@ -28,4 +28,6 @@ pub enum SequencerError {
     StarknetApi(#[from] StarknetApiError),
     #[error(transparent)]
     EntryPointExecution(#[from] EntryPointExecutionError),
+    #[error("Wait for pending transactions")]
+    PendingTransactions,
 }

crates/katana-core/src/starknet/mod.rs

@@ -28,11 +28,12 @@ use transaction::{StarknetTransaction, StarknetTransactions};
 
 use self::transaction::ExternalFunctionCall;
 use crate::accounts::PredeployedAccounts;
-use crate::block_context::block_context_from_config;
+use crate::block_context::{block_context_from_config, BlockContextGenerator};
 use crate::constants::{
     DEFAULT_PREFUNDED_ACCOUNT_BALANCE, ERC20_CONTRACT_CLASS_HASH, FEE_TOKEN_ADDRESS, UDC_ADDRESS,
     UDC_CLASS_HASH,
 };
+use crate::sequencer_error::SequencerError;
 use crate::state::DictStateReader;
 use crate::util::{
     convert_blockifier_tx_to_starknet_api_tx, convert_state_diff_to_rpc_state_diff,
@@ -55,6 +56,7 @@ pub struct StarknetWrapper {
     pub config: StarknetConfig,
     pub blocks: StarknetBlocks,
     pub block_context: BlockContext,
+    pub block_context_generator: BlockContextGenerator,
     pub transactions: StarknetTransactions,
     pub state: DictStateReader,
     pub predeployed_accounts: PredeployedAccounts,
@@ -65,6 +67,7 @@ impl StarknetWrapper {
     pub fn new(config: StarknetConfig) -> Self {
         let blocks = StarknetBlocks::default();
         let block_context = block_context_from_config(&config);
+        let block_context_generator = BlockContextGenerator::default();
         let transactions = StarknetTransactions::default();
         let mut state = DictStateReader::default();
         let pending_state = CachedState::new(state.clone());
@@ -84,6 +87,7 @@ impl StarknetWrapper {
             blocks,
             transactions,
             block_context,
+            block_context_generator,
             pending_state,
             predeployed_accounts,
         }
@@ -349,7 +353,20 @@ impl StarknetWrapper {
 
     fn update_block_context(&mut self) {
         self.block_context.block_number = self.block_context.block_number.next();
-        self.block_context.block_timestamp = BlockTimestamp(get_current_timestamp().as_secs());
+
+        let current_timestamp_secs = get_current_timestamp().as_secs() as i64;
+
+        if self.block_context_generator.next_block_start_time == 0 {
+            let block_timestamp =
+                current_timestamp_secs + self.block_context_generator.block_timestamp_offset;
+            self.block_context.block_timestamp = BlockTimestamp(block_timestamp as u64);
+        } else {
+            let block_timestamp = self.block_context_generator.next_block_start_time;
+            self.block_context_generator.block_timestamp_offset =
+                block_timestamp as i64 - current_timestamp_secs;
+            self.block_context.block_timestamp = BlockTimestamp(block_timestamp);
+            self.block_context_generator.next_block_start_time = 0;
+        }
     }
 
     // apply the pending state diff to the state
@@ -358,6 +375,29 @@ impl StarknetWrapper {
         apply_new_state(state, &mut self.pending_state);
         self.blocks.store_state(self.block_context.block_number, state.clone());
     }
+
+    pub fn set_next_block_timestamp(&mut self, timestamp: u64) -> Result<(), SequencerError> {
+        if has_pending_transactions(self) {
+            return Err(SequencerError::PendingTransactions);
+        }
+        self.block_context_generator.next_block_start_time = timestamp;
+        Ok(())
+    }
+
+    pub fn increase_next_block_timestamp(&mut self, timestamp: u64) -> Result<(), SequencerError> {
+        if has_pending_transactions(self) {
+            return Err(SequencerError::PendingTransactions);
+        }
+        self.block_context_generator.block_timestamp_offset += timestamp as i64;
+        Ok(())
+    }
+}
+
+fn has_pending_transactions(starknet: &StarknetWrapper) -> bool {
+    match starknet.blocks.pending_block {
+        Some(ref pending_block) => !pending_block.inner.body.transactions.is_empty(),
+        None => false,
+    }
 }
 
 fn apply_new_state(old_state: &mut DictStateReader, new_state: &mut CachedState<DictStateReader>) {

crates/katana-core/tests/starknet.rs

@@ -33,6 +33,48 @@ fn create_test_starknet() -> StarknetWrapper {
     starknet
 }
 
+#[test]
+fn test_next_block_timestamp_in_past() {
+    let mut starknet = create_test_starknet();
+    starknet.generate_pending_block();
+
+    let timestamp = starknet.block_context.block_timestamp;
+    starknet.set_next_block_timestamp(timestamp.0 - 1000).unwrap();
+
+    starknet.generate_pending_block();
+    let new_timestamp = starknet.block_context.block_timestamp;
+
+    assert_eq!(new_timestamp.0, timestamp.0 - 1000, "timestamp should be updated");
+}
+
+#[test]
+fn test_set_next_block_timestamp_in_future() {
+    let mut starknet = create_test_starknet();
+    starknet.generate_pending_block();
+
+    let timestamp = starknet.block_context.block_timestamp;
+    starknet.set_next_block_timestamp(timestamp.0 + 1000).unwrap();
+
+    starknet.generate_pending_block();
+    let new_timestamp = starknet.block_context.block_timestamp;
+
+    assert_eq!(new_timestamp.0, timestamp.0 + 1000, "timestamp should be updated");
+}
+
+#[test]
+fn test_increase_next_block_timestamp() {
+    let mut starknet = create_test_starknet();
+    starknet.generate_pending_block();
+
+    let timestamp = starknet.block_context.block_timestamp;
+    starknet.increase_next_block_timestamp(1000).unwrap();
+
+    starknet.generate_pending_block();
+    let new_timestamp = starknet.block_context.block_timestamp;
+
+    assert_eq!(new_timestamp.0, timestamp.0 + 1000, "timestamp should be updated");
+}
+
 #[test]
 fn test_creating_blocks() {
     let mut starknet = create_test_starknet();

crates/katana-rpc/src/katana/api.rs

@@ -5,7 +5,11 @@ use jsonrpsee::types::ErrorObject;
 use katana_core::accounts::Account;
 
 #[derive(thiserror::Error, Clone, Copy, Debug)]
-pub enum KatanaApiError {}
+#[allow(clippy::enum_variant_names)]
+pub enum KatanaApiError {
+    #[error("Failed to change next block timestamp")]
+    FailedToChangeNextBlockTimestamp = 1,
+}
 
 impl From<KatanaApiError> for Error {
     fn from(err: KatanaApiError) -> Self {
@@ -18,6 +22,15 @@ pub trait KatanaApi {
     #[method(name = "generateBlock")]
     async fn generate_block(&self) -> Result<(), Error>;
 
+    #[method(name = "nextBlockTimestamp")]
+    async fn next_block_timestamp(&self) -> Result<u64, Error>;
+
+    #[method(name = "setNextBlockTimestamp")]
+    async fn set_next_block_timestamp(&self, timestamp: u64) -> Result<(), Error>;
+
+    #[method(name = "increaseNextBlockTimestamp")]
+    async fn increase_next_block_timestamp(&self, timestamp: u64) -> Result<(), Error>;
+
     #[method(name = "predeployedAccounts")]
     async fn predeployed_accounts(&self) -> Result<Vec<Account>, Error>;
 }

crates/katana-rpc/src/katana/mod.rs

@@ -4,7 +4,7 @@ use jsonrpsee::core::{async_trait, Error};
 use katana_core::accounts::Account;
 use katana_core::sequencer::Sequencer;
 
-use self::api::KatanaApiServer;
+use self::api::{KatanaApiError, KatanaApiServer};
 
 pub mod api;
 
@@ -25,6 +25,24 @@ impl<S: Sequencer + Send + Sync + 'static> KatanaApiServer for KatanaRpc<S> {
         Ok(())
     }
 
+    async fn next_block_timestamp(&self) -> Result<u64, Error> {
+        Ok(self.sequencer.next_block_timestamp().await.0)
+    }
+
+    async fn set_next_block_timestamp(&self, timestamp: u64) -> Result<(), Error> {
+        self.sequencer
+            .set_next_block_timestamp(timestamp)
+            .await
+            .map_err(|_| Error::from(KatanaApiError::FailedToChangeNextBlockTimestamp))
+    }
+
+    async fn increase_next_block_timestamp(&self, timestamp: u64) -> Result<(), Error> {
+        self.sequencer
+            .increase_next_block_timestamp(timestamp)
+            .await
+            .map_err(|_| Error::from(KatanaApiError::FailedToChangeNextBlockTimestamp))
+    }
+
     async fn predeployed_accounts(&self) -> Result<Vec<Account>, Error> {
         let accounts: Vec<Account> = self.sequencer.predeployed_accounts().await;
         Ok(accounts)