From e7b6d7dffd6459174f02598bd8b84fe4b1cb6e72 Mon Sep 17 00:00:00 2001
From: Liam Aharon <liam.aharon@hotmail.com>
Date: Mon, 20 May 2024 03:53:12 +1000
Subject: [PATCH] `remote-externalities`: `rpc_child_get_keys` to use paged
 scraping (#4512)

Replace usage of deprecated
`substrate_rpc_client::ChildStateApi::storage_keys` with
`substrate_rpc_client::ChildStateApi::storage_keys_paged`.

Required for successful scraping of Aleph Zero state.
---
 .../frame/remote-externalities/src/lib.rs     | 53 +++++++++++++------
 1 file changed, 37 insertions(+), 16 deletions(-)

diff --git a/substrate/utils/frame/remote-externalities/src/lib.rs b/substrate/utils/frame/remote-externalities/src/lib.rs
index 201b5e176f34..0ecb98f31343 100644
--- a/substrate/utils/frame/remote-externalities/src/lib.rs
+++ b/substrate/utils/frame/remote-externalities/src/lib.rs
@@ -834,30 +834,51 @@ where
 	) -> Result<Vec<StorageKey>, &'static str> {
 		let retry_strategy =
 			FixedInterval::new(Self::KEYS_PAGE_RETRY_INTERVAL).take(Self::MAX_RETRIES);
-		let get_child_keys_closure = || {
-			#[allow(deprecated)]
-			substrate_rpc_client::ChildStateApi::storage_keys(
-				client,
-				PrefixedStorageKey::new(prefixed_top_key.as_ref().to_vec()),
-				child_prefix.clone(),
-				Some(at),
-			)
-		};
-		let child_keys =
-			Retry::spawn(retry_strategy, get_child_keys_closure).await.map_err(|e| {
-				error!(target: LOG_TARGET, "Error = {:?}", e);
-				"rpc child_get_keys failed."
-			})?;
+		let mut all_child_keys = Vec::new();
+		let mut start_key = None;
+
+		loop {
+			let get_child_keys_closure = || {
+				let top_key = PrefixedStorageKey::new(prefixed_top_key.0.clone());
+				substrate_rpc_client::ChildStateApi::storage_keys_paged(
+					client,
+					top_key,
+					Some(child_prefix.clone()),
+					Self::DEFAULT_KEY_DOWNLOAD_PAGE,
+					start_key.clone(),
+					Some(at),
+				)
+			};
+
+			let child_keys = Retry::spawn(retry_strategy.clone(), get_child_keys_closure)
+				.await
+				.map_err(|e| {
+					error!(target: LOG_TARGET, "Error = {:?}", e);
+					"rpc child_get_keys failed."
+				})?;
+
+			let keys_count = child_keys.len();
+			if keys_count == 0 {
+				break;
+			}
+
+			start_key = child_keys.last().cloned();
+			all_child_keys.extend(child_keys);
+
+			if keys_count < Self::DEFAULT_KEY_DOWNLOAD_PAGE as usize {
+				break;
+			}
+		}
 
 		debug!(
 			target: LOG_TARGET,
 			"[thread = {:?}] scraped {} child-keys of the child-bearing top key: {}",
 			std::thread::current().id(),
-			child_keys.len(),
+			all_child_keys.len(),
 			HexDisplay::from(prefixed_top_key)
 		);
 
-		Ok(child_keys)
+		Ok(all_child_keys)
 	}
 }