diff --git a/examples/examples/get_storage.rs b/examples/examples/get_storage.rs index e6db7d7ae..ddcac90f6 100644 --- a/examples/examples/get_storage.rs +++ b/examples/examples/get_storage.rs @@ -76,7 +76,7 @@ async fn main() { let storage_key_prefix = api.get_storage_map_key_prefix("System", "Account").unwrap(); let max_keys = 3; let storage_keys = api - .get_storage_keys_paged(Some(storage_key_prefix), max_keys, None, None) + .get_storage_keys_paged_limited(Some(storage_key_prefix), max_keys, None, None) .unwrap(); assert_eq!(storage_keys.len() as u32, max_keys); // Get the storage values that belong to the retrieved storage keys. @@ -91,7 +91,7 @@ async fn main() { let storage_double_map_key_prefix = api.get_storage_double_map_key_prefix("Staking", "ErasStakers", 0).unwrap(); let double_map_storage_keys = api - .get_storage_keys_paged(Some(storage_double_map_key_prefix), max_keys, None, None) + .get_storage_keys_paged_limited(Some(storage_double_map_key_prefix), max_keys, None, None) .unwrap(); // Get the storage values that belong to the retrieved storage keys. diff --git a/src/api/rpc_api/state.rs b/src/api/rpc_api/state.rs index 6101bb30b..46ccc5cd5 100644 --- a/src/api/rpc_api/state.rs +++ b/src/api/rpc_api/state.rs @@ -20,6 +20,7 @@ use ac_node_api::MetadataError; use ac_primitives::config::Config; use alloc::{borrow::ToOwned, string::String, vec, vec::Vec}; use codec::{Decode, Encode}; +use core::cmp; use log::*; use serde::de::DeserializeOwned; use sp_storage::{StorageChangeSet, StorageData, StorageKey}; @@ -99,7 +100,7 @@ pub trait GetStorage { /// If `start_key` is passed, return next keys in storage in lexicographic order. /// /// `at_block`: the state is queried at this block, set to `None` to get the state from the latest known block. - async fn get_storage_keys_paged( + async fn get_storage_keys_paged_limited( &self, prefix: Option, count: u32, @@ -113,7 +114,7 @@ pub trait GetStorage { /// If `start_key` is passed, return next keys in storage in lexicographic order. /// /// `at_block`: the state is queried at this block, set to `None` to get the state from the latest known block. - async fn get_all_storage_keys_paged_up_to_count( + async fn get_storage_keys_paged( &self, prefix: Option, count: u32, @@ -267,7 +268,7 @@ where } } - async fn get_storage_keys_paged( + async fn get_storage_keys_paged_limited( &self, storage_key_prefix: Option, count: u32, @@ -284,7 +285,7 @@ where Ok(storage) } - async fn get_all_storage_keys_paged_up_to_count( + async fn get_storage_keys_paged( &self, storage_key_prefix: Option, count: u32, @@ -292,22 +293,26 @@ where at_block: Option, ) -> Result> { let mut storage_keys: Vec = Vec::new(); - let mut still_todo = count; - let mut new_key = start_key; - - while still_todo > 0 { - let new_count = if still_todo < STORAGE_KEYS_PAGED_MAX_COUNT { - still_todo - } else { - STORAGE_KEYS_PAGED_MAX_COUNT - }; + let mut keys_left_to_fetch = count; + let mut new_start_key = start_key; + while keys_left_to_fetch > 0 { + let new_count = cmp::min(STORAGE_KEYS_PAGED_MAX_COUNT, keys_left_to_fetch); let mut keys = self - .get_storage_keys_paged(storage_key_prefix.clone(), new_count, new_key, at_block) + .get_storage_keys_paged_limited( + storage_key_prefix.clone(), + new_count, + new_start_key, + at_block, + ) .await?; + let num_keys = keys.len() as u32; storage_keys.append(&mut keys); - still_todo -= new_count; - new_key = keys.last().map(|x| x.to_owned()); + if num_keys < new_count { + break + } + keys_left_to_fetch -= new_count; + new_start_key = keys.last().map(|x| x.to_owned()); } Ok(storage_keys) diff --git a/testing/examples/state_tests.rs b/testing/examples/state_tests.rs index 9a1fd64cd..5359c5f39 100644 --- a/testing/examples/state_tests.rs +++ b/testing/examples/state_tests.rs @@ -90,22 +90,25 @@ async fn main() { let _constants: Balance = api.get_constant("Balances", "ExistentialDeposit").unwrap(); let max_keys = 2003; - let result = - api.get_storage_keys_paged(Some(storage_key_prefix.clone()), max_keys.clone(), None, None); + let result = api.get_storage_keys_paged_limited( + Some(storage_key_prefix.clone()), + max_keys.clone(), + None, + None, + ); assert!(result.is_err()); assert!(format!("{result:?}").contains("count exceeds maximum value")); let storage_keys = api - .get_all_storage_keys_paged_up_to_count(Some(storage_key_prefix), max_keys, None, None) + .get_storage_keys_paged(Some(storage_key_prefix), max_keys, None, None) .unwrap(); - assert!(storage_keys.len() as u32 > 3); - assert!(storage_keys.len() as u32 <= max_keys); + assert_eq!(storage_keys.len() as u32, 13); let max_keys = 20; - let storage_keys = api.get_storage_keys_paged(None, max_keys.clone(), None, None).unwrap(); + let storage_keys = + api.get_storage_keys_paged_limited(None, max_keys.clone(), None, None).unwrap(); assert_eq!(storage_keys.len() as u32, max_keys); - let storage_keys = - api.get_all_storage_keys_paged_up_to_count(None, max_keys, None, None).unwrap(); + let storage_keys = api.get_storage_keys_paged(None, max_keys, None, None).unwrap(); assert_eq!(storage_keys.len() as u32, max_keys); }