ensure invest pools when vault withdraw from stake pools

Author: kingsleydon

pallets/phala/src/compute/stake_pool_v2.rs

@@ -780,13 +780,6 @@ pub mod pallet {
 				who.clone(),
 				pool_info.basepool.pid,
 			)?;
-			let nft_id = maybe_nft_id.ok_or(Error::<T>::NoNftToWithdraw)?;
-			// The nft instance must be wrote to Nft storage at the end of the function
-			// this nft's property shouldn't be accessed or wrote again from storage before set_nft_attr
-			// is called. Or the property of the nft will be overwrote incorrectly.
-			let mut nft_guard =
-				base_pool::Pallet::<T>::get_nft_attr_guard(pool_info.basepool.cid, nft_id)?;
-			let nft = &mut nft_guard.attr;
 			let in_queue_shares = match pool_info
 				.basepool
 				.withdraw_queue
@@ -803,10 +796,33 @@ pub mod pallet {
 				}
 				None => Zero::zero(),
 			};
+			ensure!(maybe_nft_id.is_some() || in_queue_shares > Zero::zero(), Error::<T>::NoNftToWithdraw);
+			let nft_id = match maybe_nft_id {
+				Some(nft_id) => nft_id,
+				// An nft is necessary to initiate a smaller withdrawal
+				None => base_pool::Pallet::<T>::mint_nft(
+					pool_info.basepool.cid,
+					who.clone(),
+					Zero::zero(),
+					pool_info.basepool.pid,
+				)?,
+			};
+			// The nft instance must be wrote to Nft storage at the end of the function
+			// this nft's property shouldn't be accessed or wrote again from storage before set_nft_attr
+			// is called. Or the property of the nft will be overwrote incorrectly.
+			let mut nft_guard =
+				base_pool::Pallet::<T>::get_nft_attr_guard(pool_info.basepool.cid, nft_id)?;
+			let nft = &mut nft_guard.attr;
 			ensure!(
 				base_pool::is_nondust_balance(shares) && (shares <= nft.shares + in_queue_shares),
 				Error::<T>::InvalidWithdrawalAmount
 			);
+			if let Some(vault_pid) = as_vault {
+				let mut vault_info = ensure_vault::<T>(vault_pid)?;
+				if !vault_info.invest_pools.contains(&pid) {
+					vault_info.invest_pools.push(pid);
+				}
+			}
 			base_pool::Pallet::<T>::try_withdraw(
 				&mut pool_info.basepool,
 				nft,

pallets/phala/src/compute/vault.rs

@@ -429,7 +429,7 @@ pub mod pallet {
 					if base_pool::is_nondust_balance(property.shares) {
 						true
 					} else {
-						if (!base_pool::is_nondust_balance(total_shares)) {
+						if !base_pool::is_nondust_balance(total_shares) {
 							let _ = base_pool::Pallet::<T>::burn_nft(
 								&base_pool::pallet_id::<T::AccountId>(),
 								stake_pool.basepool.cid,

pallets/phala/src/test.rs

@@ -2363,10 +2363,6 @@ fn vault_force_withdraw_with_dust_investment() {
 
         // Verify the vault is locked
         assert!(vault::VaultLocks::<Test>::contains_key(vault1));
-
-        // After force withdrawal, the invest_pools should be empty since the only pool had no shares
-        let vault = ensure_vault::<Test>(vault1).unwrap();
-        assert!(vault.invest_pools.is_empty(), "invest_pools should be cleaned up");
     });
 }