fix: retain stake pool with withdraw in invest pools

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

@@ -425,23 +425,25 @@ pub mod pallet {
 						nft_id,
 					).expect("get nft should not fail: qed.");
 					let property = &property_guard.attr;
-					if !base_pool::is_nondust_balance(property.shares) {
-						let _ = base_pool::Pallet::<T>::burn_nft(
-							&base_pool::pallet_id::<T::AccountId>(),
-							stake_pool.basepool.cid,
-							nft_id,
-						);
-						false
+					total_shares += property.shares;
+					if base_pool::is_nondust_balance(property.shares) {
+						true
 					} else {
-						total_shares += property.shares;
-						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,
+								nft_id,
+							);
+							pools_to_remove.push(*pid);
+						}
+						false
 					}
 				} else {
 					false
 				};
 
 				if !should_withdraw {
-					pools_to_remove.push(*pid);
 					continue;
 				}
 

pallets/phala/src/snapshots/phala_pallets__test__vault_force_withdraw_with_dust_investment.snap

@@ -3,51 +3,6 @@ source: pallets/phala/src/test.rs
 expression: take_events()
 ---
 [
-    RuntimeEvent::RmrkCore(
-        Event::PropertiesRemoved {
-            collection_id: 10000,
-            maybe_nft_id: Some(
-                0,
-            ),
-        },
-    ),
-    RuntimeEvent::Uniques(
-        Event::Burned {
-            collection: 10000,
-            item: 0,
-            owner: 13009150994509951074,
-        },
-    ),
-    RuntimeEvent::RmrkCore(
-        Event::NFTBurned {
-            owner: 7813586407040180578,
-            collection_id: 10000,
-            nft_id: 0,
-        },
-    ),
-    RuntimeEvent::RmrkCore(
-        Event::PropertyRemoved {
-            collection_id: 10000,
-            maybe_nft_id: Some(
-                0,
-            ),
-            key: BoundedVec(
-                [
-                    99,
-                    114,
-                    101,
-                    97,
-                    116,
-                    101,
-                    116,
-                    105,
-                    109,
-                    101,
-                ],
-                32000,
-            ),
-        },
-    ),
     RuntimeEvent::PhalaVault(
         Event::ForceShutdown {
             pid: 1,

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");
     });
 }