chainHead: Add support for storage pagination and cancellation (#14755)

* chainHead/api: Add `chain_head_unstable_continue` method Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io> * chainHead/subscriptions: Register operations for pagination Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io> * chainHead/subscriptions: Merge limits with registered operation Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io> * chainHead/subscriptions: Expose the operation state Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io> * chain_head/storage: Generate WaitingForContinue event Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io> * chainHead: Use the continue operation Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io> * chainHead/tests: Adjust testing to the new storage interface Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io> * chainHead/config: Make pagination limit configurable Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io> * chainHead/tests: Adjust chainHeadConfig Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io> * chainHead/tests: Check pagination and continue method Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io> * chainHead/api: Add `chainHead_unstable_stopOperation` method Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io> * chainHead/subscription: Add shared atomic state for efficient alloc Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io> * chainHead: Implement operation stop Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io> * chainHead/tests: Check that storage ops can be cancelled Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io> * chainHead/storage: Change docs for query_storage_iter_pagination Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io> * chainHead/subscriptions: Fix merge conflicts Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io> * chainHead: Replace `async-channel` with `tokio::sync` Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io> * chainHead/subscription: Add comment about the sender/recv continue Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io> --------- Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io>
2026-07-07 21:17:23 +00:00 · 2023-08-24 14:32:30 +03:00
parent dbb48a5ba2
commit 2e6a2ffa8a
6 changed files with 751 additions and 154 deletions
@@ -17,12 +17,13 @@
 // along with this program. If not, see <https://www.gnu.org/licenses/>.

 use futures::channel::oneshot;
+use parking_lot::Mutex;
 use sc_client_api::Backend;
 use sc_utils::mpsc::{tracing_unbounded, TracingUnboundedReceiver, TracingUnboundedSender};
 use sp_runtime::traits::Block as BlockT;
 use std::{
 	collections::{hash_map::Entry, HashMap},
-	sync::Arc,
+	sync::{atomic::AtomicBool, Arc},
 	time::{Duration, Instant},
 };

@@ -154,12 +155,184 @@ struct PermitOperations {
 	_permit: tokio::sync::OwnedSemaphorePermit,
 }

-impl PermitOperations {
+/// The state of one operation.
+///
+/// This is directly exposed to users via `chain_head_unstable_continue` and
+/// `chain_head_unstable_stop_operation`.
+#[derive(Clone)]
+pub struct OperationState {
+	/// The shared operation state that holds information about the
+	/// `waitingForContinue` event and cancellation.
+	shared_state: Arc<SharedOperationState>,
+	/// Send notifications when the user calls `chainHead_continue` method.
+	send_continue: tokio::sync::mpsc::Sender<()>,
+}
+
+impl OperationState {
+	/// Returns true if `chainHead_continue` is called after the
+	/// `waitingForContinue` event was emitted for the associated
+	/// operation ID.
+	pub fn submit_continue(&self) -> bool {
+		// `waitingForContinue` not generated.
+		if !self.shared_state.requested_continue.load(std::sync::atomic::Ordering::Acquire) {
+			return false
+		}
+
+		// Has enough capacity for 1 message.
+		// Can fail if the `stop_operation` propagated the stop first.
+		self.send_continue.try_send(()).is_ok()
+	}
+
+	/// Stops the operation if `waitingForContinue` event was emitted for the associated
+	/// operation ID.
+	///
+	/// Returns nothing in accordance with `chainHead_unstable_stopOperation`.
+	pub fn stop_operation(&self) {
+		// `waitingForContinue` not generated.
+		if !self.shared_state.requested_continue.load(std::sync::atomic::Ordering::Acquire) {
+			return
+		}
+
+		self.shared_state
+			.operation_stopped
+			.store(true, std::sync::atomic::Ordering::Release);
+
+		// Send might not have enough capacity if `submit_continue` was sent first.
+		// However, the `operation_stopped` boolean was set.
+		let _ = self.send_continue.try_send(());
+	}
+}
+
+/// The shared operation state between the backend [`RegisteredOperation`] and frontend
+/// [`RegisteredOperation`].
+struct SharedOperationState {
+	/// True if the `chainHead` generated `waitingForContinue` event.
+	requested_continue: AtomicBool,
+	/// True if the operation was cancelled by the user.
+	operation_stopped: AtomicBool,
+}
+
+impl SharedOperationState {
+	/// Constructs a new [`SharedOperationState`].
+	///
+	/// This is efficiently cloned under a single heap allocation.
+	fn new() -> Arc<Self> {
+		Arc::new(SharedOperationState {
+			requested_continue: AtomicBool::new(false),
+			operation_stopped: AtomicBool::new(false),
+		})
+	}
+}
+
+/// The registered operation passed to the `chainHead` methods.
+///
+/// This is used internally by the `chainHead` methods.
+pub struct RegisteredOperation {
+	/// The shared operation state that holds information about the
+	/// `waitingForContinue` event and cancellation.
+	shared_state: Arc<SharedOperationState>,
+	/// Receive notifications when the user calls `chainHead_continue` method.
+	recv_continue: tokio::sync::mpsc::Receiver<()>,
+	/// The operation ID of the request.
+	operation_id: String,
+	/// Track the operations ID of this subscription.
+	operations: Arc<Mutex<HashMap<String, OperationState>>>,
+	/// Permit a number of items to be executed by this operation.
+	permit: PermitOperations,
+}
+
+impl RegisteredOperation {
+	/// Wait until the user calls `chainHead_continue` or the operation
+	/// is cancelled via `chainHead_stopOperation`.
+	pub async fn wait_for_continue(&mut self) {
+		self.shared_state
+			.requested_continue
+			.store(true, std::sync::atomic::Ordering::Release);
+
+		// The sender part of this channel is around for as long as this object exists,
+		// because it is stored in the `OperationState` of the `operations` field.
+		// The sender part is removed from tracking when this object is dropped.
+		let _ = self.recv_continue.recv().await;
+
+		self.shared_state
+			.requested_continue
+			.store(false, std::sync::atomic::Ordering::Release);
+	}
+
+	/// Returns true if the current operation was stopped.
+	pub fn was_stopped(&self) -> bool {
+		self.shared_state.operation_stopped.load(std::sync::atomic::Ordering::Acquire)
+	}
+
+	/// Get the operation ID.
+	pub fn operation_id(&self) -> String {
+		self.operation_id.clone()
+	}
+
 	/// Returns the number of reserved elements for this permit.
 	///
 	/// This can be smaller than the number of items requested via [`LimitOperations::reserve()`].
-	fn num_reserved(&self) -> usize {
-		self.num_ops
+	pub fn num_reserved(&self) -> usize {
+		self.permit.num_ops
+	}
+}
+
+impl Drop for RegisteredOperation {
+	fn drop(&mut self) {
+		let mut operations = self.operations.lock();
+		operations.remove(&self.operation_id);
+	}
+}
+
+/// The ongoing operations of a subscription.
+struct Operations {
+	/// The next operation ID to be generated.
+	next_operation_id: usize,
+	/// Limit the number of ongoing operations.
+	limits: LimitOperations,
+	/// Track the operations ID of this subscription.
+	operations: Arc<Mutex<HashMap<String, OperationState>>>,
+}
+
+impl Operations {
+	/// Constructs a new [`Operations`].
+	fn new(max_operations: usize) -> Self {
+		Operations {
+			next_operation_id: 0,
+			limits: LimitOperations::new(max_operations),
+			operations: Default::default(),
+		}
+	}
+
+	/// Register a new operation.
+	pub fn register_operation(&mut self, to_reserve: usize) -> Option<RegisteredOperation> {
+		let permit = self.limits.reserve_at_most(to_reserve)?;
+
+		let operation_id = self.next_operation_id();
+
+		// At most one message can be sent.
+		let (send_continue, recv_continue) = tokio::sync::mpsc::channel(1);
+		let shared_state = SharedOperationState::new();
+
+		let state = OperationState { send_continue, shared_state: shared_state.clone() };
+
+		// Cloned operations for removing the current ID on drop.
+		let operations = self.operations.clone();
+		operations.lock().insert(operation_id.clone(), state);
+
+		Some(RegisteredOperation { shared_state, operation_id, recv_continue, operations, permit })
+	}
+
+	/// Get the associated operation state with the ID.
+	pub fn get_operation(&self, id: &str) -> Option<OperationState> {
+		self.operations.lock().get(id).map(|state| state.clone())
+	}
+
+	/// Generate the next operation ID for this subscription.
+	fn next_operation_id(&mut self) -> String {
+		let op_id = self.next_operation_id;
+		self.next_operation_id += 1;
+		op_id.to_string()
 	}
 }

@@ -180,10 +353,8 @@ struct SubscriptionState<Block: BlockT> {
 	///
 	/// This object is cloned between methods.
 	response_sender: TracingUnboundedSender<FollowEvent<Block::Hash>>,
-	/// Limit the number of ongoing operations.
-	limits: LimitOperations,
-	/// The next operation ID.
-	next_operation_id: usize,
+	/// The ongoing operations of a subscription.
+	operations: Operations,
 	/// Track the block hashes available for this subscription.
 	///
 	/// This implementation assumes:
@@ -296,18 +467,16 @@ impl<Block: BlockT> SubscriptionState<Block> {
 		timestamp
 	}

-	/// Generate the next operation ID for this subscription.
-	fn next_operation_id(&mut self) -> usize {
-		let op_id = self.next_operation_id;
-		self.next_operation_id = self.next_operation_id.wrapping_add(1);
-		op_id
+	/// Register a new operation.
+	///
+	/// The registered operation can execute at least one item and at most the requested items.
+	fn register_operation(&mut self, to_reserve: usize) -> Option<RegisteredOperation> {
+		self.operations.register_operation(to_reserve)
 	}

-	/// Reserves capacity to execute at least one operation and at most the requested items.
-	///
-	/// For more details see [`PermitOperations`].
-	fn reserve_at_most(&self, to_reserve: usize) -> Option<PermitOperations> {
-		self.limits.reserve_at_most(to_reserve)
+	/// Get the associated operation state with the ID.
+	pub fn get_operation(&self, id: &str) -> Option<OperationState> {
+		self.operations.get_operation(id)
 	}
 }

@@ -318,8 +487,7 @@ pub struct BlockGuard<Block: BlockT, BE: Backend<Block>> {
 	hash: Block::Hash,
 	with_runtime: bool,
 	response_sender: TracingUnboundedSender<FollowEvent<Block::Hash>>,
-	operation_id: String,
-	permit_operations: PermitOperations,
+	operation: RegisteredOperation,
 	backend: Arc<BE>,
 }

@@ -337,22 +505,14 @@ impl<Block: BlockT, BE: Backend<Block>> BlockGuard<Block, BE> {
 		hash: Block::Hash,
 		with_runtime: bool,
 		response_sender: TracingUnboundedSender<FollowEvent<Block::Hash>>,
-		operation_id: usize,
-		permit_operations: PermitOperations,
+		operation: RegisteredOperation,
 		backend: Arc<BE>,
 	) -> Result<Self, SubscriptionManagementError> {
 		backend
 			.pin_block(hash)
 			.map_err(|err| SubscriptionManagementError::Custom(err.to_string()))?;

-		Ok(Self {
-			hash,
-			with_runtime,
-			response_sender,
-			operation_id: operation_id.to_string(),
-			permit_operations,
-			backend,
-		})
+		Ok(Self { hash, with_runtime, response_sender, operation, backend })
 	}

 	/// The `with_runtime` flag of the subscription.
@@ -365,16 +525,9 @@ impl<Block: BlockT, BE: Backend<Block>> BlockGuard<Block, BE> {
 		self.response_sender.clone()
 	}

-	/// The operation ID of this method.
-	pub fn operation_id(&self) -> String {
-		self.operation_id.clone()
-	}
-
-	/// Returns the number of reserved elements for this permit.
-	///
-	/// This can be smaller than the number of items requested.
-	pub fn num_reserved(&self) -> usize {
-		self.permit_operations.num_reserved()
+	/// Get the details of the registered operation.
+	pub fn operation(&mut self) -> &mut RegisteredOperation {
+		&mut self.operation
 	}
 }

@@ -445,9 +598,8 @@ impl<Block: BlockT, BE: Backend<Block>> SubscriptionsInner<Block, BE> {
 				with_runtime,
 				tx_stop: Some(tx_stop),
 				response_sender,
-				limits: LimitOperations::new(self.max_ongoing_operations),
-				next_operation_id: 0,
 				blocks: Default::default(),
+				operations: Operations::new(self.max_ongoing_operations),
 			};
 			entry.insert(state);

@@ -631,21 +783,24 @@ impl<Block: BlockT, BE: Backend<Block>> SubscriptionsInner<Block, BE> {
 			return Err(SubscriptionManagementError::BlockHashAbsent)
 		}

-		let Some(permit_operations) = sub.reserve_at_most(to_reserve) else {
+		let Some(operation) = sub.register_operation(to_reserve) else {
 			// Error when the server cannot execute at least one operation.
 			return Err(SubscriptionManagementError::ExceededLimits)
 		};

-		let operation_id = sub.next_operation_id();
 		BlockGuard::new(
 			hash,
 			sub.with_runtime,
 			sub.response_sender.clone(),
-			operation_id,
-			permit_operations,
+			operation,
 			self.backend.clone(),
 		)
 	}
+
+	pub fn get_operation(&mut self, sub_id: &str, id: &str) -> Option<OperationState> {
+		let state = self.subs.get(sub_id)?;
+		state.get_operation(id)
+	}
 }

 #[cfg(test)]
@@ -758,8 +913,7 @@ mod tests {
 			with_runtime: false,
 			tx_stop: None,
 			response_sender,
-			next_operation_id: 0,
-			limits: LimitOperations::new(MAX_OPERATIONS_PER_SUB),
+			operations: Operations::new(MAX_OPERATIONS_PER_SUB),
 			blocks: Default::default(),
 		};

@@ -788,9 +942,8 @@ mod tests {
 			with_runtime: false,
 			tx_stop: None,
 			response_sender,
-			next_operation_id: 0,
-			limits: LimitOperations::new(MAX_OPERATIONS_PER_SUB),
 			blocks: Default::default(),
+			operations: Operations::new(MAX_OPERATIONS_PER_SUB),
 		};

 		let hash = H256::random();
@@ -1107,12 +1260,12 @@ mod tests {

 		// One operation is reserved.
 		let permit_one = ops.reserve_at_most(1).unwrap();
-		assert_eq!(permit_one.num_reserved(), 1);
+		assert_eq!(permit_one.num_ops, 1);

 		// Request 2 operations, however there is capacity only for one.
 		let permit_two = ops.reserve_at_most(2).unwrap();
 		// Number of reserved permits is smaller than provided.
-		assert_eq!(permit_two.num_reserved(), 1);
+		assert_eq!(permit_two.num_ops, 1);

 		// Try to reserve operations when there's no space.
 		let permit = ops.reserve_at_most(1);
@@ -1123,6 +1276,6 @@ mod tests {

 		// Can reserve again
 		let permit_three = ops.reserve_at_most(1).unwrap();
-		assert_eq!(permit_three.num_reserved(), 1);
+		assert_eq!(permit_three.num_ops, 1);
 	}
 }
@@ -25,6 +25,8 @@ mod error;
 mod inner;

 use self::inner::SubscriptionsInner;
+
+pub use self::inner::OperationState;
 pub use error::SubscriptionManagementError;
 pub use inner::{BlockGuard, InsertedSubscriptionData};

@@ -126,4 +128,10 @@ impl<Block: BlockT, BE: Backend<Block>> SubscriptionManagement<Block, BE> {
 		let mut inner = self.inner.write();
 		inner.lock_block(sub_id, hash, to_reserve)
 	}
+
+	/// Get the operation state.
+	pub fn get_operation(&self, sub_id: &str, operation_id: &str) -> Option<OperationState> {
+		let mut inner = self.inner.write();
+		inner.get_operation(sub_id, operation_id)
+	}
 }