Voting hook (#7703)

# Description
This is a continuation of the following
[PR](#5735 (comment))

---------

Co-authored-by: Bastian Köcher <[email protected]>

Author: ndkazu

polkadot/runtime/rococo/src/governance/mod.rs

@@ -48,6 +48,7 @@ impl pallet_conviction_voting::Config for Runtime {
 		frame_support::traits::tokens::currency::ActiveIssuanceOf<Balances, Self::AccountId>;
 	type Polls = Referenda;
 	type BlockNumberProvider = System;
+	type VotingHooks = ();
 }
 
 parameter_types! {

polkadot/runtime/westend/src/governance/mod.rs

@@ -45,6 +45,7 @@ impl pallet_conviction_voting::Config for Runtime {
 		frame_support::traits::tokens::currency::ActiveIssuanceOf<Balances, Self::AccountId>;
 	type Polls = Referenda;
 	type BlockNumberProvider = System;
+	type VotingHooks = ();
 }
 
 parameter_types! {

prdoc/pr_7703.prdoc

@@ -0,0 +1,37 @@
+title: 'Add voting hooks to Conviction_Voting'
+doc:
+  - audience: Runtime Dev
+    description: |
+      This change introduces voting hooks to the conviction-voting pallet, enabling developers to 
+      customize behavior during various stages of the voting process. These hooks provide a mechanism 
+      to execute specific logic before a vote is recorded, before a vote is removed, or when a vote 
+      fails to be recorded, while maintaining compatibility with the existing conviction-voting pallet.
+
+      The key hooks include:
+      - `on_vote`: Called before a vote is recorded. This hook allows developers to validate or 
+      perform actions based on the vote. If it returns an error, the voting operation is reverted. 
+      However, any storage modifications made by this hook will persist even if the vote fails later.
+      - `on_remove_vote`: Called before a vote is removed. This hook cannot fail and is useful for 
+      cleanup or additional logic when a vote is removed.
+      - `lock_balance_on_unsuccessful_vote`: Called when a vote fails to be recorded, such as due to 
+      insufficient balance. It allows locking a specific balance amount as part of the failure handling.
+
+      Advantages of using voting hooks:
+      - Flexibility: Developers can implement custom logic to extend or modify the behavior of the 
+      conviction-voting pallet.
+      - Control: Hooks provide fine-grained control over different stages of the voting process.
+      - Error Handling: The `on_vote` hook enables early validation, preventing (to some extent) invalid votes from 
+      being recorded.
+
+      How to use:
+      - Implement the `VotingHooks` trait in your runtime or custom module.
+      - Define the desired behavior for each hook method, such as validation logic in `on_vote` or 
+      cleanup actions in `on_remove_vote`.
+      - Integrate the implementation with the conviction-voting pallet to enable the hooks.
+crates: 
+- name: pallet-conviction-voting
+  bump: major
+- name: rococo-runtime
+  bump: minor
+- name: westend-runtime
+  bump: minor

substrate/bin/node/runtime/src/lib.rs

@@ -1015,6 +1015,7 @@ impl pallet_conviction_voting::Config for Runtime {
 	type MaxTurnout = frame_support::traits::TotalIssuanceOf<Balances, Self::AccountId>;
 	type Polls = Referenda;
 	type BlockNumberProvider = System;
+	type VotingHooks = ();
 }
 
 parameter_types! {

substrate/frame/conviction-voting/src/benchmarking.rs

@@ -73,6 +73,8 @@ benchmarks_instance_pallet! {
 		whitelist_account!(caller);
 		let account_vote = account_vote::<T, I>(100u32.into());
 
+		T::VotingHooks::on_vote_worst_case(&caller);
+
 		let (class, all_polls) = fill_voting::<T, I>();
 		let polls = &all_polls[&class];
 		let r = polls.len() - 1;
@@ -100,6 +102,8 @@ benchmarks_instance_pallet! {
 		whitelist_account!(caller);
 		let old_account_vote = account_vote::<T, I>(100u32.into());
 
+		T::VotingHooks::on_vote_worst_case(&caller);
+
 		let (class, all_polls) = fill_voting::<T, I>();
 		let polls = &all_polls[&class];
 		let r = polls.len();
@@ -128,6 +132,8 @@ benchmarks_instance_pallet! {
 		whitelist_account!(caller);
 		let old_account_vote = account_vote::<T, I>(100u32.into());
 
+		T::VotingHooks::on_vote_worst_case(&caller);
+
 		let (class, all_polls) = fill_voting::<T, I>();
 		let polls = &all_polls[&class];
 		let r = polls.len();
@@ -157,6 +163,8 @@ benchmarks_instance_pallet! {
 		whitelist_account!(caller);
 		let old_account_vote = account_vote::<T, I>(100u32.into());
 
+		T::VotingHooks::on_vote_worst_case(&caller);
+
 		let (class, all_polls) = fill_voting::<T, I>();
 		let polls = &all_polls[&class];
 		let r = polls.len();

substrate/frame/conviction-voting/src/lib.rs

@@ -43,13 +43,15 @@ use sp_runtime::{
 };
 
 mod conviction;
+mod traits;
 mod types;
 mod vote;
 pub mod weights;
 
 pub use self::{
 	conviction::Conviction,
 	pallet::*,
+	traits::{Status, VotingHooks},
 	types::{Delegations, Tally, UnvoteScope},
 	vote::{AccountVote, Casting, Delegating, Vote, Voting},
 	weights::WeightInfo,
@@ -142,6 +144,18 @@ pub mod pallet {
 		type VoteLockingPeriod: Get<BlockNumberFor<Self, I>>;
 		/// Provider for the block number. Normally this is the `frame_system` pallet.
 		type BlockNumberProvider: BlockNumberProvider;
+		/// Hooks are called when a new vote is registered or an existing vote is removed.
+		///
+		/// The trait does not expose weight information.
+		/// The weight of each hook is assumed to be benchmarked as part of the function that calls
+		/// it. Hooks should never recursively call into functions that called,
+		/// directly or indirectly, the function that called them.
+		/// This could lead to infinite recursion and stack overflow.
+		/// Note that this also means to not call into other generic functionality like batch or
+		/// similar. Also, anything that a hook did will be subject to the transactional semantics
+		/// of the calling function. This means that if the calling function fails, the hook will
+		/// be rolled back without further notice.
+		type VotingHooks: VotingHooks<Self::AccountId, PollIndexOf<Self, I>, BalanceOf<Self, I>>;
 	}
 
 	/// All voting for a particular voter in a particular voting class. We store the balance for the
@@ -410,6 +424,9 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 			vote.balance() <= T::Currency::total_balance(who),
 			Error::<T, I>::InsufficientFunds
 		);
+		// Call on_vote hook
+		T::VotingHooks::on_before_vote(who, poll_index, vote)?;
+
 		T::Polls::try_access_poll(poll_index, |poll_status| {
 			let (tally, class) = poll_status.ensure_ongoing().ok_or(Error::<T, I>::NotOngoing)?;
 			VotingFor::<T, I>::try_mutate(who, &class, |voting| {
@@ -435,7 +452,7 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 						tally.increase(approve, *delegations);
 					}
 				} else {
-					return Err(Error::<T, I>::AlreadyDelegating.into())
+					return Err(Error::<T, I>::AlreadyDelegating.into());
 				}
 				// Extend the lock to `balance` (rather than setting it) since we don't know what
 				// other votes are in place.
@@ -477,10 +494,13 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 							tally.reduce(approve, *delegations);
 						}
 						Self::deposit_event(Event::VoteRemoved { who: who.clone(), vote: v.1 });
+						T::VotingHooks::on_remove_vote(who, poll_index, Status::Ongoing);
 						Ok(())
 					},
 					PollStatus::Completed(end, approved) => {
-						if let Some((lock_periods, balance)) = v.1.locked_if(approved) {
+						if let Some((lock_periods, balance)) =
+							v.1.locked_if(vote::LockedIf::Status(approved))
+						{
 							let unlock_at = end.saturating_add(
 								T::VoteLockingPeriod::get().saturating_mul(lock_periods.into()),
 							);
@@ -492,10 +512,37 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 								);
 								prior.accumulate(unlock_at, balance)
 							}
+						} else if v.1.as_standard().is_some_and(|vote| vote != approved) {
+							// Unsuccessful vote, use special hook to lock the funds too in case of
+							// conviction.
+							if let Some(to_lock) =
+								T::VotingHooks::lock_balance_on_unsuccessful_vote(who, poll_index)
+							{
+								if let AccountVote::Standard { vote, .. } = v.1 {
+									let unlock_at = end.saturating_add(
+										T::VoteLockingPeriod::get()
+											.saturating_mul(vote.conviction.lock_periods().into()),
+									);
+									let now = T::BlockNumberProvider::current_block_number();
+									if now < unlock_at {
+										ensure!(
+											matches!(scope, UnvoteScope::Any),
+											Error::<T, I>::NoPermissionYet
+										);
+										prior.accumulate(unlock_at, to_lock)
+									}
+								}
+							}
 						}
+						// Call on_remove_vote hook
+						T::VotingHooks::on_remove_vote(who, poll_index, Status::Completed);
+						Ok(())
+					},
+					PollStatus::None => {
+						// Poll was cancelled.
+						T::VotingHooks::on_remove_vote(who, poll_index, Status::None);
 						Ok(())
 					},
-					PollStatus::None => Ok(()), // Poll was cancelled.
 				})
 			} else {
 				Ok(())

substrate/frame/conviction-voting/src/tests.rs

@@ -17,7 +17,7 @@
 
 //! The crate's tests.
 
-use std::collections::BTreeMap;
+use std::{cell::RefCell, collections::BTreeMap};
 
 use frame_support::{
 	assert_noop, assert_ok, derive_impl, parameter_types,
@@ -155,6 +155,7 @@ impl Config for Test {
 	type MaxTurnout = frame_support::traits::TotalIssuanceOf<Balances, Self::AccountId>;
 	type Polls = TestPolls;
 	type BlockNumberProvider = System;
+	type VotingHooks = HooksHandler;
 }
 
 pub fn new_test_ext() -> sp_io::TestExternalities {
@@ -913,3 +914,140 @@ fn errors_with_remove_vote_work() {
 		);
 	});
 }
+
+thread_local! {
+	static LAST_ON_VOTE_DATA: RefCell<Option<(u64, u8, AccountVote<u64>)>> = RefCell::new(None);
+	static LAST_ON_REMOVE_VOTE_DATA: RefCell<Option<(u64, u8, Status)>> = RefCell::new(None);
+	static LAST_LOCKED_IF_UNSUCCESSFUL_VOTE_DATA: RefCell<Option<(u64, u8)>> = RefCell::new(None);
+	static REMOVE_VOTE_LOCKED_AMOUNT: RefCell<Option<u64>> = RefCell::new(None);
+}
+
+pub struct HooksHandler;
+
+impl HooksHandler {
+	fn last_on_vote_data() -> Option<(u64, u8, AccountVote<u64>)> {
+		LAST_ON_VOTE_DATA.with(|data| *data.borrow())
+	}
+
+	fn last_on_remove_vote_data() -> Option<(u64, u8, Status)> {
+		LAST_ON_REMOVE_VOTE_DATA.with(|data| *data.borrow())
+	}
+
+	fn last_locked_if_unsuccessful_vote_data() -> Option<(u64, u8)> {
+		LAST_LOCKED_IF_UNSUCCESSFUL_VOTE_DATA.with(|data| *data.borrow())
+	}
+
+	fn reset() {
+		LAST_ON_VOTE_DATA.with(|data| *data.borrow_mut() = None);
+		LAST_ON_REMOVE_VOTE_DATA.with(|data| *data.borrow_mut() = None);
+		LAST_LOCKED_IF_UNSUCCESSFUL_VOTE_DATA.with(|data| *data.borrow_mut() = None);
+		REMOVE_VOTE_LOCKED_AMOUNT.with(|data| *data.borrow_mut() = None);
+	}
+
+	fn with_remove_locked_amount(v: u64) {
+		REMOVE_VOTE_LOCKED_AMOUNT.with(|data| *data.borrow_mut() = Some(v));
+	}
+}
+
+impl VotingHooks<u64, u8, u64> for HooksHandler {
+	fn on_before_vote(who: &u64, ref_index: u8, vote: AccountVote<u64>) -> DispatchResult {
+		LAST_ON_VOTE_DATA.with(|data| {
+			*data.borrow_mut() = Some((*who, ref_index, vote));
+		});
+		Ok(())
+	}
+
+	fn on_remove_vote(who: &u64, ref_index: u8, ongoing: Status) {
+		LAST_ON_REMOVE_VOTE_DATA.with(|data| {
+			*data.borrow_mut() = Some((*who, ref_index, ongoing));
+		});
+	}
+
+	fn lock_balance_on_unsuccessful_vote(who: &u64, ref_index: u8) -> Option<u64> {
+		LAST_LOCKED_IF_UNSUCCESSFUL_VOTE_DATA.with(|data| {
+			*data.borrow_mut() = Some((*who, ref_index));
+
+			REMOVE_VOTE_LOCKED_AMOUNT.with(|data| *data.borrow())
+		})
+	}
+
+	#[cfg(feature = "runtime-benchmarks")]
+	fn on_vote_worst_case(_who: &u64) {}
+
+	#[cfg(feature = "runtime-benchmarks")]
+	fn on_remove_vote_worst_case(_who: &u64) {}
+}
+
+#[test]
+fn voting_hooks_are_called_correctly() {
+	new_test_ext().execute_with(|| {
+		let c = class(3);
+
+		let usable_balance_1 = Balances::usable_balance(1);
+		dbg!(usable_balance_1);
+
+		// Voting
+		assert_ok!(Voting::vote(RuntimeOrigin::signed(1), 3, aye(1, 1)));
+		assert_eq!(
+			HooksHandler::last_on_vote_data(),
+			Some((
+				1,
+				3,
+				AccountVote::Standard {
+					vote: Vote { aye: true, conviction: Conviction::Locked1x },
+					balance: 1
+				}
+			))
+		);
+		assert_ok!(Voting::vote(RuntimeOrigin::signed(2), 3, nay(20, 2)));
+		assert_eq!(
+			HooksHandler::last_on_vote_data(),
+			Some((
+				2,
+				3,
+				AccountVote::Standard {
+					vote: Vote { aye: false, conviction: Conviction::Locked2x },
+					balance: 20
+				}
+			))
+		);
+		HooksHandler::reset();
+
+		// removing vote while ongoing
+		assert_ok!(Voting::vote(RuntimeOrigin::signed(3), 3, nay(20, 0)));
+		assert_eq!(
+			HooksHandler::last_on_vote_data(),
+			Some((
+				3,
+				3,
+				AccountVote::Standard {
+					vote: Vote { aye: false, conviction: Conviction::None },
+					balance: 20
+				}
+			))
+		);
+		assert_ok!(Voting::remove_vote(RuntimeOrigin::signed(3), Some(c), 3));
+		assert_eq!(HooksHandler::last_on_remove_vote_data(), Some((3, 3, Status::Ongoing)));
+		HooksHandler::reset();
+
+		Polls::set(vec![(3, Completed(3, false))].into_iter().collect());
+
+		// removing successful vote while completed
+		assert_ok!(Voting::remove_vote(RuntimeOrigin::signed(2), Some(c), 3));
+		assert_eq!(HooksHandler::last_on_remove_vote_data(), Some((2, 3, Status::Completed)));
+		assert_eq!(HooksHandler::last_locked_if_unsuccessful_vote_data(), None);
+
+		HooksHandler::reset();
+
+		HooksHandler::with_remove_locked_amount(5);
+
+		// removing unsuccessful vote when completed
+		assert_ok!(Voting::remove_vote(RuntimeOrigin::signed(1), Some(c), 3));
+		assert_eq!(HooksHandler::last_on_remove_vote_data(), Some((1, 3, Status::Completed)));
+		assert_eq!(HooksHandler::last_locked_if_unsuccessful_vote_data(), Some((1, 3)));
+
+		// Removing unsuccessful vote when completed should lock if given amount from the hook
+		assert_ok!(Voting::unlock(RuntimeOrigin::signed(1), c, 1));
+		assert_eq!(Balances::usable_balance(1), 5);
+	});
+}