Introduce weight refund

bin/node/executor/tests/basic.rs

@@ -342,11 +342,6 @@ fn full_native_block_import_works() {
 				)),
 				topics: vec![],
 			},
-			EventRecord {
-				phase: Phase::ApplyExtrinsic(1),
-				event: Event::pallet_treasury(pallet_treasury::RawEvent::Deposit(fees * 8 / 10)),
-				topics: vec![],
-			},
 			EventRecord {
 				phase: Phase::ApplyExtrinsic(1),
 				event: Event::pallet_balances(pallet_balances::RawEvent::Transfer(
@@ -356,6 +351,11 @@ fn full_native_block_import_works() {
 				)),
 				topics: vec![],
 			},
+			EventRecord {
+				phase: Phase::ApplyExtrinsic(1),
+				event: Event::pallet_treasury(pallet_treasury::RawEvent::Deposit(fees * 8 / 10)),
+				topics: vec![],
+			},
 			EventRecord {
 				phase: Phase::ApplyExtrinsic(1),
 				event: Event::frame_system(frame_system::RawEvent::ExtrinsicSuccess(
@@ -395,11 +395,6 @@ fn full_native_block_import_works() {
 				)),
 				topics: vec![],
 			},
-			EventRecord {
-				phase: Phase::ApplyExtrinsic(1),
-				event: Event::pallet_treasury(pallet_treasury::RawEvent::Deposit(fees * 8 / 10)),
-				topics: vec![],
-			},
 			EventRecord {
 				phase: Phase::ApplyExtrinsic(1),
 				event: Event::pallet_balances(
@@ -413,14 +408,14 @@ fn full_native_block_import_works() {
 			},
 			EventRecord {
 				phase: Phase::ApplyExtrinsic(1),
-				event: Event::frame_system(frame_system::RawEvent::ExtrinsicSuccess(
-					DispatchInfo { weight: 1000000, class: DispatchClass::Normal, pays_fee: true }
-				)),
+				event: Event::pallet_treasury(pallet_treasury::RawEvent::Deposit(fees * 8 / 10)),
 				topics: vec![],
 			},
 			EventRecord {
-				phase: Phase::ApplyExtrinsic(2),
-				event: Event::pallet_treasury(pallet_treasury::RawEvent::Deposit(fees * 8 / 10)),
+				phase: Phase::ApplyExtrinsic(1),
+				event: Event::frame_system(frame_system::RawEvent::ExtrinsicSuccess(
+					DispatchInfo { weight: 1000000, class: DispatchClass::Normal, pays_fee: true }
+				)),
 				topics: vec![],
 			},
 			EventRecord {
@@ -434,6 +429,11 @@ fn full_native_block_import_works() {
 				),
 				topics: vec![],
 			},
+			EventRecord {
+				phase: Phase::ApplyExtrinsic(2),
+				event: Event::pallet_treasury(pallet_treasury::RawEvent::Deposit(fees * 8 / 10)),
+				topics: vec![],
+			},
 			EventRecord {
 				phase: Phase::ApplyExtrinsic(2),
 				event: Event::frame_system(frame_system::RawEvent::ExtrinsicSuccess(

frame/support/src/weights.rs

@@ -139,6 +139,21 @@ pub struct PostDispatchInfo {
 	pub actual_weight: Option<Weight>,
 }
 
+impl PostDispatchInfo {
+	/// Calculate how much (if any) weight was not used by the `Dispatchable`.
+	pub fn calc_unspent(&self, info: &DispatchInfo) -> Weight {
+		if let Some(actual_weight) = self.actual_weight {
+			if actual_weight >= info.weight {
+				0
+			} else {
+				info.weight - actual_weight
+			}
+		} else {
+			0
+		}
+	}
+}
+
 impl From<Option<Weight>> for PostDispatchInfo {
 	fn from(actual_weight: Option<Weight>) -> Self {
 		Self {

frame/system/src/lib.rs

@@ -120,7 +120,7 @@ use frame_support::{
 		Contains, Get, ModuleToIndex, OnNewAccount, OnKilledAccount, IsDeadAccount, Happened,
 		StoredMap, EnsureOrigin,
 	},
-	weights::{Weight, DispatchInfo, DispatchClass, SimpleDispatchInfo, FunctionOf}
+	weights::{Weight, DispatchInfo, PostDispatchInfo, DispatchClass, SimpleDispatchInfo, FunctionOf}
 };
 use codec::{Encode, Decode, FullCodec, EncodeLike};
 
@@ -1169,7 +1169,7 @@ pub fn split_inner<T, R, S>(option: Option<T>, splitter: impl FnOnce(T) -> (R, S
 pub struct CheckWeight<T: Trait + Send + Sync>(PhantomData<T>);
 
 impl<T: Trait + Send + Sync> CheckWeight<T> where
-	T::Call: Dispatchable<Info=DispatchInfo>
+	T::Call: Dispatchable<Info=DispatchInfo, PostInfo=PostDispatchInfo>
 {
 	/// Get the quota ratio of each dispatch class type. This indicates that all operational
 	/// dispatches can use the full capacity of any resource, while user-triggered ones can consume
@@ -1264,7 +1264,7 @@ impl<T: Trait + Send + Sync> CheckWeight<T> where
 }
 
 impl<T: Trait + Send + Sync> SignedExtension for CheckWeight<T> where
-	T::Call: Dispatchable<Info=DispatchInfo>
+	T::Call: Dispatchable<Info=DispatchInfo, PostInfo=PostDispatchInfo>
 {
 	type AccountId = T::AccountId;
 	type Call = T::Call;
@@ -1319,7 +1319,7 @@ impl<T: Trait + Send + Sync> SignedExtension for CheckWeight<T> where
 	fn post_dispatch(
 		_pre: Self::Pre,
 		info: &DispatchInfoOf<Self::Call>,
-		_post_info: &PostDispatchInfoOf<Self::Call>,
+		post_info: &PostDispatchInfoOf<Self::Call>,
 		_len: usize,
 		result: &DispatchResult,
 	) -> Result<(), TransactionValidityError> {
@@ -1329,6 +1329,14 @@ impl<T: Trait + Send + Sync> SignedExtension for CheckWeight<T> where
 		if info.class == DispatchClass::Mandatory && result.is_err() {
 			Err(InvalidTransaction::BadMandatory)?
 		}
+
+		let unspent = post_info.calc_unspent(info);
+		if unspent > 0 {
+			AllExtrinsicsWeight::mutate(|weight| {
+				*weight = weight.map(|w| w.saturating_sub(unspent));
+			})
+		}
+
 		Ok(())
 	}
 }
@@ -1624,7 +1632,7 @@ mod tests {
 		type Origin = ();
 		type Trait = ();
 		type Info = DispatchInfo;
-		type PostInfo = ();
+		type PostInfo = PostDispatchInfo;
 		fn dispatch(self, _origin: Self::Origin)
 			-> sp_runtime::DispatchResultWithInfo<Self::PostInfo> {
 				panic!("Do not use dummy implementation for dispatch.");

frame/transaction-payment/src/lib.rs

@@ -36,7 +36,8 @@ use codec::{Encode, Decode};
 use frame_support::{
 	decl_storage, decl_module,
 	traits::{Currency, Get, OnUnbalanced, ExistenceRequirement, WithdrawReason, Imbalance},
-	weights::{Weight, DispatchInfo, GetDispatchInfo},
+	weights::{Weight, DispatchInfo, PostDispatchInfo, GetDispatchInfo},
+	dispatch::DispatchResult,
 };
 use sp_runtime::{
 	Fixed64,
@@ -46,7 +47,7 @@ use sp_runtime::{
 	},
 	traits::{
 		Zero, Saturating, SignedExtension, SaturatedConversion, Convert, Dispatchable,
-		DispatchInfoOf,
+		DispatchInfoOf, PostDispatchInfoOf,
 	},
 };
 use pallet_transaction_payment_rpc_runtime_api::RuntimeDispatchInfo;
@@ -102,7 +103,7 @@ decl_module! {
 }
 
 impl<T: Trait> Module<T> where
-	T::Call: Dispatchable<Info=DispatchInfo>,
+	T::Call: Dispatchable<Info=DispatchInfo, PostInfo=PostDispatchInfo>,
 {
 	/// Query the data that we know about the fee of a given `call`.
 	///
@@ -140,7 +141,8 @@ impl<T: Trait> Module<T> where
 pub struct ChargeTransactionPayment<T: Trait + Send + Sync>(#[codec(compact)] BalanceOf<T>);
 
 impl<T: Trait + Send + Sync> ChargeTransactionPayment<T> where
-	T::Call: Dispatchable<Info=DispatchInfo>,
+	T::Call: Dispatchable<Info=DispatchInfo, PostInfo=PostDispatchInfo>,
+	BalanceOf<T>: Send + Sync,
 {
 	/// utility constructor. Used only in client/factory code.
 	pub fn from(fee: BalanceOf<T>) -> Self {
@@ -165,22 +167,12 @@ impl<T: Trait + Send + Sync> ChargeTransactionPayment<T> where
 		len: u32,
 		info: &DispatchInfoOf<T::Call>,
 		tip: BalanceOf<T>,
-	) -> BalanceOf<T>
-	where
-		BalanceOf<T>: Sync + Send,
-	{
+	) -> BalanceOf<T> {
 		if info.pays_fee {
 			let len = <BalanceOf<T>>::from(len);
 			let per_byte = T::TransactionByteFee::get();
 			let len_fee = per_byte.saturating_mul(len);
-
-			let weight_fee = {
-				// cap the weight to the maximum defined in runtime, otherwise it will be the
-				// `Bounded` maximum of its data type, which is not desired.
-				let capped_weight = info.weight
-					.min(<T as frame_system::Trait>::MaximumBlockWeight::get());
-				T::WeightToFee::convert(capped_weight)
-			};
+			let weight_fee = Self::compute_weight_fee(info.weight);
 
 			// the adjustable part of the fee
 			let adjustable_fee = len_fee.saturating_add(weight_fee);
@@ -194,6 +186,42 @@ impl<T: Trait + Send + Sync> ChargeTransactionPayment<T> where
 			tip
 		}
 	}
+
+	fn compute_weight_fee(weight: Weight) -> BalanceOf<T> {
+		// cap the weight to the maximum defined in runtime, otherwise it will be the
+		// `Bounded` maximum of its data type, which is not desired.
+		let capped_weight = weight.min(<T as frame_system::Trait>::MaximumBlockWeight::get());
+		T::WeightToFee::convert(capped_weight)
+	}
+
+	fn withdraw_fee(
+		&self,
+		who: &T::AccountId,
+		info: &DispatchInfoOf<T::Call>,
+		len: usize,
+	) -> Result<(BalanceOf<T>, Option<NegativeImbalanceOf<T>>), TransactionValidityError> {
+		let tip = self.0;
+		let fee = Self::compute_fee(len as u32, info, tip);
+
+		// Only mess with balances if fee is not zero.
+		if fee.is_zero() {
+			return Ok((fee, None));
+		}
+
+		match T::Currency::withdraw(
+			who,
+			fee,
+			if tip.is_zero() {
+				WithdrawReason::TransactionPayment.into()
+			} else {
+				WithdrawReason::TransactionPayment | WithdrawReason::Tip
+			},
+			ExistenceRequirement::KeepAlive,
+		) {
+			Ok(imbalance) => Ok((fee, Some(imbalance))),
+			Err(_) => Err(InvalidTransaction::Payment.into()),
+		}
+	}
 }
 
 impl<T: Trait + Send + Sync> sp_std::fmt::Debug for ChargeTransactionPayment<T> {
@@ -209,13 +237,13 @@ impl<T: Trait + Send + Sync> sp_std::fmt::Debug for ChargeTransactionPayment<T>
 
 impl<T: Trait + Send + Sync> SignedExtension for ChargeTransactionPayment<T> where
 	BalanceOf<T>: Send + Sync,
-	T::Call: Dispatchable<Info=DispatchInfo>,
+	T::Call: Dispatchable<Info=DispatchInfo, PostInfo=PostDispatchInfo>,
 {
 	const IDENTIFIER: &'static str = "ChargeTransactionPayment";
 	type AccountId = T::AccountId;
 	type Call = T::Call;
 	type AdditionalSigned = ();
-	type Pre = ();
+	type Pre = (BalanceOf<T>, Self::AccountId, Option<NegativeImbalanceOf<T>>);
 	fn additional_signed(&self) -> sp_std::result::Result<(), TransactionValidityError> { Ok(()) }
 
 	fn validate(
@@ -225,35 +253,55 @@ impl<T: Trait + Send + Sync> SignedExtension for ChargeTransactionPayment<T> whe
 		info: &DispatchInfoOf<Self::Call>,
 		len: usize,
 	) -> TransactionValidity {
-		// pay any fees.
-		let tip = self.0;
-		let fee = Self::compute_fee(len as u32, info, tip);
-		// Only mess with balances if fee is not zero.
-		if !fee.is_zero() {
-			let imbalance = match T::Currency::withdraw(
-				who,
-				fee,
-				if tip.is_zero() {
-					WithdrawReason::TransactionPayment.into()
-				} else {
-					WithdrawReason::TransactionPayment | WithdrawReason::Tip
-				},
-				ExistenceRequirement::KeepAlive,
-			) {
-				Ok(imbalance) => imbalance,
-				Err(_) => return InvalidTransaction::Payment.into(),
-			};
-			let imbalances = imbalance.split(tip);
-			T::OnTransactionPayment::on_unbalanceds(Some(imbalances.0).into_iter()
-				.chain(Some(imbalances.1)));
-		}
+		let (fee, _) = self.withdraw_fee(who, info, len)?;
 
 		let mut r = ValidTransaction::default();
 		// NOTE: we probably want to maximize the _fee (of any type) per weight unit_ here, which
 		// will be a bit more than setting the priority to tip. For now, this is enough.
 		r.priority = fee.saturated_into::<TransactionPriority>();
 		Ok(r)
 	}
+
+	fn pre_dispatch(
+		self,
+		who: &Self::AccountId,
+		_call: &Self::Call,
+		info: &DispatchInfoOf<Self::Call>,
+		len: usize
+	) -> Result<Self::Pre, TransactionValidityError> {
+		let (_, imbalance) = self.withdraw_fee(who, info, len)?;
+		Ok((self.0, who.clone(), imbalance))
+	}
+
+	fn post_dispatch(
+		pre: Self::Pre,
+		info: &DispatchInfoOf<Self::Call>,
+		post_info: &PostDispatchInfoOf<Self::Call>,
+		_len: usize,
+		_result: &DispatchResult,
+	) -> Result<(), TransactionValidityError> {
+		let (tip, who, imbalance) = pre;
+		if let Some(payed) = imbalance {
+			let refund = Self::compute_weight_fee(post_info.calc_unspent(info));
+			let actual_payment = match T::Currency::deposit_into_existing(&who, refund) {
+				Ok(refund_imbalance) => {
+					// The refund cannot be larger than the up front payed max weight.
+					// `PostDispatchInfo::calc_unspent` guards against such a case.
+					match payed.offset(refund_imbalance) {
+						Ok(actual_payment) => actual_payment,
+						Err(_) => return Err(InvalidTransaction::Payment.into()),
+					}
+				}
+				// We do not recreate the account using the refund. The up front payment
+				// is gone in that case.
+				Err(_) => payed,
+			};
+			let imbalances = actual_payment.split(tip);
+			T::OnTransactionPayment::on_unbalanceds(Some(imbalances.0).into_iter()
+				.chain(Some(imbalances.1)));
+		}
+		Ok(())
+	}
 }
 
 #[cfg(test)]