Separate the logic for retrieving the current block for different tok…

pallets/vtoken-voting/src/agents/bifrost_agent/agent.rs

@@ -147,4 +147,8 @@ impl<T: Config> VotingAgent<T> for BifrostAgent<T> {
 		Ok(<BifrostCall<T> as ConvictionVotingCall<T>>::remove_vote(Some(class), poll_index)
 			.encode())
 	}
+
+	fn block_number(&self) -> BlockNumberFor<T> {
+		T::LocalBlockNumberProvider::current_block_number()
+	}
 }

pallets/vtoken-voting/src/agents/relaychain_agent/agent.rs

@@ -117,4 +117,8 @@ impl<T: Config> VotingAgent<T> for RelaychainAgent<T> {
 		)
 		.encode())
 	}
+
+	fn block_number(&self) -> BlockNumberFor<T> {
+		BlockNumberFor::<T>::from(T::RelaychainBlockNumberProvider::current_block_number())
+	}
 }

pallets/vtoken-voting/src/lib.rs

@@ -71,6 +71,7 @@ pub use weights::WeightInfo;
 use xcm::v4::{prelude::*, Location, Weight as XcmWeight};
 
 const CONVICTION_VOTING_ID: LockIdentifier = *b"vtvoting";
+const SUPPORTED_VTOKENS: &[CurrencyId] = &[VKSM, VDOT, VBNC];
 
 type PollIndex = u32;
 type PollClass = u16;
@@ -156,6 +157,8 @@ pub mod pallet {
 		type WeightInfo: WeightInfo;
 
 		type PalletsOrigin: CallerTrait<Self::AccountId>;
+
+		type LocalBlockNumberProvider: BlockNumberProvider<BlockNumber = BlockNumberFor<Self>>;
 	}
 
 	#[pallet::event]
@@ -417,7 +420,7 @@ pub mod pallet {
 		_,
 		Twox64Concat,
 		BlockNumberFor<T>,
-		BoundedVec<(CurrencyIdOf<T>, PollIndex), ConstU32<50>>,
+		BoundedVec<(CurrencyIdOf<T>, PollIndex), ConstU32<100>>,
 		ValueQuery,
 	>;
 
@@ -459,39 +462,56 @@ pub mod pallet {
 
 	#[pallet::hooks]
 	impl<T: Config> Hooks<BlockNumberFor<T>> for Pallet<T> {
-		fn on_idle(n: BlockNumberFor<T>, remaining_weight: Weight) -> Weight {
+		fn on_idle(
+			bifrost_current_block_number: BlockNumberFor<T>,
+			remaining_weight: Weight,
+		) -> Weight {
 			let db_weight = T::DbWeight::get();
 			let mut used_weight = db_weight.reads(3);
 			if remaining_weight.any_lt(used_weight) ||
-				n % T::ReferendumCheckInterval::get() != Zero::zero()
+				bifrost_current_block_number % T::ReferendumCheckInterval::get() != Zero::zero()
 			{
 				return Weight::zero();
 			}
 			let relay_current_block_number =
 				T::RelaychainBlockNumberProvider::current_block_number();
 
-			for relay_block_number in ReferendumTimeout::<T>::iter_keys() {
-				if relay_current_block_number >= relay_block_number {
-					let info_list = ReferendumTimeout::<T>::get(relay_block_number);
+			for block_number in ReferendumTimeout::<T>::iter_keys() {
+				if relay_current_block_number >= block_number ||
+					bifrost_current_block_number >= block_number
+				{
+					let info_list = ReferendumTimeout::<T>::get(block_number);
 					let len = info_list.len() as u64;
 					let temp_weight = db_weight.reads_writes(len, len) + db_weight.writes(1);
 					if remaining_weight.any_lt(used_weight + temp_weight) {
 						return used_weight;
 					}
 					used_weight += temp_weight;
 					for (vtoken, poll_index) in info_list.iter() {
+						let current_block_number = match Self::get_agent_block_number(&vtoken) {
+							Ok(block_number) => block_number,
+							Err(e) => {
+								log::error!(
+									"Failed to get {:?} current block number: {:?}",
+									vtoken,
+									e
+								);
+								// exit the current loop
+								continue;
+							},
+						};
+
 						ReferendumInfoFor::<T>::mutate(vtoken, poll_index, |maybe_info| {
 							match maybe_info {
 								Some(info) =>
 									if let ReferendumInfo::Ongoing(_) = info {
-										*info =
-											ReferendumInfo::Completed(relay_current_block_number);
+										*info = ReferendumInfo::Completed(current_block_number);
 									},
 								None => {},
 							}
 						});
 					}
-					ReferendumTimeout::<T>::remove(relay_block_number);
+					ReferendumTimeout::<T>::remove(block_number);
 				}
 			}
 
@@ -637,10 +657,12 @@ pub mod pallet {
 			Self::ensure_vtoken(&vtoken)?;
 			Self::ensure_referendum_completed(vtoken, poll_index)?;
 
+			let current_block_number = Self::get_agent_block_number(&vtoken)?;
+
 			ReferendumInfoFor::<T>::insert(
 				vtoken,
 				poll_index,
-				ReferendumInfo::Killed(T::RelaychainBlockNumberProvider::current_block_number()),
+				ReferendumInfo::Killed(current_block_number),
 			);
 
 			Self::deposit_event(Event::<T>::ReferendumKilled { vtoken, poll_index });
@@ -752,41 +774,7 @@ pub mod pallet {
 				Self::deposit_event(Event::<T>::VoteNotified { vtoken, poll_index, success });
 			}
 
-			if let Some((vtoken, poll_index)) = PendingReferendumInfo::<T>::get(query_id) {
-				if success {
-					ReferendumInfoFor::<T>::try_mutate_exists(
-						vtoken,
-						poll_index,
-						|maybe_info| -> DispatchResult {
-							if let Some(info) = maybe_info {
-								if let ReferendumInfo::Ongoing(status) = info {
-									let relay_current_block_number =
-										T::RelaychainBlockNumberProvider::current_block_number();
-									status.submitted = Some(relay_current_block_number);
-									ReferendumTimeout::<T>::mutate(
-										relay_current_block_number.saturating_add(
-											UndecidingTimeout::<T>::get(vtoken)
-												.ok_or(Error::<T>::NoData)?,
-										),
-										|ref_vec| {
-											ref_vec
-												.try_push((vtoken, poll_index))
-												.map_err(|_| Error::<T>::TooMany)
-										},
-									)?;
-									Self::deposit_event(Event::<T>::ReferendumInfoCreated {
-										vtoken,
-										poll_index,
-										info: info.clone(),
-									});
-								}
-							}
-							Ok(())
-						},
-					)?;
-				} else {
-					ReferendumInfoFor::<T>::remove(vtoken, poll_index);
-				}
+			if let Some((_, _)) = PendingReferendumInfo::<T>::get(query_id) {
 				PendingReferendumInfo::<T>::remove(query_id);
 			}
 
@@ -862,6 +850,7 @@ pub mod pallet {
 				if let Some((old_vote, vtoken_balance)) = maybe_old_vote {
 					Self::try_vote(&who, vtoken, poll_index, old_vote, vtoken_balance)?;
 				}
+				ReferendumInfoFor::<T>::remove(vtoken, poll_index);
 			} else {
 				if !VoteDelegatorFor::<T>::contains_key((&who, vtoken, poll_index)) {
 					VoteDelegatorFor::<T>::insert((&who, vtoken, poll_index), derivative_index);
@@ -876,6 +865,36 @@ pub mod pallet {
 					}
 					Ok(())
 				})?;
+
+				ReferendumInfoFor::<T>::try_mutate_exists(
+					vtoken,
+					poll_index,
+					|maybe_info| -> DispatchResult {
+						if let Some(info) = maybe_info {
+							if let ReferendumInfo::Ongoing(status) = info {
+								let current_block_number = Self::get_agent_block_number(&vtoken)?;
+								status.submitted = Some(current_block_number);
+								ReferendumTimeout::<T>::mutate(
+									current_block_number.saturating_add(
+										UndecidingTimeout::<T>::get(vtoken)
+											.ok_or(Error::<T>::NoData)?,
+									),
+									|ref_vec| {
+										ref_vec
+											.try_push((vtoken, poll_index))
+											.map_err(|_| Error::<T>::TooMany)
+									},
+								)?;
+								Self::deposit_event(Event::<T>::ReferendumInfoCreated {
+									vtoken,
+									poll_index,
+									info: info.clone(),
+								});
+							}
+						}
+						Ok(())
+					},
+				)?;
 			}
 
 			Ok(())
@@ -1211,7 +1230,7 @@ pub mod pallet {
 		}
 
 		fn ensure_vtoken(vtoken: &CurrencyIdOf<T>) -> Result<(), DispatchError> {
-			ensure!([VKSM, VDOT, VBNC].contains(vtoken), Error::<T>::VTokenNotSupport);
+			ensure!(SUPPORTED_VTOKENS.contains(vtoken), Error::<T>::VTokenNotSupport);
 			Ok(())
 		}
 
@@ -1439,6 +1458,13 @@ pub mod pallet {
 			}
 		}
 
+		pub(crate) fn get_agent_block_number(
+			currency_id: &CurrencyIdOf<T>,
+		) -> Result<BlockNumberFor<T>, Error<T>> {
+			let voting_agent = Self::get_voting_agent(&currency_id)?;
+			Ok(voting_agent.block_number())
+		}
+
 		pub(crate) fn convert_vtoken_to_dest_location(
 			vtoken: CurrencyId,
 		) -> Result<Location, Error<T>> {

pallets/vtoken-voting/src/mock.rs

@@ -442,6 +442,7 @@ impl vtoken_voting::Config for Runtime {
 	type ReferendumCheckInterval = ReferendumCheckInterval;
 	type WeightInfo = ();
 	type PalletsOrigin = OriginCaller;
+	type LocalBlockNumberProvider = System;
 }
 
 pub fn new_test_ext() -> sp_io::TestExternalities {

pallets/vtoken-voting/src/tests/common_test.rs

@@ -896,7 +896,7 @@ fn notify_vote_success_max_works() {
 fn notify_vote_success_exceed_max_fail() {
 	for &vtoken in TOKENS {
 		new_test_ext().execute_with(|| {
-			for poll_index in 0..50 {
+			for poll_index in 0..100 {
 				assert_ok!(VtokenVoting::vote(
 					RuntimeOrigin::signed(ALICE),
 					vtoken,
@@ -909,7 +909,7 @@ fn notify_vote_success_exceed_max_fail() {
 					response_success()
 				));
 			}
-			let poll_index = 50;
+			let poll_index = 100;
 			assert_ok!(VtokenVoting::vote(
 				RuntimeOrigin::signed(ALICE),
 				vtoken,

pallets/vtoken-voting/src/tests/vbnc_test.rs

@@ -507,7 +507,7 @@ fn errors_with_vote_works() {
 				Error::<Runtime>::InsufficientFunds
 			);
 
-			for poll_index in 0..256 {
+			for poll_index in 0..100 {
 				assert_ok!(VtokenVoting::vote(
 					RuntimeOrigin::signed(1),
 					vtoken,
@@ -516,8 +516,8 @@ fn errors_with_vote_works() {
 				));
 			}
 			assert_noop!(
-				VtokenVoting::vote(RuntimeOrigin::signed(1), vtoken, 256, aye(10, 0)),
-				Error::<Runtime>::MaxVotesReached
+				VtokenVoting::vote(RuntimeOrigin::signed(1), vtoken, 100, aye(10, 0)),
+				Error::<Runtime>::TooMany
 			);
 		});
 	}

pallets/vtoken-voting/src/traits.rs

@@ -108,6 +108,9 @@ pub trait VotingAgent<T: Config> {
 		poll_index: PollIndex,
 		derivative_index: DerivativeIndex,
 	) -> Result<Vec<u8>, Error<T>>;
+
+	/// Get current agent block number.
+	fn block_number(&self) -> BlockNumberFor<T>;
 }
 
 /// Trait defining calls related to conviction voting mechanisms.

runtime/bifrost-kusama/src/lib.rs

@@ -1434,6 +1434,7 @@ impl bifrost_vtoken_voting::Config for Runtime {
 	type ReferendumCheckInterval = ReferendumCheckInterval;
 	type WeightInfo = weights::bifrost_vtoken_voting::BifrostWeight<Runtime>;
 	type PalletsOrigin = OriginCaller;
+	type LocalBlockNumberProvider = System;
 }
 
 // Bifrost modules end

runtime/bifrost-polkadot/src/lib.rs

@@ -1334,6 +1334,7 @@ impl bifrost_vtoken_voting::Config for Runtime {
 	type ReferendumCheckInterval = ReferendumCheckInterval;
 	type WeightInfo = weights::bifrost_vtoken_voting::BifrostWeight<Runtime>;
 	type PalletsOrigin = OriginCaller;
+	type LocalBlockNumberProvider = System;
 }
 
 // Bifrost modules end