pre and post dispatch weights of OnDeliveryConfirmed callback (#1040)

* pre and post dispatch weights of OnDeliveryConfirmed callback

* Update modules/messages/README.md

Co-authored-by: Tomasz Drwięga <tomusdrw@users.noreply.github.com>

* clippy + compilation

* fix test issue from parallel PR

Co-authored-by: Tomasz Drwięga <tomusdrw@users.noreply.github.com>

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

@@ -688,6 +688,7 @@ mod tests {
 			bp_millau::ADDITIONAL_MESSAGE_BYTE_DELIVERY_WEIGHT,
 			bp_millau::MAX_SINGLE_MESSAGE_DELIVERY_CONFIRMATION_TX_WEIGHT,
 			bp_millau::PAY_INBOUND_DISPATCH_FEE_WEIGHT,
+			DbWeight::get(),
 		);
 
 		let max_incoming_message_proof_size = bp_rialto::EXTRA_STORAGE_PROOF_SIZE.saturating_add(
@@ -712,6 +713,7 @@ mod tests {
 			max_incoming_inbound_lane_data_proof_size,
 			bp_rialto::MAX_UNREWARDED_RELAYER_ENTRIES_AT_INBOUND_LANE,
 			bp_rialto::MAX_UNCONFIRMED_MESSAGES_AT_INBOUND_LANE,
+			DbWeight::get(),
 		);
 	}
 }

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

@@ -1125,6 +1125,7 @@ mod tests {
 			bp_rialto::ADDITIONAL_MESSAGE_BYTE_DELIVERY_WEIGHT,
 			bp_rialto::MAX_SINGLE_MESSAGE_DELIVERY_CONFIRMATION_TX_WEIGHT,
 			bp_rialto::PAY_INBOUND_DISPATCH_FEE_WEIGHT,
+			DbWeight::get(),
 		);
 
 		let max_incoming_message_proof_size = bp_millau::EXTRA_STORAGE_PROOF_SIZE.saturating_add(
@@ -1149,6 +1150,7 @@ mod tests {
 			max_incoming_inbound_lane_data_proof_size,
 			bp_millau::MAX_UNREWARDED_RELAYER_ENTRIES_AT_INBOUND_LANE,
 			bp_millau::MAX_UNCONFIRMED_MESSAGES_AT_INBOUND_LANE,
+			DbWeight::get(),
 		);
 	}
 

bridges/modules/messages/README.md

@@ -354,7 +354,7 @@ Both conditions are verified by `pallet_bridge_messages::ensure_weights_are_corr
 `pallet_bridge_messages::ensure_able_to_receive_messages` functions, which must be called from every
 runtime's tests.
 
-### Post-dispatch weight refunds of the `receive_messages_proof` call
+#### Post-dispatch weight refunds of the `receive_messages_proof` call
 
 Weight formula of the `receive_messages_proof` call assumes that the dispatch fee of every message is
 paid at the target chain (where call is executed), that every message will be dispatched and that
@@ -388,6 +388,7 @@ The weight formula is:
 Weight = BaseWeight + MessagesCount * MessageConfirmationWeight
        + RelayersCount * RelayerRewardWeight
        + Max(0, ActualProofSize - ExpectedProofSize) * ProofByteDeliveryWeight
+       + MessagesCount * (DbReadWeight + DbWriteWeight)
 ```
 
 Where:
@@ -403,6 +404,15 @@ Where:
 | `ExpectedProofSize`       | `EXTRA_STORAGE_PROOF_SIZE`                                                                                            | Size of proof that we are expecting                                                                                                                                                                     |
 | `ProofByteDeliveryWeight` | `(receive_single_message_proof_16_kb - receive_single_message_proof_1_kb) / (15 * 1024)`                              | Weight of processing every additional proof byte over `ExpectedProofSize` limit. We're using the same formula, as for message delivery, because proof mechanism is assumed to be the same in both cases |
 
+#### Post-dispatch weight refunds of the `receive_messages_delivery_proof` call
+
+Weight formula of the `receive_messages_delivery_proof` call assumes that all messages in the proof
+are actually delivered (so there are no already confirmed messages) and every messages is processed
+by the `OnDeliveryConfirmed` callback. This means that for every message, we're adding single db read
+weight and single db write weight. If, by some reason, messages are not processed by the
+`OnDeliveryConfirmed` callback, or their processing is faster than that additional weight, the
+difference is refunded to the submitter.
+
 #### Why we're always able to craft `receive_messages_delivery_proof` transaction?
 
 There can be at most `<PeerRuntime as pallet_bridge_messages::Config>::MaxUnconfirmedMessagesAtInboundLane`

bridges/modules/messages/src/lib.rs

@@ -594,14 +594,35 @@ decl_module! {
 		}
 
 		/// Receive messages delivery proof from bridged chain.
-		#[weight = T::WeightInfo::receive_messages_delivery_proof_weight(proof, relayers_state)]
+		#[weight = T::WeightInfo::receive_messages_delivery_proof_weight(
+			proof,
+			relayers_state,
+			T::DbWeight::get(),
+		)]
 		pub fn receive_messages_delivery_proof(
 			origin,
 			proof: MessagesDeliveryProofOf<T, I>,
 			relayers_state: UnrewardedRelayersState,
-		) -> DispatchResult {
+		) -> DispatchResultWithPostInfo {
 			ensure_not_halted::<T, I>()?;
 
+			// why do we need to know the weight of this (`receive_messages_delivery_proof`) call? Because
+			// we may want to return some funds for messages that are not processed by the delivery callback,
+			// or if their actual processing weight is less than accounted by weight formula.
+			// So to refund relayer, we need to:
+			//
+			// ActualWeight = DeclaredWeight - UnspentCallbackWeight
+			//
+			// The DeclaredWeight is exactly what's computed here. Unfortunately it is impossible
+			// to get pre-computed value (and it has been already computed by the executive).
+			let single_message_callback_overhead = T::WeightInfo::single_message_callback_overhead(T::DbWeight::get());
+			let declared_weight = T::WeightInfo::receive_messages_delivery_proof_weight(
+				&proof,
+				&relayers_state,
+				T::DbWeight::get(),
+			);
+			let mut actual_weight = declared_weight;
+
 			let confirmation_relayer = ensure_signed(origin)?;
 			let (lane_id, lane_data) = T::TargetHeaderChain::verify_messages_delivery_proof(proof).map_err(|err| {
 				log::trace!(
@@ -641,7 +662,37 @@ decl_module! {
 			};
 			if let Some(confirmed_messages) = confirmed_messages {
 				// handle messages delivery confirmation
-				T::OnDeliveryConfirmed::on_messages_delivered(&lane_id, &confirmed_messages);
+				let preliminary_callback_overhead = relayers_state.total_messages.saturating_mul(
+					single_message_callback_overhead
+				);
+				let actual_callback_weight = T::OnDeliveryConfirmed::on_messages_delivered(
+					&lane_id,
+					&confirmed_messages,
+				);
+				match preliminary_callback_overhead.checked_sub(actual_callback_weight) {
+					Some(difference) if difference == 0 => (),
+					Some(difference) => {
+						log::trace!(
+							target: "runtime::bridge-messages",
+							"Messages delivery callback has returned unspent weight to refund the submitter: \
+							{} - {} = {}",
+							preliminary_callback_overhead,
+							actual_callback_weight,
+							difference,
+						);
+						actual_weight -= difference;
+					},
+					None => {
+						debug_assert!(false, "The delivery confirmation callback is wrong");
+						log::trace!(
+							target: "runtime::bridge-messages",
+							"Messages delivery callback has returned more weight that it may spent: \
+							{} vs {}",
+							preliminary_callback_overhead,
+							actual_callback_weight,
+						);
+					}
+				}
 
 				// emit 'delivered' event
 				let received_range = confirmed_messages.begin..=confirmed_messages.end;
@@ -684,7 +735,10 @@ decl_module! {
 				lane_id,
 			);
 
-			Ok(())
+			Ok(PostDispatchInfo {
+				actual_weight: Some(actual_weight),
+				pays_fee: Pays::Yes,
+			})
 		}
 	}
 }
@@ -962,8 +1016,8 @@ mod tests {
 	use crate::mock::{
 		message, message_payload, run_test, unrewarded_relayer, Event as TestEvent, Origin,
 		TestMessageDeliveryAndDispatchPayment, TestMessagesDeliveryProof, TestMessagesParameter, TestMessagesProof,
-		TestRuntime, TokenConversionRate, PAYLOAD_REJECTED_BY_TARGET_CHAIN, REGULAR_PAYLOAD, TEST_LANE_ID,
-		TEST_RELAYER_A, TEST_RELAYER_B,
+		TestOnDeliveryConfirmed1, TestOnDeliveryConfirmed2, TestRuntime, TokenConversionRate,
+		PAYLOAD_REJECTED_BY_TARGET_CHAIN, REGULAR_PAYLOAD, TEST_LANE_ID, TEST_RELAYER_A, TEST_RELAYER_B,
 	};
 	use bp_messages::{UnrewardedRelayer, UnrewardedRelayersState};
 	use frame_support::{assert_noop, assert_ok};
@@ -1260,10 +1314,14 @@ mod tests {
 						TEST_LANE_ID,
 						InboundLaneData {
 							last_confirmed_nonce: 1,
-							..Default::default()
+							relayers: vec![unrewarded_relayer(1, 1, TEST_RELAYER_A)].into_iter().collect(),
 						},
 					))),
-					Default::default(),
+					UnrewardedRelayersState {
+						unrewarded_relayer_entries: 1,
+						messages_in_oldest_entry: 1,
+						total_messages: 1,
+					},
 				),
 				Error::<TestRuntime, DefaultInstance>::Halted,
 			);
@@ -1309,10 +1367,14 @@ mod tests {
 					TEST_LANE_ID,
 					InboundLaneData {
 						last_confirmed_nonce: 1,
-						..Default::default()
+						relayers: vec![unrewarded_relayer(1, 1, TEST_RELAYER_A)].into_iter().collect(),
 					},
 				))),
-				Default::default(),
+				UnrewardedRelayersState {
+					unrewarded_relayer_entries: 1,
+					messages_in_oldest_entry: 1,
+					total_messages: 1,
+				},
 			));
 		});
 	}
@@ -1923,10 +1985,70 @@ mod tests {
 			));
 
 			// ensure that both callbacks have been called twice: for 1+2, then for 3
-			crate::mock::TestOnDeliveryConfirmed1::ensure_called(&TEST_LANE_ID, &delivered_messages_1_and_2);
-			crate::mock::TestOnDeliveryConfirmed1::ensure_called(&TEST_LANE_ID, &delivered_message_3);
-			crate::mock::TestOnDeliveryConfirmed2::ensure_called(&TEST_LANE_ID, &delivered_messages_1_and_2);
-			crate::mock::TestOnDeliveryConfirmed2::ensure_called(&TEST_LANE_ID, &delivered_message_3);
+			TestOnDeliveryConfirmed1::ensure_called(&TEST_LANE_ID, &delivered_messages_1_and_2);
+			TestOnDeliveryConfirmed1::ensure_called(&TEST_LANE_ID, &delivered_message_3);
+			TestOnDeliveryConfirmed2::ensure_called(&TEST_LANE_ID, &delivered_messages_1_and_2);
+			TestOnDeliveryConfirmed2::ensure_called(&TEST_LANE_ID, &delivered_message_3);
+		});
+	}
+
+	fn confirm_3_messages_delivery() -> (Weight, Weight) {
+		send_regular_message();
+		send_regular_message();
+		send_regular_message();
+
+		let proof = TestMessagesDeliveryProof(Ok((
+			TEST_LANE_ID,
+			InboundLaneData {
+				last_confirmed_nonce: 0,
+				relayers: vec![unrewarded_relayer(1, 3, TEST_RELAYER_A)].into_iter().collect(),
+			},
+		)));
+		let relayers_state = UnrewardedRelayersState {
+			unrewarded_relayer_entries: 1,
+			total_messages: 3,
+			..Default::default()
+		};
+		let pre_dispatch_weight = <TestRuntime as Config>::WeightInfo::receive_messages_delivery_proof_weight(
+			&proof,
+			&relayers_state,
+			crate::mock::DbWeight::get(),
+		);
+		let post_dispatch_weight =
+			Pallet::<TestRuntime>::receive_messages_delivery_proof(Origin::signed(1), proof, relayers_state)
+				.expect("confirmation has failed")
+				.actual_weight
+				.expect("receive_messages_delivery_proof always returns Some");
+		(pre_dispatch_weight, post_dispatch_weight)
+	}
+
+	#[test]
+	fn receive_messages_delivery_proof_refunds_zero_weight() {
+		run_test(|| {
+			let (pre_dispatch_weight, post_dispatch_weight) = confirm_3_messages_delivery();
+			assert_eq!(pre_dispatch_weight, post_dispatch_weight);
+		});
+	}
+
+	#[test]
+	fn receive_messages_delivery_proof_refunds_non_zero_weight() {
+		run_test(|| {
+			TestOnDeliveryConfirmed1::set_consumed_weight_per_message(crate::mock::DbWeight::get().writes(1));
+
+			let (pre_dispatch_weight, post_dispatch_weight) = confirm_3_messages_delivery();
+			assert_eq!(
+				pre_dispatch_weight.saturating_sub(post_dispatch_weight),
+				crate::mock::DbWeight::get().reads(1) * 3
+			);
+		});
+	}
+
+	#[test]
+	#[should_panic]
+	fn receive_messages_panics_in_debug_mode_if_callback_is_wrong() {
+		run_test(|| {
+			TestOnDeliveryConfirmed1::set_consumed_weight_per_message(crate::mock::DbWeight::get().reads_writes(2, 2));
+			confirm_3_messages_delivery()
 		});
 	}
 }

bridges/modules/messages/src/mock.rs

@@ -31,7 +31,10 @@ use bp_messages::{
 };
 use bp_runtime::{messages::MessageDispatchResult, Size};
 use codec::{Decode, Encode};
-use frame_support::{parameter_types, weights::Weight};
+use frame_support::{
+	parameter_types,
+	weights::{RuntimeDbWeight, Weight},
+};
 use sp_core::H256;
 use sp_runtime::{
 	testing::Header as SubstrateHeader,
@@ -87,6 +90,7 @@ parameter_types! {
 	pub const MaximumBlockWeight: Weight = 1024;
 	pub const MaximumBlockLength: u32 = 2 * 1024;
 	pub const AvailableBlockRatio: Perbill = Perbill::one();
+	pub const DbWeight: RuntimeDbWeight = RuntimeDbWeight { read: 1, write: 2 };
 }
 
 impl frame_system::Config for TestRuntime {
@@ -110,7 +114,7 @@ impl frame_system::Config for TestRuntime {
 	type SystemWeightInfo = ();
 	type BlockWeights = ();
 	type BlockLength = ();
-	type DbWeight = ();
+	type DbWeight = DbWeight;
 	type SS58Prefix = ();
 	type OnSetCode = ();
 }
@@ -358,12 +362,25 @@ impl TestOnDeliveryConfirmed1 {
 		let key = (b"TestOnDeliveryConfirmed1", lane, messages).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"TestOnDeliveryConfirmed1_Weight", &weight);
+	}
+
+	/// Get consumed weight returned by the callback.
+	pub fn get_consumed_weight_per_message() -> Option<Weight> {
+		frame_support::storage::unhashed::get(b"TestOnDeliveryConfirmed1_Weight")
+	}
 }
 
 impl OnDeliveryConfirmed for TestOnDeliveryConfirmed1 {
-	fn on_messages_delivered(lane: &LaneId, messages: &DeliveredMessages) {
+	fn on_messages_delivered(lane: &LaneId, messages: &DeliveredMessages) -> Weight {
 		let key = (b"TestOnDeliveryConfirmed1", lane, messages).encode();
 		frame_support::storage::unhashed::put(&key, &true);
+		Self::get_consumed_weight_per_message()
+			.unwrap_or_else(|| DbWeight::get().reads_writes(1, 1))
+			.saturating_mul(messages.total_messages())
 	}
 }
 
@@ -380,9 +397,10 @@ impl TestOnDeliveryConfirmed2 {
 }
 
 impl OnDeliveryConfirmed for TestOnDeliveryConfirmed2 {
-	fn on_messages_delivered(lane: &LaneId, messages: &DeliveredMessages) {
+	fn on_messages_delivered(lane: &LaneId, messages: &DeliveredMessages) -> Weight {
 		let key = (b"TestOnDeliveryConfirmed2", lane, messages).encode();
 		frame_support::storage::unhashed::put(&key, &true);
+		0
 	}
 }
 

bridges/modules/messages/src/weights_ext.rs

@@ -20,7 +20,7 @@ use crate::weights::WeightInfo;
 
 use bp_messages::{MessageNonce, UnrewardedRelayersState};
 use bp_runtime::{PreComputedSize, Size};
-use frame_support::weights::Weight;
+use frame_support::weights::{RuntimeDbWeight, Weight};
 
 /// Size of the message being delivered in benchmarks.
 pub const EXPECTED_DEFAULT_MESSAGE_LENGTH: u32 = 128;
@@ -35,6 +35,7 @@ pub fn ensure_weights_are_correct<W: WeightInfoExt>(
 	expected_additional_byte_delivery_weight: Weight,
 	expected_messages_delivery_confirmation_tx_weight: Weight,
 	expected_pay_inbound_dispatch_fee_weight: Weight,
+	db_weight: RuntimeDbWeight,
 ) {
 	// verify `send_message` weight components
 	assert_ne!(W::send_message_overhead(), 0);
@@ -82,6 +83,7 @@ pub fn ensure_weights_are_correct<W: WeightInfoExt>(
 			total_messages: 1,
 			..Default::default()
 		},
+		db_weight,
 	);
 	assert!(
 		actual_messages_delivery_confirmation_tx_weight <= expected_messages_delivery_confirmation_tx_weight,
@@ -138,6 +140,7 @@ pub fn ensure_able_to_receive_confirmation<W: WeightInfoExt>(
 	max_inbound_lane_data_proof_size_from_peer_chain: u32,
 	max_unrewarded_relayer_entries_at_peer_inbound_lane: MessageNonce,
 	max_unconfirmed_messages_at_inbound_lane: MessageNonce,
+	db_weight: RuntimeDbWeight,
 ) {
 	// verify that we're able to receive confirmation of maximal-size
 	let max_confirmation_transaction_size =
@@ -158,6 +161,7 @@ pub fn ensure_able_to_receive_confirmation<W: WeightInfoExt>(
 			total_messages: max_unconfirmed_messages_at_inbound_lane,
 			..Default::default()
 		},
+		db_weight,
 	);
 	assert!(
 		max_confirmation_transaction_dispatch_weight <= max_extrinsic_weight,
@@ -212,7 +216,11 @@ pub trait WeightInfoExt: WeightInfo {
 	}
 
 	/// Weight of confirmation delivery extrinsic.
-	fn receive_messages_delivery_proof_weight(proof: &impl Size, relayers_state: &UnrewardedRelayersState) -> Weight {
+	fn receive_messages_delivery_proof_weight(
+		proof: &impl Size,
+		relayers_state: &UnrewardedRelayersState,
+		db_weight: RuntimeDbWeight,
+	) -> Weight {
 		// basic components of extrinsic weight
 		let transaction_overhead = Self::receive_messages_delivery_proof_overhead();
 		let messages_overhead = Self::receive_messages_delivery_proof_messages_overhead(relayers_state.total_messages);
@@ -225,10 +233,16 @@ pub trait WeightInfoExt: WeightInfo {
 		let proof_size_overhead =
 			Self::storage_proof_size_overhead(actual_proof_size.saturating_sub(expected_proof_size));
 
+		// and cost of calling `OnDeliveryConfirmed::on_messages_delivered()` for every confirmed message
+		let callback_overhead = relayers_state
+			.total_messages
+			.saturating_mul(Self::single_message_callback_overhead(db_weight));
+
 		transaction_overhead
 			.saturating_add(messages_overhead)
 			.saturating_add(relayers_overhead)
 			.saturating_add(proof_size_overhead)
+			.saturating_add(callback_overhead)
 	}
 
 	// Functions that are used by extrinsics weights formulas.
@@ -321,6 +335,11 @@ pub trait WeightInfoExt: WeightInfo {
 	fn pay_inbound_dispatch_fee_overhead() -> Weight {
 		Self::receive_single_message_proof().saturating_sub(Self::receive_single_prepaid_message_proof())
 	}
+
+	/// Returns pre-dispatch weight of single message delivery callback call.
+	fn single_message_callback_overhead(db_weight: RuntimeDbWeight) -> Weight {
+		db_weight.reads_writes(1, 1)
+	}
 }
 
 impl WeightInfoExt for () {

bridges/primitives/messages/src/lib.rs

@@ -226,6 +226,15 @@ impl DeliveredMessages {
 		}
 	}
 
+	/// Return total count of delivered messages.
+	pub fn total_messages(&self) -> MessageNonce {
+		if self.end >= self.begin {
+			self.end - self.begin + 1
+		} else {
+			0
+		}
+	}
+
 	/// Note new dispatched message.
 	pub fn note_dispatched_message(&mut self, dispatch_result: bool) {
 		self.end += 1;

bridges/primitives/messages/src/source_chain.rs

@@ -19,7 +19,7 @@
 use crate::{DeliveredMessages, InboundLaneData, LaneId, MessageNonce, OutboundLaneData};
 
 use bp_runtime::Size;
-use frame_support::{Parameter, RuntimeDebug};
+use frame_support::{weights::Weight, Parameter, RuntimeDebug};
 use sp_std::{collections::btree_map::BTreeMap, fmt::Debug};
 
 /// The sender of the message on the source chain.
@@ -136,12 +136,33 @@ pub trait MessageDeliveryAndDispatchPayment<AccountId, Balance> {
 }
 
 /// Handler for messages delivery confirmation.
-#[impl_trait_for_tuples::impl_for_tuples(30)]
 pub trait OnDeliveryConfirmed {
 	/// Called when we receive confirmation that our messages have been delivered to the
 	/// target chain. The confirmation also has single bit dispatch result for every
-	/// confirmed message (see `DeliveredMessages` for details).
-	fn on_messages_delivered(_lane: &LaneId, _messages: &DeliveredMessages) {}
+	/// confirmed message (see `DeliveredMessages` for details). Guaranteed to be called
+	/// only when at least one message is delivered.
+	///
+	/// Should return total weight consumed by the call.
+	///
+	/// NOTE: messages pallet assumes that maximal weight that may be spent on processing
+	/// single message is single DB read + single DB write. So this function shall never
+	/// return weight that is larger than total number of messages * (db read + db write).
+	/// If your pallet needs more time for processing single message, please do it
+	/// from `on_initialize` call(s) of the next block(s).
+	fn on_messages_delivered(_lane: &LaneId, _messages: &DeliveredMessages) -> Weight;
+}
+
+#[impl_trait_for_tuples::impl_for_tuples(30)]
+impl OnDeliveryConfirmed for Tuple {
+	fn on_messages_delivered(lane: &LaneId, messages: &DeliveredMessages) -> Weight {
+		let mut total_weight: Weight = 0;
+		for_tuples!(
+			#(
+				total_weight = total_weight.saturating_add(Tuple::on_messages_delivered(lane, messages));
+			)*
+		);
+		total_weight
+	}
 }
 
 /// Structure that may be used in place of `TargetHeaderChain`, `LaneMessageVerifier` and