Add OnMessageAccepted callback (#1134)

* Add `OnMessageAccepted` config type

* Update actual weight

* Add unit test

* Update weight

* Remove old comment & update wrong test data

* Make ci happy

* Add lane_id param

* update test case

* Make log info more readable

* Use saturating_sub

* Update docs

bridges/bin/millau/runtime/src/lib.rs

@@ -394,6 +394,7 @@ impl pallet_bridge_messages::Config<WithRialtoMessagesInstance> for Runtime {
 		GetDeliveryConfirmationTransactionFee,
 		RootAccountForPayments,
 	>;
+	type OnMessageAccepted = ();
 	type OnDeliveryConfirmed = pallet_bridge_token_swap::Pallet<Runtime, WithRialtoTokenSwapInstance>;
 
 	type SourceHeaderChain = crate::rialto_messages::Rialto;

bridges/bin/rialto/runtime/src/lib.rs

@@ -528,6 +528,7 @@ impl pallet_bridge_messages::Config<WithMillauMessagesInstance> for Runtime {
 		GetDeliveryConfirmationTransactionFee,
 		RootAccountForPayments,
 	>;
+	type OnMessageAccepted = ();
 	type OnDeliveryConfirmed = ();
 
 	type SourceHeaderChain = crate::millau_messages::Millau;

bridges/modules/messages/src/lib.rs

@@ -48,7 +48,8 @@ use crate::weights::WeightInfo;
 
 use bp_messages::{
 	source_chain::{
-		LaneMessageVerifier, MessageDeliveryAndDispatchPayment, OnDeliveryConfirmed, RelayersRewards, TargetHeaderChain,
+		LaneMessageVerifier, MessageDeliveryAndDispatchPayment, OnDeliveryConfirmed, OnMessageAccepted,
+		RelayersRewards, TargetHeaderChain,
 	},
 	target_chain::{DispatchMessage, MessageDispatch, ProvedLaneMessages, ProvedMessages, SourceHeaderChain},
 	total_unrewarded_messages, DeliveredMessages, InboundLaneData, LaneId, MessageData, MessageKey, MessageNonce,
@@ -160,6 +161,8 @@ pub mod pallet {
 			Self::AccountId,
 			Self::OutboundMessageFee,
 		>;
+		/// Handler for accepted messages.
+		type OnMessageAccepted: OnMessageAccepted;
 		/// Handler for delivered messages.
 		type OnDeliveryConfirmed: OnDeliveryConfirmed;
 
@@ -238,7 +241,7 @@ pub mod pallet {
 		}
 
 		/// Send message over lane.
-		#[pallet::weight(T::WeightInfo::send_message_weight(payload))]
+		#[pallet::weight(T::WeightInfo::send_message_weight(payload, T::DbWeight::get()))]
 		pub fn send_message(
 			origin: OriginFor<T>,
 			lane_id: LaneId,
@@ -557,19 +560,19 @@ pub mod pallet {
 					Some(difference) => {
 						log::trace!(
 							target: "runtime::bridge-messages",
-							"Messages delivery callback has returned unspent weight to refund the submitter: \
+							"T::OnDeliveryConfirmed callback has spent less weight than expected. Refunding: \
 							{} - {} = {}",
 							preliminary_callback_overhead,
 							actual_callback_weight,
 							difference,
 						);
-						actual_weight -= difference;
+						actual_weight = actual_weight.saturating_sub(difference);
 					}
 					None => {
-						debug_assert!(false, "The delivery confirmation callback is wrong");
-						log::trace!(
+						debug_assert!(false, "T::OnDeliveryConfirmed callback consumed too much weight.");
+						log::error!(
 							target: "runtime::bridge-messages",
-							"Messages delivery callback has returned more weight that it may spent: \
+							"T::OnDeliveryConfirmed callback has spent more weight that it is allowed to: \
 							{} vs {}",
 							preliminary_callback_overhead,
 							actual_callback_weight,
@@ -859,7 +862,7 @@ fn send_message<T: Config<I>, I: 'static>(
 	ensure_normal_operating_mode::<T, I>()?;
 
 	// initially, actual (post-dispatch) weight is equal to pre-dispatch weight
-	let mut actual_weight = T::WeightInfo::send_message_weight(&payload);
+	let mut actual_weight = T::WeightInfo::send_message_weight(&payload, T::DbWeight::get());
 
 	// let's first check if message can be delivered to target chain
 	T::TargetHeaderChain::verify_message(&payload).map_err(|err| {
@@ -913,6 +916,36 @@ fn send_message<T: Config<I>, I: 'static>(
 		payload: encoded_payload,
 		fee: delivery_and_dispatch_fee,
 	});
+	// Guaranteed to be called outside only when the message is accepted.
+	// We assume that the maximum weight call back used is `single_message_callback_overhead`, so do not perform
+	// complex db operation in callback. If you want to, put these magic logic in outside pallet and control
+	// the weight there.
+	let single_message_callback_overhead = T::WeightInfo::single_message_callback_overhead(T::DbWeight::get());
+	let actual_callback_weight = T::OnMessageAccepted::on_messages_accepted(&lane_id, &nonce);
+	match single_message_callback_overhead.checked_sub(actual_callback_weight) {
+		Some(difference) if difference == 0 => (),
+		Some(difference) => {
+			log::trace!(
+				target: "runtime::bridge-messages",
+				"T::OnMessageAccepted callback has spent less weight than expected. Refunding: \
+				{} - {} = {}",
+				single_message_callback_overhead,
+				actual_callback_weight,
+				difference,
+			);
+			actual_weight = actual_weight.saturating_sub(difference);
+		}
+		None => {
+			debug_assert!(false, "T::OnMessageAccepted callback consumed too much weight.");
+			log::error!(
+				target: "runtime::bridge-messages",
+				"T::OnMessageAccepted callback has spent more weight that it is allowed to: \
+				{} vs {}",
+				single_message_callback_overhead,
+				actual_callback_weight,
+			);
+		}
+	}
 
 	// message sender pays for pruning at most `MaxMessagesToPruneAtOnce` messages
 	// the cost of pruning every message is roughly single db write
@@ -1114,7 +1147,7 @@ mod tests {
 	use crate::mock::{
 		message, message_payload, run_test, unrewarded_relayer, Event as TestEvent, Origin,
 		TestMessageDeliveryAndDispatchPayment, TestMessagesDeliveryProof, TestMessagesParameter, TestMessagesProof,
-		TestOnDeliveryConfirmed1, TestOnDeliveryConfirmed2, TestRuntime, TokenConversionRate,
+		TestOnDeliveryConfirmed1, TestOnDeliveryConfirmed2, TestOnMessageAccepted, TestRuntime, TokenConversionRate,
 		PAYLOAD_REJECTED_BY_TARGET_CHAIN, REGULAR_PAYLOAD, TEST_LANE_ID, TEST_RELAYER_A, TEST_RELAYER_B,
 	};
 	use bp_messages::{UnrewardedRelayer, UnrewardedRelayersState};
@@ -2259,4 +2292,40 @@ mod tests {
 			);
 		});
 	}
+
+	#[test]
+	fn message_accepted_callbacks_are_called() {
+		run_test(|| {
+			send_regular_message();
+			TestOnMessageAccepted::ensure_called(&TEST_LANE_ID, &1);
+		});
+	}
+
+	#[test]
+	#[should_panic]
+	fn message_accepted_panics_in_debug_mode_if_callback_is_wrong() {
+		run_test(|| {
+			TestOnMessageAccepted::set_consumed_weight_per_message(crate::mock::DbWeight::get().reads_writes(2, 2));
+			send_regular_message();
+		});
+	}
+
+	#[test]
+	fn message_accepted_refunds_non_zero_weight() {
+		run_test(|| {
+			TestOnMessageAccepted::set_consumed_weight_per_message(crate::mock::DbWeight::get().writes(1));
+			let actual_callback_weight = send_regular_message();
+			let pre_dispatch_weight = <TestRuntime as Config>::WeightInfo::send_message_weight(
+				&REGULAR_PAYLOAD,
+				crate::mock::DbWeight::get(),
+			);
+			let prune_weight =
+				crate::mock::DbWeight::get().writes(<TestRuntime as Config>::MaxMessagesToPruneAtOnce::get());
+
+			assert_eq!(
+				pre_dispatch_weight.saturating_sub(actual_callback_weight),
+				crate::mock::DbWeight::get().reads(1).saturating_add(prune_weight)
+			);
+		});
+	}
 }

bridges/modules/messages/src/mock.rs

@@ -22,8 +22,8 @@ use crate::Config;
 use bitvec::prelude::*;
 use bp_messages::{
 	source_chain::{
-		LaneMessageVerifier, MessageDeliveryAndDispatchPayment, OnDeliveryConfirmed, RelayersRewards, Sender,
-		TargetHeaderChain,
+		LaneMessageVerifier, MessageDeliveryAndDispatchPayment, OnDeliveryConfirmed, OnMessageAccepted,
+		RelayersRewards, Sender, TargetHeaderChain,
 	},
 	target_chain::{DispatchMessage, MessageDispatch, ProvedLaneMessages, ProvedMessages, SourceHeaderChain},
 	DeliveredMessages, InboundLaneData, LaneId, Message, MessageData, MessageKey, MessageNonce, OutboundLaneData,
@@ -178,6 +178,7 @@ impl Config for TestRuntime {
 	type TargetHeaderChain = TestTargetHeaderChain;
 	type LaneMessageVerifier = TestLaneMessageVerifier;
 	type MessageDeliveryAndDispatchPayment = TestMessageDeliveryAndDispatchPayment;
+	type OnMessageAccepted = TestOnMessageAccepted;
 	type OnDeliveryConfirmed = (TestOnDeliveryConfirmed1, TestOnDeliveryConfirmed2);
 
 	type SourceHeaderChain = TestSourceHeaderChain;
@@ -356,6 +357,35 @@ impl MessageDeliveryAndDispatchPayment<AccountId, TestMessageFee> for TestMessag
 	}
 }
 
+#[derive(Debug)]
+pub struct TestOnMessageAccepted;
+
+impl TestOnMessageAccepted {
+	/// Verify that the callback has been called when the message is accepted.
+	pub fn ensure_called(lane: &LaneId, message: &MessageNonce) {
+		let key = (b"TestOnMessageAccepted", lane, message).encode();
+		assert_eq!(frame_support::storage::unhashed::get(&key), Some(true));
+	}
+
+	/// Set consumed weight returned by the callback.
+	pub fn set_consumed_weight_per_message(weight: Weight) {
+		frame_support::storage::unhashed::put(b"TestOnMessageAccepted_Weight", &weight);
+	}
+
+	/// Get consumed weight returned by the callback.
+	pub fn get_consumed_weight_per_message() -> Option<Weight> {
+		frame_support::storage::unhashed::get(b"TestOnMessageAccepted_Weight")
+	}
+}
+
+impl OnMessageAccepted for TestOnMessageAccepted {
+	fn on_messages_accepted(lane: &LaneId, message: &MessageNonce) -> Weight {
+		let key = (b"TestOnMessageAccepted", lane, message).encode();
+		frame_support::storage::unhashed::put(&key, &true);
+		Self::get_consumed_weight_per_message().unwrap_or_else(|| DbWeight::get().reads_writes(1, 1))
+	}
+}
+
 /// First on-messages-delivered callback.
 #[derive(Debug)]
 pub struct TestOnDeliveryConfirmed1;

bridges/modules/messages/src/weights_ext.rs

@@ -189,11 +189,14 @@ pub trait WeightInfoExt: WeightInfo {
 	// Functions that are directly mapped to extrinsics weights.
 
 	/// Weight of message send extrinsic.
-	fn send_message_weight(message: &impl Size) -> Weight {
+	fn send_message_weight(message: &impl Size, db_weight: RuntimeDbWeight) -> Weight {
 		let transaction_overhead = Self::send_message_overhead();
 		let message_size_overhead = Self::send_message_size_overhead(message.size_hint());
+		let call_back_overhead = Self::single_message_callback_overhead(db_weight);
 
-		transaction_overhead.saturating_add(message_size_overhead)
+		transaction_overhead
+			.saturating_add(message_size_overhead)
+			.saturating_add(call_back_overhead)
 	}
 
 	/// Weight of message delivery extrinsic.
@@ -341,7 +344,11 @@ pub trait WeightInfoExt: WeightInfo {
 		Self::receive_single_message_proof().saturating_sub(Self::receive_single_prepaid_message_proof())
 	}
 
-	/// Returns pre-dispatch weight of single message delivery callback call.
+	/// Returns pre-dispatch weight of single callback call.
+	///
+	/// When benchmarking the weight please take into consideration both the `OnMessageAccepted` and
+	/// `OnDeliveryConfirmed` callbacks. The method should return the greater of the two, because it's
+	/// used to estimate the weight in both contexts.
 	fn single_message_callback_overhead(db_weight: RuntimeDbWeight) -> Weight {
 		db_weight.reads_writes(1, 1)
 	}

bridges/primitives/messages/src/source_chain.rs

@@ -173,6 +173,18 @@ impl OnDeliveryConfirmed for Tuple {
 	}
 }
 
+/// Handler for messages have been accepted
+pub trait OnMessageAccepted {
+	/// Called when a message has been accepted by message pallet.
+	fn on_messages_accepted(lane: &LaneId, message: &MessageNonce) -> Weight;
+}
+
+impl OnMessageAccepted for () {
+	fn on_messages_accepted(_lane: &LaneId, _message: &MessageNonce) -> Weight {
+		0
+	}
+}
+
 /// Structure that may be used in place of `TargetHeaderChain`, `LaneMessageVerifier` and
 /// `MessageDeliveryAndDispatchPayment` on chains, where outbound messages are forbidden.
 pub struct ForbidOutboundMessages;