ensure invest pools when vault withdraw from stake pools

Phala-Network · Dec 16, 2024 · 9b6bd90 · 9b6bd90
1 parent cb723b1
commit 9b6bd90
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 57 deletions.
diff --git a/pallets/phala/src/compute/stake_pool_v2.rs b/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,

diff --git a/pallets/phala/src/compute/vault.rs b/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,

diff --git a/...s/phala/src/snapshots/phala_pallets__test__vault_force_withdraw_with_dust_investment.snap b/...s/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,

diff --git a/pallets/phala/src/test.rs b/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");
     });
 }