Make Dispatchable return the actual weight consumed

frame/recovery/src/lib.rs

@@ -161,6 +161,7 @@ use frame_support::{
 	decl_module, decl_event, decl_storage, decl_error, ensure,
 	Parameter, RuntimeDebug, weights::{GetDispatchInfo, SimpleDispatchInfo, FunctionOf},
 	traits::{Currency, ReservableCurrency, Get, BalanceStatus},
+	dispatch::PostDispatchInfo,
 };
 use frame_system::{self as system, ensure_signed, ensure_root};
 
@@ -178,7 +179,7 @@ pub trait Trait: frame_system::Trait {
 	type Event: From<Event<Self>> + Into<<Self as frame_system::Trait>::Event>;
 
 	/// The overarching call type.
-	type Call: Parameter + Dispatchable<Origin=Self::Origin> + GetDispatchInfo;
+	type Call: Parameter + Dispatchable<Origin=Self::Origin, PostInfo=PostDispatchInfo> + GetDispatchInfo;
 
 	/// The currency mechanism.
 	type Currency: ReservableCurrency<Self::AccountId>;
@@ -348,6 +349,7 @@ decl_module! {
 			let target = Self::proxy(&who).ok_or(Error::<T>::NotAllowed)?;
 			ensure!(&target == &account, Error::<T>::NotAllowed);
 			call.dispatch(frame_system::RawOrigin::Signed(account).into())
+				.map(|_| ()).map_err(|e| e.error)
 		}
 
 		/// Allow ROOT to bypass the recovery process and set an a rescuer account

frame/scheduler/src/lib.rs

@@ -150,7 +150,11 @@ decl_module! {
 								Lookup::<T>::remove(id);
 							}
 						}
-						Self::deposit_event(RawEvent::Dispatched((now, index), maybe_id, r));
+						Self::deposit_event(RawEvent::Dispatched(
+							(now, index),
+							maybe_id,
+							r.map(|_| ()).map_err(|e| e.error)
+						));
 						result = cumulative_weight;
 						None
 					} else {

frame/sudo/src/lib.rs

@@ -87,7 +87,7 @@
 #![cfg_attr(not(feature = "std"), no_std)]
 
 use sp_std::prelude::*;
-use sp_runtime::{traits::{StaticLookup, Dispatchable}, DispatchError};
+use sp_runtime::traits::{StaticLookup, Dispatchable};
 
 use frame_support::{
 	Parameter, decl_module, decl_event, decl_storage, decl_error, ensure,
@@ -133,7 +133,6 @@ decl_module! {
 			let res = match call.dispatch(frame_system::RawOrigin::Root.into()) {
 				Ok(_) => true,
 				Err(e) => {
-					let e: DispatchError = e.into();
 					sp_runtime::print(e);
 					false
 				}
@@ -192,7 +191,6 @@ decl_module! {
 			let res = match call.dispatch(frame_system::RawOrigin::Signed(who).into()) {
 				Ok(_) => true,
 				Err(e) => {
-					let e: DispatchError = e.into();
 					sp_runtime::print(e);
 					false
 				}

frame/support/src/dispatch.rs

@@ -25,18 +25,35 @@ pub use frame_metadata::{
 };
 pub use crate::weights::{
 	SimpleDispatchInfo, GetDispatchInfo, DispatchInfo, WeighData, ClassifyDispatch,
-	TransactionPriority, Weight, PaysFee,
+	TransactionPriority, Weight, PaysFee, PostDispatchInfo, WithPostDispatchInfo,
 };
-pub use sp_runtime::{traits::Dispatchable, DispatchError, DispatchResult};
+pub use sp_runtime::{traits::Dispatchable, DispatchError};
 pub use crate::traits::{CallMetadata, GetCallMetadata, GetCallName};
 
+/// The return typ of a `Dispatchable` in frame. When returned explicitly from
+/// a dispatchable function it allows overriding the default `PostDispatchInfo`
+/// returned from a dispatch.
+pub type DispatchResultWithPostInfo =
+	sp_runtime::DispatchResultWithInfo<crate::weights::PostDispatchInfo>;
+
+/// Unaugmented version of `DispatchResultWithPostInfo` that can be returned from
+/// dispatchable functions and is automatically converted to the augmented type. Should be
+/// used whenever the `PostDispatchInfo` does not need to be overwritten. As this should
+/// be the common case it is the implicit return type when none is specified.
+pub type DispatchResult = Result<(), sp_runtime::DispatchError>;
+
+/// The error type contained in a `DispatchResultWithPostInfo`.
+pub type DispatchErrorWithPostInfo =
+	sp_runtime::DispatchErrorWithPostInfo<crate::weights::PostDispatchInfo>;
+
+
 /// A type that cannot be instantiated.
 pub enum Never {}
 
 /// Serializable version of Dispatchable.
 /// This value can be used as a "function" in an extrinsic.
 pub trait Callable<T> {
-	type Call: Dispatchable + Codec + Clone + PartialEq + Eq;
+	type Call: Dispatchable<PostInfo=PostDispatchInfo> + Codec + Clone + PartialEq + Eq;
 }
 
 // dirty hack to work around serde_derive issue
@@ -119,6 +136,44 @@ impl<T> Parameter for T where T: Codec + EncodeLike + Clone + Eq + fmt::Debug {}
 /// # fn main() {}
 /// ```
 ///
+/// ### Consuming only portions of the annotated static weight
+///
+/// Per default a callable function consumes all of its static weight as declared via
+/// the #[weight] attribute. However, there are use cases where only a portion of this
+/// weight should be consumed. In that case the static weight is charged pre dispatch and
+/// the difference is refunded post dispatch.
+///
+/// In order to make use of this feature the function must return `DispatchResultWithPostInfo`
+/// in place of the default `DispatchResult`. Then the actually consumed weight can be returned.
+/// To consume a non default weight while returning an error
+/// [`WithPostDispatchInfo::with_weight`](./weight/trait.WithPostDispatchInfo.html) can be used
+/// to augment any error with custom weight information.
+///
+/// ```
+/// # #[macro_use]
+/// # extern crate frame_support;
+/// # use frame_support::dispatch::{DispatchResultWithPostInfo, WithPostDispatchInfo};
+/// # use frame_support::weights::SimpleDispatchInfo;
+/// # use frame_system::{self as system, Trait, ensure_signed};
+/// decl_module! {
+/// 	pub struct Module<T: Trait> for enum Call where origin: T::Origin {
+/// 		#[weight = SimpleDispatchInfo::FixedNormal(1_000_000)]
+/// 		fn my_long_function(origin, do_expensive_calc: bool) -> DispatchResultWithPostInfo {
+/// 			ensure_signed(origin).map_err(|e| e.with_weight(100_000))?;
+/// 			if do_expensive_calc {
+/// 				// do the expensive calculation
+/// 				// ...
+/// 				// return None to indicate that we are using all weight (the default)
+/// 				return Ok(None.into());
+/// 			}
+/// 			// expensive calculation not executed: use only a portion of the weight
+/// 			Ok(Some(100_000).into())
+/// 		}
+/// 	}
+/// }
+/// # fn main() {}
+/// ```
+///
 /// ### Privileged Function Example
 ///
 /// A privileged function checks that the origin of the call is `ROOT`.
@@ -938,7 +993,7 @@ macro_rules! decl_module {
 		$ignore:ident
 		$mod_type:ident<$trait_instance:ident $(, $instance:ident)?> $fn_name:ident $origin:ident $system:ident [ $( $param_name:ident),* ]
 	) => {
-		<$mod_type<$trait_instance $(, $instance)?>>::$fn_name( $origin $(, $param_name )* )
+		<$mod_type<$trait_instance $(, $instance)?>>::$fn_name( $origin $(, $param_name )* ).map(Into::into).map_err(Into::into)
 	};
 
 	// no `deposit_event` function wanted
@@ -1538,7 +1593,8 @@ macro_rules! decl_module {
 		{
 			type Trait = $trait_instance;
 			type Origin = $origin_type;
-			fn dispatch(self, _origin: Self::Origin) -> $crate::sp_runtime::DispatchResult {
+			type PostInfo = $crate::weights::PostDispatchInfo;
+			fn dispatch(self, _origin: Self::Origin) -> $crate::dispatch::DispatchResultWithPostInfo {
 				match self {
 					$(
 						$call_type::$fn_name( $( $param_name ),* ) => {
@@ -1563,10 +1619,10 @@ macro_rules! decl_module {
 			where $( $other_where_bounds )*
 		{
 			#[doc(hidden)]
-			pub fn dispatch<D: $crate::dispatch::Dispatchable<Trait = $trait_instance>>(
+			pub fn dispatch<D: $crate::dispatch::Dispatchable<Trait = $trait_instance, PostInfo = $crate::weights::PostDispatchInfo>>(
 				d: D,
 				origin: D::Origin
-			) -> $crate::sp_runtime::DispatchResult {
+			) -> $crate::dispatch::DispatchResultWithPostInfo {
 				d.dispatch(origin)
 			}
 		}
@@ -1664,10 +1720,11 @@ macro_rules! impl_outer_dispatch {
 		impl $crate::dispatch::Dispatchable for $call_type {
 			type Origin = $origin;
 			type Trait = $call_type;
+			type PostInfo = $crate::weights::PostDispatchInfo;
 			fn dispatch(
 				self,
 				origin: $origin,
-			) -> $crate::sp_runtime::DispatchResult {
+			) -> $crate::dispatch::DispatchResultWithPostInfo {
 				$crate::impl_outer_dispatch! {
 					@DISPATCH_MATCH
 					self

frame/support/src/lib.rs

@@ -209,7 +209,11 @@ macro_rules! assert_err {
 #[cfg(feature = "std")]
 macro_rules! assert_ok {
 	( $x:expr $(,)? ) => {
-		assert_eq!($x, Ok(()));
+		let is = $x;
+		match is {
+			Ok(_) => (),
+			_ => assert!(false, "Expected Ok(_). Got {:#?}", is),
+		}
 	};
 	( $x:expr, $y:expr $(,)? ) => {
 		assert_eq!($x, Ok($y));

frame/support/src/weights.rs

@@ -44,6 +44,7 @@ use sp_runtime::{
 	traits::SignedExtension,
 	generic::{CheckedExtrinsic, UncheckedExtrinsic},
 };
+use crate::dispatch::{DispatchErrorWithPostInfo, DispatchError};
 
 /// Re-export priority as type
 pub use sp_runtime::transaction_validity::TransactionPriority;
@@ -116,6 +117,64 @@ pub struct DispatchInfo {
 	pub pays_fee: bool,
 }
 
+/// Weight information that is only available post dispatch.
+#[derive(Clone, Copy, Eq, PartialEq, Default, RuntimeDebug, Encode, Decode)]
+pub struct PostDispatchInfo {
+	/// Actual weight consumed by a call or `None` which stands for the worst case static weight.
+	pub actual_weight: Option<Weight>,
+}
+
+impl From<Option<Weight>> for PostDispatchInfo {
+	fn from(actual_weight: Option<Weight>) -> Self {
+		Self {
+			actual_weight,
+		}
+	}
+}
+
+impl From<()> for PostDispatchInfo {
+	fn from(_: ()) -> Self {
+		Self {
+			actual_weight: None,
+		}
+	}
+}
+
+impl sp_runtime::traits::Printable for PostDispatchInfo {
+	fn print(&self) {
+		"actual_weight=".print();
+		match self.actual_weight {
+			Some(weight) => weight.print(),
+			None => "max-weight".print(),
+		}
+	}
+}
+
+/// Allows easy conversion from `DispatchError` to `DispatchErrorWithPostInfo` for dispatchables
+/// that want to return a custom a posteriori weight on error.
+pub trait WithPostDispatchInfo {
+	/// Call this on your modules custom errors type in order to return a custom weight on error.
+	///
+	/// # Example
+	///
+	/// ```ignore
+	/// let who = ensure_signed(origin).map_err(|e| e.with_weight(100))?;
+	/// ensure!(who == me, Error::<T>::NotMe.with_weight(200_000));
+	/// ```
+	fn with_weight(self, actual_weight: Weight) -> DispatchErrorWithPostInfo;
+}
+
+impl<T> WithPostDispatchInfo for T where
+	T: Into<DispatchError>
+{
+	fn with_weight(self, actual_weight: Weight) -> DispatchErrorWithPostInfo {
+		DispatchErrorWithPostInfo {
+			post_info: PostDispatchInfo { actual_weight: Some(actual_weight) },
+			error: self.into(),
+		}
+	}
+}
+
 /// A `Dispatchable` function (aka transaction) that can carry some static information along with
 /// it, using the `#[weight]` attribute.
 pub trait GetDispatchInfo {

frame/utility/src/lib.rs

@@ -68,6 +68,7 @@ use sp_io::hashing::blake2_256;
 use frame_support::{decl_module, decl_event, decl_error, decl_storage, Parameter, ensure, RuntimeDebug};
 use frame_support::{traits::{Get, ReservableCurrency, Currency},
 	weights::{GetDispatchInfo, DispatchClass,FunctionOf},
+	dispatch::PostDispatchInfo,
 };
 use frame_system::{self as system, ensure_signed};
 use sp_runtime::{DispatchError, DispatchResult, traits::Dispatchable};
@@ -83,7 +84,7 @@ pub trait Trait: frame_system::Trait {
 	type Event: From<Event<Self>> + Into<<Self as frame_system::Trait>::Event>;
 
 	/// The overarching call type.
-	type Call: Parameter + Dispatchable<Origin=Self::Origin> + GetDispatchInfo + From<frame_system::Call<Self>>;
+	type Call: Parameter + Dispatchable<Origin=Self::Origin, PostInfo=PostDispatchInfo> + GetDispatchInfo + From<frame_system::Call<Self>>;
 
 	/// The currency mechanism.
 	type Currency: ReservableCurrency<Self::AccountId>;
@@ -246,7 +247,7 @@ decl_module! {
 			for (index, call) in calls.into_iter().enumerate() {
 				let result = call.dispatch(origin.clone());
 				if let Err(e) = result {
-					Self::deposit_event(Event::<T>::BatchInterrupted(index as u32, e));
+					Self::deposit_event(Event::<T>::BatchInterrupted(index as u32, e.error));
 					return Ok(());
 				}
 			}
@@ -269,6 +270,7 @@ decl_module! {
 			let who = ensure_signed(origin)?;
 			let pseudonym = Self::sub_account_id(who, index);
 			call.dispatch(frame_system::RawOrigin::Signed(pseudonym).into())
+				.map(|_| ()).map_err(|e| e.error)
 		}
 
 		/// Register approval for a dispatch to be made from a deterministic composite account if
@@ -357,7 +359,9 @@ decl_module! {
 				let result = call.dispatch(frame_system::RawOrigin::Signed(id.clone()).into());
 				let _ = T::Currency::unreserve(&m.depositor, m.deposit);
 				<Multisigs<T>>::remove(&id, call_hash);
-				Self::deposit_event(RawEvent::MultisigExecuted(who, timepoint, id, result));
+				Self::deposit_event(RawEvent::MultisigExecuted(
+					who, timepoint, id, result.map(|_| ()).map_err(|e| e.error)
+				));
 			} else {
 				ensure!(maybe_timepoint.is_none(), Error::<T>::UnexpectedTimepoint);
 				if threshold > 1 {
@@ -373,6 +377,7 @@ decl_module! {
 					Self::deposit_event(RawEvent::NewMultisig(who, id));
 				} else {
 					return call.dispatch(frame_system::RawOrigin::Signed(id).into())
+						.map(|_| ()).map_err(|e| e.error)
 				}
 			}
 			Ok(())

primitives/runtime/src/generic/checked_extrinsic.rs

@@ -79,7 +79,7 @@ where
 			(None, pre)
 		};
 		let res = self.function.dispatch(Origin::from(maybe_who));
-		Extra::post_dispatch(pre, info.clone(), len);
-		Ok(res.map_err(Into::into))
+		Extra::post_dispatch(pre, info, len);
+		Ok(res.map(|_| ()).map_err(|e| e.error))
 	}
 }