add benchmarked extension weight data to runtime and update extension…

… weight

crates/pallet-subspace/src/extensions/mod.rs

@@ -4,21 +4,23 @@
 pub mod benchmarking;
 pub mod weights;
 
+use crate::extensions::weights::WeightInfo;
 use crate::pallet::Call as SubspaceCall;
 use crate::{Config, Pallet as Subspace};
 use codec::{Decode, Encode};
 use frame_support::pallet_prelude::{PhantomData, TypeInfo, Weight};
+use frame_support::RuntimeDebugNoBound;
 use frame_system::pallet_prelude::{BlockNumberFor, RuntimeCallFor};
 use scale_info::prelude::fmt;
 use sp_consensus_subspace::SignedVote;
-use sp_runtime::impl_tx_ext_default;
 use sp_runtime::traits::{
     AsSystemOriginSigner, DispatchInfoOf, DispatchOriginOf, Dispatchable, Implication,
-    TransactionExtension, ValidateResult,
+    PostDispatchInfoOf, TransactionExtension, ValidateResult,
 };
 use sp_runtime::transaction_validity::{
-    InvalidTransaction, TransactionSource, TransactionValidity, ValidTransaction,
+    InvalidTransaction, TransactionSource, TransactionValidityError, ValidTransaction,
 };
+use sp_runtime::DispatchResult;
 
 /// Trait to convert Runtime call to possible Subspace call.
 pub trait MaybeSubspaceCall<Runtime>
@@ -28,6 +30,15 @@ where
     fn maybe_subspace_call(&self) -> Option<&SubspaceCall<Runtime>>;
 }
 
+/// Weight info used by this extension
+#[derive(RuntimeDebugNoBound)]
+pub enum ExtensionWeightData {
+    /// Represents the validated call's used weight
+    Validated(Weight),
+    /// Skipped validation
+    Skipped,
+}
+
 /// Extensions for pallet-subspace unsigned extrinsics.
 #[derive(Encode, Decode, Clone, Eq, PartialEq, TypeInfo)]
 pub struct SubspaceExtension<Runtime>(PhantomData<Runtime>);
@@ -63,14 +74,20 @@ where
     fn do_check_vote(
         signed_vote: &SignedVote<BlockNumberFor<Runtime>, Runtime::Hash, Runtime::AccountId>,
         source: TransactionSource,
-    ) -> TransactionValidity {
+    ) -> Result<(ValidTransaction, Weight), TransactionValidityError> {
         let pre_dispatch = source == TransactionSource::InBlock;
         if pre_dispatch {
-            Subspace::<Runtime>::pre_dispatch_vote(signed_vote).map(|_| ValidTransaction::default())
+            Subspace::<Runtime>::pre_dispatch_vote(signed_vote)
+                .map(|weight| (ValidTransaction::default(), weight))
         } else {
             Subspace::<Runtime>::validate_vote(signed_vote)
         }
     }
+
+    fn max_weight() -> Weight {
+        <Runtime as Config>::ExtensionWeightInfo::vote()
+            .max(<Runtime as Config>::ExtensionWeightInfo::vote_with_equivocation())
+    }
 }
 
 impl<Runtime> TransactionExtension<RuntimeCallFor<Runtime>> for SubspaceExtension<Runtime>
@@ -82,12 +99,19 @@ where
 {
     const IDENTIFIER: &'static str = "SubspaceExtension";
     type Implicit = ();
-    type Val = ();
-    type Pre = ();
+    type Val = ExtensionWeightData;
+    type Pre = ExtensionWeightData;
 
-    // TODO: return correct weights
-    fn weight(&self, _call: &RuntimeCallFor<Runtime>) -> Weight {
-        Weight::zero()
+    fn weight(&self, call: &RuntimeCallFor<Runtime>) -> Weight {
+        let subspace_call = match call.maybe_subspace_call() {
+            Some(subspace_call) => subspace_call,
+            None => return Weight::zero(),
+        };
+
+        match subspace_call {
+            SubspaceCall::vote { .. } => Self::max_weight(),
+            _ => Weight::zero(),
+        }
     }
 
     fn validate(
@@ -102,22 +126,61 @@ where
     ) -> ValidateResult<Self::Val, RuntimeCallFor<Runtime>> {
         // we only care about unsigned calls
         if origin.as_system_origin_signer().is_some() {
-            return Ok((ValidTransaction::default(), (), origin));
+            return Ok((
+                ValidTransaction::default(),
+                ExtensionWeightData::Skipped,
+                origin,
+            ));
         };
 
         let subspace_call = match call.maybe_subspace_call() {
             Some(subspace_call) => subspace_call,
-            None => return Ok((ValidTransaction::default(), (), origin)),
+            None => {
+                return Ok((
+                    ValidTransaction::default(),
+                    ExtensionWeightData::Skipped,
+                    origin,
+                ))
+            }
         };
 
-        let validity = match subspace_call {
+        let (validity, weight_used) = match subspace_call {
             SubspaceCall::vote { signed_vote } => Self::do_check_vote(signed_vote, source)?,
             _ => return Err(InvalidTransaction::Call.into()),
         };
 
-        Ok((validity, (), origin))
+        Ok((
+            validity,
+            ExtensionWeightData::Validated(weight_used),
+            origin,
+        ))
     }
 
-    // nothing to prepare since vote is already checked
-    impl_tx_ext_default!(RuntimeCallFor<Runtime>; prepare);
+    fn prepare(
+        self,
+        val: Self::Val,
+        _origin: &DispatchOriginOf<RuntimeCallFor<Runtime>>,
+        _call: &RuntimeCallFor<Runtime>,
+        _info: &DispatchInfoOf<RuntimeCallFor<Runtime>>,
+        _len: usize,
+    ) -> Result<Self::Pre, TransactionValidityError> {
+        Ok(val)
+    }
+
+    fn post_dispatch_details(
+        pre: Self::Pre,
+        _info: &DispatchInfoOf<RuntimeCallFor<Runtime>>,
+        _post_info: &PostDispatchInfoOf<RuntimeCallFor<Runtime>>,
+        _len: usize,
+        _result: &DispatchResult,
+    ) -> Result<Weight, TransactionValidityError> {
+        match pre {
+            // return the unused weight for a validated call.
+            ExtensionWeightData::Validated(used_weight) => {
+                Ok(Self::max_weight().saturating_sub(used_weight))
+            }
+            // return no weight since this call is not validated and took no weight.
+            ExtensionWeightData::Skipped => Ok(Weight::zero()),
+        }
+    }
 }

crates/pallet-subspace/src/lib.rs

@@ -17,6 +17,7 @@ mod benchmarking;
 pub mod extensions;
 pub mod weights;
 
+use crate::extensions::weights::WeightInfo as ExtensionWeightInfo;
 #[cfg(not(feature = "std"))]
 use alloc::string::String;
 use codec::{Decode, Encode, MaxEncodedLen};
@@ -41,6 +42,7 @@ use sp_runtime::transaction_validity::{
     InvalidTransaction, TransactionPriority, TransactionSource, TransactionValidity,
     TransactionValidityError, ValidTransaction,
 };
+use sp_runtime::Weight;
 use sp_std::collections::btree_map::BTreeMap;
 use sp_std::prelude::*;
 use subspace_core_primitives::pieces::PieceOffset;
@@ -87,7 +89,7 @@ struct VoteVerificationData {
 
 #[frame_support::pallet]
 pub mod pallet {
-    use super::{EraChangeTrigger, VoteVerificationData};
+    use super::{EraChangeTrigger, ExtensionWeightInfo, VoteVerificationData};
     use crate::weights::WeightInfo;
     use frame_support::pallet_prelude::*;
     use frame_system::pallet_prelude::*;
@@ -225,6 +227,9 @@ pub mod pallet {
         /// Maximum number of block number to block slot mappings to keep (oldest pruned first).
         #[pallet::constant]
         type BlockSlotCount: Get<u32>;
+
+        /// Extension weight information for the pallet's extensions.
+        type ExtensionWeightInfo: ExtensionWeightInfo;
     }
 
     #[derive(Debug, Default, Encode, Decode, TypeInfo)]
@@ -1251,7 +1256,7 @@ impl<T: Config> Pallet<T> {
 
     fn validate_vote(
         signed_vote: &SignedVote<BlockNumberFor<T>, T::Hash, T::AccountId>,
-    ) -> TransactionValidity {
+    ) -> Result<(ValidTransaction, Weight), TransactionValidityError> {
         check_vote::<T>(signed_vote, false)?;
 
         ValidTransaction::with_tag_prefix("SubspaceVote")
@@ -1261,16 +1266,17 @@ impl<T: Config> Pallet<T> {
             .longevity(2)
             .and_provides(signed_vote.signature)
             .build()
+            .map(|validity| (validity, T::ExtensionWeightInfo::vote()))
     }
 
     fn pre_dispatch_vote(
         signed_vote: &SignedVote<BlockNumberFor<T>, T::Hash, T::AccountId>,
-    ) -> Result<(), TransactionValidityError> {
+    ) -> Result<Weight, TransactionValidityError> {
         match check_vote::<T>(signed_vote, true) {
-            Ok(()) => Ok(()),
+            Ok(()) => Ok(T::ExtensionWeightInfo::vote()),
             Err(CheckVoteError::Equivocated { .. }) => {
                 // Return Ok such that changes from this pre-dispatch are persisted
-                Ok(())
+                Ok(T::ExtensionWeightInfo::vote_with_equivocation())
             }
             Err(error) => Err(error.into()),
         }

crates/pallet-subspace/src/mock.rs

@@ -146,6 +146,7 @@ impl Config for Test {
     type BlockSlotCount = BlockSlotCount;
 
     type WeightInfo = ();
+    type ExtensionWeightInfo = crate::extensions::weights::SubstrateWeight<Test>;
 }
 
 pub fn go_to_block(

crates/subspace-runtime/src/lib.rs

@@ -302,6 +302,7 @@ impl pallet_subspace::Config for Runtime {
     type EraChangeTrigger = pallet_subspace::NormalEraChange;
     type BlockSlotCount = BlockSlotCount;
     type WeightInfo = pallet_subspace::weights::SubstrateWeight<Runtime>;
+    type ExtensionWeightInfo = pallet_subspace::extensions::weights::SubstrateWeight<Runtime>;
 }
 
 impl pallet_timestamp::Config for Runtime {

test/subspace-test-runtime/src/lib.rs

@@ -317,6 +317,7 @@ impl pallet_subspace::Config for Runtime {
     type EraChangeTrigger = pallet_subspace::NormalEraChange;
     type BlockSlotCount = BlockSlotCount;
     type WeightInfo = pallet_subspace::weights::SubstrateWeight<Runtime>;
+    type ExtensionWeightInfo = pallet_subspace::extensions::weights::SubstrateWeight<Runtime>;
 }
 
 impl pallet_timestamp::Config for Runtime {