Fix the storage_size/state_getStorageSize RPC call (#13154)

* Have `KeyIterator` clone the `prefix` it receives

* Stream keys in `storage_size` RPC and add a runtime limit

* Update client/rpc/Cargo.toml

Co-authored-by: Bastian Köcher <git@kchr.de>

* Update client/rpc/src/state/utils.rs

Co-authored-by: Bastian Köcher <git@kchr.de>

* Rename the types to signify that the cancellation is due to a timeout

* Move the test into a `mod tests`

* Add a comment regarding `biased` in `tokio::select`

* Make the `clone` explicit when calling `KeyIterator::{new, new_child}`

Co-authored-by: Bastian Köcher <git@kchr.de>

substrate/client/rpc/src/state/utils.rs

@@ -0,0 +1,140 @@
+// This file is part of Substrate.
+
+// Copyright (C) 2017-2023 Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with this program. If not, see <https://www.gnu.org/licenses/>.
+
+use std::{
+	sync::{
+		atomic::{AtomicBool, Ordering},
+		Arc,
+	},
+	time::Duration,
+};
+
+/// An error signifying that a task has been cancelled due to a timeout.
+#[derive(Debug)]
+pub struct Timeout;
+
+impl std::error::Error for Timeout {}
+impl std::fmt::Display for Timeout {
+	fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
+		fmt.write_str("task has been running too long")
+	}
+}
+
+/// A handle which can be used to check whether the task has been cancelled due to a timeout.
+#[repr(transparent)]
+pub struct IsTimedOut(Arc<AtomicBool>);
+
+impl IsTimedOut {
+	#[must_use]
+	pub fn check_if_timed_out(&self) -> std::result::Result<(), Timeout> {
+		if self.0.load(Ordering::Relaxed) {
+			Err(Timeout)
+		} else {
+			Ok(())
+		}
+	}
+}
+
+/// An error for a task which either panicked, or has been cancelled due to a timeout.
+#[derive(Debug)]
+pub enum SpawnWithTimeoutError {
+	JoinError(tokio::task::JoinError),
+	Timeout,
+}
+
+impl std::error::Error for SpawnWithTimeoutError {}
+impl std::fmt::Display for SpawnWithTimeoutError {
+	fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
+		match self {
+			SpawnWithTimeoutError::JoinError(error) => error.fmt(fmt),
+			SpawnWithTimeoutError::Timeout => Timeout.fmt(fmt),
+		}
+	}
+}
+
+struct CancelOnDrop(Arc<AtomicBool>);
+impl Drop for CancelOnDrop {
+	fn drop(&mut self) {
+		self.0.store(true, Ordering::Relaxed)
+	}
+}
+
+/// Spawns a new blocking task with a given `timeout`.
+///
+/// The `callback` should continuously call [`IsTimedOut::check_if_timed_out`],
+/// which will return an error once the task runs for longer than `timeout`.
+///
+/// If `timeout` is `None` then this works just as a regular `spawn_blocking`.
+pub async fn spawn_blocking_with_timeout<R>(
+	timeout: Option<Duration>,
+	callback: impl FnOnce(IsTimedOut) -> std::result::Result<R, Timeout> + Send + 'static,
+) -> Result<R, SpawnWithTimeoutError>
+where
+	R: Send + 'static,
+{
+	let is_timed_out_arc = Arc::new(AtomicBool::new(false));
+	let is_timed_out = IsTimedOut(is_timed_out_arc.clone());
+	let _cancel_on_drop = CancelOnDrop(is_timed_out_arc);
+	let task = tokio::task::spawn_blocking(move || callback(is_timed_out));
+
+	let result = if let Some(timeout) = timeout {
+		tokio::select! {
+			// Shouldn't really matter, but make sure the task is polled before the timeout,
+			// in case the task finishes after the timeout and the timeout is really short.
+			biased;
+
+			task_result = task => task_result,
+			_ = tokio::time::sleep(timeout) => Ok(Err(Timeout))
+		}
+	} else {
+		task.await
+	};
+
+	match result {
+		Ok(Ok(result)) => Ok(result),
+		Ok(Err(Timeout)) => Err(SpawnWithTimeoutError::Timeout),
+		Err(error) => Err(SpawnWithTimeoutError::JoinError(error)),
+	}
+}
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+
+	#[tokio::test]
+	async fn spawn_blocking_with_timeout_works() {
+		let task: Result<(), SpawnWithTimeoutError> =
+			spawn_blocking_with_timeout(Some(Duration::from_millis(100)), |is_timed_out| {
+				std::thread::sleep(Duration::from_millis(200));
+				is_timed_out.check_if_timed_out()?;
+				unreachable!();
+			})
+			.await;
+
+		assert_matches::assert_matches!(task, Err(SpawnWithTimeoutError::Timeout));
+
+		let task = spawn_blocking_with_timeout(Some(Duration::from_millis(100)), |is_timed_out| {
+			std::thread::sleep(Duration::from_millis(20));
+			is_timed_out.check_if_timed_out()?;
+			Ok(())
+		})
+		.await;
+
+		assert_matches::assert_matches!(task, Ok(()));
+	}
+}