fix: adjust reward processing to occur after expired positions

pallets/bb-bnc/src/incentive.rs

@@ -312,25 +312,36 @@ impl<T: Config> Pallet<T> {
 		who: &AccountIdOf<T>,
 		expiry_block: BlockNumberFor<T>,
 	) -> DispatchResult {
+		let reward_block = expiry_block.saturating_sub(One::one());
+		Self::ensure_expiry_reward_snapshot(expiry_block)?;
+		let Some(reward_per_token_stored) = ExpiryRewardPerToken::<T>::get(expiry_block) else {
+			return Ok(());
+		};
+
+		let bbbnc_balance = Self::balance_of(who, Some(reward_block))?;
+		Self::update_account_reward(who, bbbnc_balance, reward_per_token_stored, None)
+	}
+
+	pub(crate) fn ensure_expiry_reward_snapshot(expiry_block: BlockNumberFor<T>) -> DispatchResult {
+		if ExpiryRewardPerToken::<T>::contains_key(expiry_block) {
+			return Ok(());
+		}
+
 		let reward_block = expiry_block.saturating_sub(One::one());
 		let conf = IncentiveConfigs::<T>::get(BB_BNC_SYSTEM_POOL_ID);
 		if conf.incentive_controller.is_none() || conf.reward_rate.is_empty() {
 			return Ok(());
 		}
-		let reward_per_token_stored = if let Some(reward_per_token_stored) =
-			ExpiryRewardPerToken::<T>::get(expiry_block)
-		{
-			reward_per_token_stored
-		} else if reward_block <= conf.last_update_time {
+
+		let reward_per_token_stored = if reward_block <= conf.last_update_time {
 			if conf.last_update_time > expiry_block {
 				log::warn!(
-					target: "bb-bnc::update_reward_before_expiry",
+					target: "bb-bnc::ensure_expiry_reward_snapshot",
 					"Missing reward-per-token snapshot for expired position settlement: expiry_block={expiry_block:?}, last_update_time={:?}",
 					conf.last_update_time,
 				);
 				return Ok(());
 			}
-			ExpiryRewardPerToken::<T>::insert(expiry_block, conf.reward_per_token_stored.clone());
 			conf.reward_per_token_stored
 		} else {
 			let reward_per_token_stored =
@@ -341,12 +352,11 @@ impl<T: Config> Pallet<T> {
 				item.reward_per_token_stored = reward_per_token_stored.clone();
 				item.last_update_time = reward_block.min(item.period_finish);
 			});
-			ExpiryRewardPerToken::<T>::insert(expiry_block, reward_per_token_stored.clone());
 			reward_per_token_stored
 		};
 
-		let bbbnc_balance = Self::balance_of(who, Some(reward_block))?;
-		Self::update_account_reward(who, bbbnc_balance, reward_per_token_stored, None)
+		ExpiryRewardPerToken::<T>::insert(expiry_block, reward_per_token_stored);
+		Ok(())
 	}
 
 	// Used to update reward when notify_reward or user call

pallets/bb-bnc/src/lib.rs

@@ -606,26 +606,6 @@ pub mod pallet {
 				T::BlockNumberProvider::current_block_number();
 			let mut weight = T::WeightInfo::on_initialize();
 
-			// Process existing rewards
-			let conf = IncentiveConfigs::<T>::get(BB_BNC_SYSTEM_POOL_ID);
-			if current_block_number == conf.period_finish {
-				if let Some(e) = Self::notify_reward_amount(
-					BB_BNC_SYSTEM_POOL_ID,
-					&conf.incentive_controller,
-					conf.last_reward.clone(),
-				)
-				.err()
-				{
-					log::error!(
-						target: "bb-bnc::notify_reward_amount",
-						"Received invalid justification for {e:?}",
-					);
-					Self::deposit_event(Event::NotifyRewardFailed {
-						rewards: conf.last_reward,
-					});
-				}
-			}
-
 			// Process expired positions
 			let next_expiring = NextExpiringBlock::<T>::get();
 
@@ -729,6 +709,37 @@ pub mod pallet {
 				// Update the next expiring block
 				NextExpiringBlock::<T>::set(next_block);
 				weight = weight.saturating_add(T::DbWeight::get().reads_writes(1, 1));
+
+				if blocks_exhausted && block <= current_block_number {
+					if let Err(e) = Self::ensure_expiry_reward_snapshot(block) {
+						log::warn!(
+							target: "bb-bnc::on_initialize",
+							"Failed to create reward snapshot for throttled expiry block {block:?}: {e:?}",
+						);
+					}
+					weight = weight.saturating_add(T::DbWeight::get().reads_writes(2, 1));
+				}
+			}
+
+			// Process existing rewards after expired positions so auto-withdraw can settle
+			// expiry rewards against the previous block before a period rollover advances state.
+			let conf = IncentiveConfigs::<T>::get(BB_BNC_SYSTEM_POOL_ID);
+			if current_block_number == conf.period_finish {
+				if let Some(e) = Self::notify_reward_amount(
+					BB_BNC_SYSTEM_POOL_ID,
+					&conf.incentive_controller,
+					conf.last_reward.clone(),
+				)
+				.err()
+				{
+					log::error!(
+						target: "bb-bnc::notify_reward_amount",
+						"Received invalid justification for {e:?}",
+					);
+					Self::deposit_event(Event::NotifyRewardFailed {
+						rewards: conf.last_reward,
+					});
+				}
 			}
 
 			weight

pallets/bb-bnc/src/tests.rs

@@ -2545,6 +2545,151 @@ fn update_reward_before_expiry_settles_multiple_users_at_same_expiry() {
 		});
 }
 
+#[test]
+fn auto_withdraw_uses_expiry_minus_one_when_reward_period_finishes_at_expiry() {
+	ExtBuilder::default()
+		.one_hundred_for_alice_n_bob()
+		.build()
+		.execute_with(|| {
+			asset_registry();
+			System::set_block_number(1);
+
+			assert_ok!(BbBNC::set_config(
+				RuntimeOrigin::root(),
+				Some(0),
+				Some(7 * DAYS),
+				Some(10)
+			));
+			assert_ok!(BbBNC::create_lock(
+				RuntimeOrigin::signed(BOB),
+				10_000_000_000_000,
+				7 * DAYS
+			));
+			let position = Position::<Runtime>::get() - 1;
+			let real_unlock_time = ((1 + 7 * DAYS) / Week::get() + 1) * Week::get();
+			assert_eq!(Locked::<Runtime>::get(position).end, real_unlock_time);
+
+			assert_ok!(Tokens::deposit(DOT, &ALICE, 10_000_000_000));
+			assert_ok!(BbBNC::notify_rewards(
+				RuntimeOrigin::root(),
+				ALICE,
+				Some(real_unlock_time - 1),
+				vec![DOT]
+			));
+			assert_eq!(
+				IncentiveConfigs::<Runtime>::get(BB_BNC_SYSTEM_POOL_ID).period_finish,
+				real_unlock_time
+			);
+
+			System::set_block_number(real_unlock_time - 1);
+			let expected_reward = BbBNC::query_pending_rewards(&BOB)
+				.unwrap()
+				.into_iter()
+				.find(|(currency_id, _)| *currency_id == DOT)
+				.map(|(_, amount)| amount)
+				.unwrap_or_default();
+			assert!(expected_reward > 0);
+
+			NextExpiringBlock::<Runtime>::set(real_unlock_time);
+			System::set_block_number(real_unlock_time);
+			BbBNC::on_initialize(real_unlock_time);
+
+			assert_eq!(Locked::<Runtime>::get(position).amount, 0);
+			assert_eq!(
+				Rewards::<Runtime>::get(BOB)
+					.unwrap()
+					.get(&DOT)
+					.copied()
+					.unwrap_or_default(),
+				expected_reward
+			);
+		});
+}
+
+#[test]
+fn throttled_auto_withdraw_after_period_rollover_uses_expiry_minus_one() {
+	ExtBuilder::default()
+		.one_hundred_for_alice_n_bob()
+		.build()
+		.execute_with(|| {
+			asset_registry();
+			System::set_block_number(1);
+
+			assert_ok!(BbBNC::set_config(
+				RuntimeOrigin::root(),
+				Some(0),
+				Some(7 * DAYS),
+				Some(1)
+			));
+			assert_ok!(BbBNC::create_lock(
+				RuntimeOrigin::signed(BOB),
+				10_000_000_000_000,
+				7 * DAYS
+			));
+			assert_ok!(BbBNC::create_lock(
+				RuntimeOrigin::signed(CHARLIE),
+				10_000_000_000_000,
+				7 * DAYS
+			));
+			let real_unlock_time = ((1 + 7 * DAYS) / Week::get() + 1) * Week::get();
+			let ordered_positions = expiring_positions(real_unlock_time);
+			assert_eq!(ordered_positions.len(), 2);
+			let orphaned_position = ordered_positions[0];
+			let delayed_position = ordered_positions[1];
+			let delayed_owner =
+				PositionOwner::<Runtime>::get(delayed_position).expect("owner should exist");
+			assert_eq!(
+				Locked::<Runtime>::get(delayed_position).end,
+				real_unlock_time
+			);
+
+			assert_ok!(Tokens::deposit(DOT, &ALICE, 10_000_000_000));
+			assert_ok!(BbBNC::notify_rewards(
+				RuntimeOrigin::root(),
+				ALICE,
+				Some(real_unlock_time - 1),
+				vec![DOT]
+			));
+			assert_eq!(
+				IncentiveConfigs::<Runtime>::get(BB_BNC_SYSTEM_POOL_ID).period_finish,
+				real_unlock_time
+			);
+
+			System::set_block_number(real_unlock_time - 1);
+			let expected_reward = BbBNC::query_pending_rewards(&delayed_owner)
+				.unwrap()
+				.into_iter()
+				.find(|(currency_id, _)| *currency_id == DOT)
+				.map(|(_, amount)| amount)
+				.unwrap_or_default();
+			assert!(expected_reward > 0);
+
+			PositionOwner::<Runtime>::remove(orphaned_position);
+			NextExpiringBlock::<Runtime>::set(real_unlock_time);
+			System::set_block_number(real_unlock_time);
+			BbBNC::on_initialize(real_unlock_time);
+
+			assert!(expiring_contains(real_unlock_time, delayed_position));
+			assert!(ExpiryRewardPerToken::<Runtime>::contains_key(
+				real_unlock_time
+			));
+			assert_eq!(NextExpiringBlock::<Runtime>::get(), real_unlock_time);
+
+			System::set_block_number(real_unlock_time + 1);
+			BbBNC::on_initialize(real_unlock_time + 1);
+
+			assert_eq!(Locked::<Runtime>::get(delayed_position).amount, 0);
+			assert_eq!(
+				Rewards::<Runtime>::get(delayed_owner)
+					.unwrap()
+					.get(&DOT)
+					.copied()
+					.unwrap_or_default(),
+				expected_reward
+			);
+		});
+}
+
 #[test]
 fn test_multiple_expired_positions() {
 	ExtBuilder::default()