feat(core): replace callbacks in RpcModule and Methods

core/src/server/rpc_module.rs

@@ -243,6 +243,46 @@ impl Methods {
 		}
 	}
 
+	/// Remove a method by name.
+	pub fn remove(&mut self, name: &'static str) -> Option<MethodCallback> {
+		self.mut_callbacks().remove(name)
+	}
+
+	/// Inserts the method callback for a given name, replacing any existing
+	/// method with the same name.
+	pub fn insert(&mut self, name: &'static str, callback: MethodCallback) -> Option<MethodCallback> {
+		self.mut_callbacks().insert(name, callback)
+	}
+
+	/// Inserts the method callback for a given name. If a method with the same
+	/// name already exists, evaluates the provided `cond` with the method name
+	/// to determine whether to replace the existing method.
+	///
+	/// Returns the previous method callback if it was replaced.
+	///
+	/// Note: If a conflict exists, and `cond` evaluates to false, the method
+	/// will not be replaced and the new method will be dropped.
+	pub fn replace_if(
+		&mut self,
+		name: &'static str,
+		callback: MethodCallback,
+		cond: impl Fn(&'static str) -> bool,
+	) -> Option<MethodCallback> {
+		match self.mut_callbacks().entry(name) {
+			Entry::Occupied(mut entry) => {
+				if cond(name) {
+					Some(entry.insert(callback))
+				} else {
+					None
+				}
+			}
+			Entry::Vacant(vacant) => {
+				vacant.insert(callback);
+				None
+			}
+		}
+	}
+
 	/// Helper for obtaining a mut ref to the callbacks HashMap.
 	fn mut_callbacks(&mut self) -> &mut FxHashMap<&'static str, MethodCallback> {
 		Arc::make_mut(&mut self.callbacks)
@@ -266,6 +306,50 @@ impl Methods {
 		Ok(())
 	}
 
+	/// Merge two [`Methods`]'s by adding all [`MethodCallback`]s from `other`
+	/// into `self`, removing and returning any existing methods with the same
+	/// name.
+	pub fn merge_replace(&mut self, other: impl Into<Methods>) -> Vec<(&'static str, MethodCallback)> {
+		let mut other = other.into();
+		let callbacks = self.mut_callbacks();
+
+		let mut removed = Vec::with_capacity(other.callbacks.len());
+		for (name, callback) in other.mut_callbacks().drain() {
+			if let Some(prev) = callbacks.remove(name) {
+				removed.push((name, prev));
+			}
+			callbacks.insert(name, callback);
+		}
+		removed
+	}
+
+	/// Merge two [`Methods`]'s by adding all [`MethodCallback`]s from `other`
+	/// into `self`. If a method with the same name already exists, evaluates
+	/// the provided `cond` with the method name to determine whether to
+	/// replace the existing method. Uses [`Self::replace_if`]
+	/// internally.
+	///
+	/// Returns a list of removed methods.
+	///
+	/// Note: If a conflict exists, and `cond` evaluates to false, the method
+	/// will not be replaced and the new method will be dropped. Methods
+	/// dropped this way will **not** be included in the returned list.
+	pub fn merge_replace_if(
+		&mut self,
+		other: impl Into<Methods>,
+		cond: impl Fn(&'static str) -> bool,
+	) -> Vec<(&'static str, MethodCallback)> {
+		let mut other = other.into();
+
+		let mut removed = Vec::with_capacity(other.callbacks.len());
+		for (name, callback) in other.mut_callbacks().drain() {
+			if let Some(prev) = self.replace_if(name, callback, &cond) {
+				removed.push((name, prev));
+			}
+		}
+		removed
+	}
+
 	/// Returns the method callback.
 	pub fn method(&self, method_name: &str) -> Option<&MethodCallback> {
 		self.callbacks.get(method_name)
@@ -521,6 +605,16 @@ impl<Context> From<RpcModule<Context>> for Methods {
 }
 
 impl<Context: Send + Sync + 'static> RpcModule<Context> {
+	/// Return a mutable reference to the currently registered methods.
+	pub fn methods_mut(&mut self) -> &mut Methods {
+		&mut self.methods
+	}
+
+	/// Remove a method by name.
+	pub fn remove_method(&mut self, method_name: &'static str) -> Option<MethodCallback> {
+		self.methods.remove(method_name)
+	}
+
 	/// Register a new synchronous RPC method, which computes the response with the given callback.
 	///
 	/// ## Examples
@@ -551,6 +645,20 @@ impl<Context: Send + Sync + 'static> RpcModule<Context> {
 		)
 	}
 
+	/// As [`Self::register_method`], but replaces and returns the method if
+	/// it already exists.
+	pub fn replace_method<R, F>(&mut self, method_name: &'static str, callback: F) -> Option<MethodCallback>
+	where
+		R: IntoResponse + 'static,
+		F: Fn(Params, &Context, &Extensions) -> R + Send + Sync + 'static,
+	{
+		let prev = self.remove_method(method_name);
+		// Errors here can be ignored, as we know the method is not already
+		// registered (we just removed it).
+		let _ = self.register_method(method_name, callback);
+		prev
+	}
+
 	/// Register a new asynchronous RPC method, which computes the response with the given callback.
 	///
 	/// ## Examples
@@ -589,6 +697,25 @@ impl<Context: Send + Sync + 'static> RpcModule<Context> {
 		)
 	}
 
+	/// As [`Self::register_async_method`], but replaces and returns the method
+	/// if it already exists.
+	pub fn replace_async_method<R, Fun, Fut>(
+		&mut self,
+		method_name: &'static str,
+		callback: Fun,
+	) -> Option<MethodCallback>
+	where
+		R: IntoResponse + 'static,
+		Fut: Future<Output = R> + Send,
+		Fun: (Fn(Params<'static>, Arc<Context>, Extensions) -> Fut) + Clone + Send + Sync + 'static,
+	{
+		let prev = self.remove_method(method_name);
+		// Errors here can be ignored, as we know the method is not already
+		// registered (we just removed it).
+		let _ = self.register_async_method(method_name, callback);
+		prev
+	}
+
 	/// Register a new **blocking** synchronous RPC method, which computes the response with the given callback.
 	/// Unlike the regular [`register_method`](RpcModule::register_method), this method can block its thread and perform
 	/// expensive computations.
@@ -632,6 +759,19 @@ impl<Context: Send + Sync + 'static> RpcModule<Context> {
 		Ok(callback)
 	}
 
+	/// As [`Self::register_blocking_method`], but replaces and returns the method if it already exists.
+	pub fn replace_blocking_method<R, F>(&mut self, method_name: &'static str, callback: F) -> Option<MethodCallback>
+	where
+		R: IntoResponse + 'static,
+		F: Fn(Params, Arc<Context>, &Extensions) -> R + Clone + Send + Sync + 'static,
+	{
+		let prev = self.remove_method(method_name);
+		// Errors here can be ignored, as we know the method is not already
+		// registered (we just removed it).
+		let _ = self.register_blocking_method(method_name, callback);
+		prev
+	}
+
 	/// Register a new publish/subscribe interface using JSON-RPC notifications.
 	///
 	/// It implements the [ethereum pubsub specification](https://geth.ethereum.org/docs/rpc/pubsub)
@@ -820,6 +960,34 @@ impl<Context: Send + Sync + 'static> RpcModule<Context> {
 		Ok(callback)
 	}
 
+	/// As [`Self::register_subscription`] but replaces and returns the method
+	/// if it already exists.
+	pub fn replace_subscription<R, F, Fut>(
+		&mut self,
+		subscribe_method_name: &'static str,
+		notif_method_name: &'static str,
+		unsubscribe_method_name: &'static str,
+		callback: F,
+	) -> Option<MethodCallback>
+	where
+		Context: Send + Sync + 'static,
+		F: (Fn(Params<'static>, PendingSubscriptionSink, Arc<Context>, Extensions) -> Fut)
+			+ Send
+			+ Sync
+			+ Clone
+			+ 'static,
+		Fut: Future<Output = R> + Send + 'static,
+		R: IntoSubscriptionCloseResponse + Send,
+	{
+		let prev = self.methods.remove(subscribe_method_name);
+		self.methods.remove(unsubscribe_method_name);
+
+		// Errors here can be ignored, as we know the method is not already
+		// registered (we just removed it).
+		let _ = self.register_subscription(subscribe_method_name, notif_method_name, unsubscribe_method_name, callback);
+		prev
+	}
+
 	/// Similar to [`RpcModule::register_subscription`] but a little lower-level API
 	/// where handling the subscription is managed the user i.e, polling the subscription
 	/// such as spawning a separate task to do so.