ext-stk: immediate release when val unbonded

contracts/provider/external-staking/src/contract.rs

@@ -1,6 +1,6 @@
 use cosmwasm_std::{
-    coin, ensure, ensure_eq, to_binary, Coin, Decimal, DepsMut, Env, Event, IbcMsg, Order,
-    Response, StdResult, Storage, Uint128, Uint256, WasmMsg,
+    coin, ensure, ensure_eq, to_binary, Addr, Coin, CosmosMsg, Decimal, DepsMut, Env, Event,
+    IbcMsg, Order, Response, StdResult, Storage, Uint128, Uint256, WasmMsg,
 };
 use cw2::set_contract_version;
 use cw_storage_plus::{Bounder, Item, Map};
@@ -17,7 +17,7 @@ use mesh_apis::ibc::{AddValidator, ProviderPacket};
 use mesh_apis::vault_api::{SlashInfo, VaultApiHelper};
 use mesh_sync::{Tx, ValueRange};
 
-use crate::crdt::CrdtState;
+use crate::crdt::{CrdtState, State};
 use crate::error::ContractError;
 use crate::ibc::{packet_timeout, IBC_CHANNEL};
 use crate::msg::{
@@ -327,10 +327,10 @@ impl ExternalStakingContract<'_> {
     /// In non-test code, this is called from `ibc_packet_ack`
     pub(crate) fn commit_unstake(
         &self,
-        deps: DepsMut,
+        mut deps: DepsMut,
         env: Env,
         tx_id: u64,
-    ) -> Result<(), ContractError> {
+    ) -> Result<Option<CosmosMsg>, ContractError> {
         use crate::state::PendingUnbond;
 
         // Load tx
@@ -370,9 +370,18 @@ impl ExternalStakingContract<'_> {
         let amount = min(tx_amount, stake.stake.high());
         stake.stake.commit_sub(amount);
 
+        let immediate_release = matches!(
+            self.val_set.validator_state(deps.storage, &tx_validator)?,
+            State::Unbonded {}
+        );
+
         // FIXME? Release period being computed after successful IBC tx
         // (Note: this is good for now, but can be revisited in v1 design)
-        let release_at = env.block.time.plus_seconds(config.unbonding_period);
+        let release_at = if immediate_release {
+            env.block.time
+        } else {
+            env.block.time.plus_seconds(config.unbonding_period)
+        };
         let unbond = PendingUnbond { amount, release_at };
         stake.pending_unbonds.push(unbond);
 
@@ -391,9 +400,18 @@ impl ExternalStakingContract<'_> {
         self.distribution
             .save(deps.storage, &tx_validator, &distribution)?;
 
+        let unbond_msg = if immediate_release {
+            // since the validator is unbonded, we shouldn't have to
+            // wait for the unbonding period to finish to unstake
+            self.withdraw_unbonded_inner(deps.branch(), env, &tx_user)?
+                .0
+        } else {
+            None
+        };
+
         // Remove tx
         self.pending_txs.remove(deps.storage, tx_id);
-        Ok(())
+        Ok(unbond_msg)
     }
 
     /// In test code, this is called from `test_rollback_unstake`.
@@ -570,26 +588,46 @@ impl ExternalStakingContract<'_> {
     pub fn withdraw_unbonded(&self, ctx: ExecCtx) -> Result<Response, ContractError> {
         nonpayable(&ctx.info)?;
 
-        let config = self.config.load(ctx.deps.storage)?;
+        let mut resp = Response::new()
+            .add_attribute("action", "withdraw_unbonded")
+            .add_attribute("owner", ctx.info.sender.to_string());
+
+        let (maybe_msg, released) =
+            self.withdraw_unbonded_inner(ctx.deps, ctx.env, &ctx.info.sender)?;
+
+        if let Some(msg) = maybe_msg {
+            resp = resp.add_message(msg);
+        }
+
+        let resp = resp.add_attribute("amount", released.to_string());
+
+        Ok(resp)
+    }
+
+    pub fn withdraw_unbonded_inner(
+        &self,
+        deps: DepsMut,
+        env: Env,
+        account: &Addr,
+    ) -> Result<(Option<CosmosMsg>, Uint128), ContractError> {
+        let config = self.config.load(deps.storage)?;
 
         let stakes: Vec<_> = self
             .stakes
             .stake
-            .prefix(&ctx.info.sender)
-            .range(ctx.deps.storage, None, None, Order::Ascending)
+            .prefix(account)
+            .range(deps.storage, None, None, Order::Ascending)
             .collect::<Result<_, _>>()?;
 
         let released: Uint128 = stakes
             .into_iter()
             .map(|(validator, mut stake)| -> Result<_, ContractError> {
-                let released = stake.release_pending(&ctx.env.block);
+                let released = stake.release_pending(&env.block);
 
                 if !released.is_zero() {
-                    self.stakes.stake.save(
-                        ctx.deps.storage,
-                        (&ctx.info.sender, &validator),
-                        &stake,
-                    )?
+                    self.stakes
+                        .stake
+                        .save(deps.storage, (account, &validator), &stake)?
                 }
 
                 Ok(released)
@@ -598,22 +636,19 @@ impl ExternalStakingContract<'_> {
                 released.map(|released| released + acc)
             })?;
 
-        let mut resp = Response::new()
-            .add_attribute("action", "withdraw_unbonded")
-            .add_attribute("owner", ctx.info.sender.to_string())
-            .add_attribute("amount", released.to_string());
-
-        if !released.is_zero() {
+        let maybe_msg = if !released.is_zero() {
             let release_msg = config.vault.release_cross_stake(
-                ctx.info.sender.into_string(),
+                account.to_string(),
                 coin(released.u128(), &config.denom),
                 vec![],
             )?;
 
-            resp = resp.add_message(release_msg);
-        }
+            Some(release_msg.into())
+        } else {
+            None
+        };
 
-        Ok(resp)
+        Ok((maybe_msg, released))
     }
 
     /// Distributes reward among users staking via particular validator. Distribution is performed

contracts/provider/external-staking/src/crdt.rs

@@ -28,7 +28,7 @@ impl ValidatorState {
         if self.is_empty() {
             State::Unknown {}
         } else {
-            self.0[0].state.clone()
+            self.0[0].state
         }
     }
 
@@ -66,6 +66,7 @@ pub struct ValState {
 }
 
 #[cw_serde]
+#[derive(Copy)]
 pub enum State {
     /// Validator is part of the validator set.
     Active {},
@@ -393,6 +394,14 @@ impl<'a> CrdtState<'a> {
         self.validators.save(storage, valoper, &validator_state)?;
         Ok(())
     }
+
+    pub fn validator_state(&self, storage: &dyn Storage, valoper: &str) -> StdResult<State> {
+        Ok(self
+            .validators
+            .may_load(storage, valoper)?
+            .map(|state| state.get_state())
+            .unwrap_or(State::Unknown {}))
+    }
 }
 
 #[cfg(test)]

contracts/provider/external-staking/src/ibc.rs

@@ -197,7 +197,9 @@ pub fn ibc_packet_ack(
                 .add_attribute("tx_id", tx_id.to_string());
         }
         (ProviderPacket::Unstake { tx_id, .. }, AckWrapper::Result(_)) => {
-            contract.commit_unstake(deps, env, tx_id)?;
+            if let Some(msg) = contract.commit_unstake(deps, env, tx_id)? {
+                resp = resp.add_message(msg);
+            }
             resp = resp
                 .add_attribute("success", "true")
                 .add_attribute("tx_id", tx_id.to_string());

contracts/provider/external-staking/src/multitest.rs

@@ -453,6 +453,43 @@ fn unstaking() {
     assert_eq!(claim.amount.val().unwrap().u128(), 240);
 }
 
+#[test]
+fn immediate_unstake_if_unbonded_validator() {
+    let user = "user1";
+
+    let app = App::new_with_balances(&[(user, &coins(200, OSMO))]);
+
+    let owner = "owner";
+
+    let (vault, contract) = setup(&app, owner, 100).unwrap();
+
+    let validators = contract.activate_validators(["validator1"]);
+
+    vault
+        .bond()
+        .with_funds(&coins(200, OSMO))
+        .call(user)
+        .unwrap();
+    vault.stake(&contract, user, validators[0], coin(200, OSMO));
+
+    contract.remove_validator(validators[0]);
+
+    contract
+        .unstake(validators[0].to_string(), coin(200, OSMO))
+        .call(user)
+        .unwrap();
+    contract
+        .test_methods_proxy()
+        .test_commit_unstake(get_last_external_staking_pending_tx_id(&contract).unwrap())
+        .call("test")
+        .unwrap();
+
+    let claim = vault
+        .claim(user.to_string(), contract.contract_addr.to_string())
+        .unwrap();
+    assert_eq!(claim.amount.val().unwrap().u128(), 0);
+}
+
 #[test]
 fn distribution() {
     let owner = "owner";

contracts/provider/external-staking/src/multitest/utils.rs

@@ -64,6 +64,8 @@ pub(crate) trait ContractExt {
         validators: [&'static str; N],
     ) -> [&'static str; N];
 
+    fn remove_validator(&self, validator: &'static str);
+
     fn distribute_batch(
         &self,
         caller: impl AsRef<str>,
@@ -89,6 +91,14 @@ impl ContractExt for Contract<'_> {
         validators
     }
 
+    #[track_caller]
+    fn remove_validator(&self, validator: &'static str) {
+        self.test_methods_proxy()
+            .test_remove_validator(validator.to_string(), 101, 1235)
+            .call("test")
+            .unwrap();
+    }
+
     #[track_caller]
     fn distribute_batch(
         &self,

contracts/provider/external-staking/src/test_methods.rs

@@ -28,6 +28,16 @@ pub trait TestMethods {
         time: u64,
     ) -> Result<Response, Self::Error>;
 
+    /// Sets validator as `unbonded`.
+    #[msg(exec)]
+    fn test_remove_validator(
+        &self,
+        ctx: ExecCtx,
+        valoper: String,
+        height: u64,
+        time: u64,
+    ) -> Result<Response, Self::Error>;
+
     /// Commits a pending unstake.
     #[msg(exec)]
     fn test_commit_unstake(&self, ctx: ExecCtx, tx_id: u64) -> Result<Response, Self::Error>;

contracts/provider/external-staking/src/test_methods_impl.rs

@@ -67,13 +67,38 @@ impl TestMethods for ExternalStakingContract<'_> {
         }
     }
 
+    /// Sets validator as `unbonded`.
+    #[msg(exec)]
+    fn test_remove_validator(
+        &self,
+        ctx: ExecCtx,
+        valoper: String,
+        height: u64,
+        time: u64,
+    ) -> Result<Response, ContractError> {
+        #[cfg(any(feature = "mt", test))]
+        {
+            self.val_set
+                .remove_validator(ctx.deps.storage, &valoper, height, time)?;
+            Ok(Response::new())
+        }
+        #[cfg(not(any(feature = "mt", test)))]
+        {
+            let _ = (ctx, valoper, height, time);
+            Err(ContractError::Unauthorized {})
+        }
+    }
+
     /// Commits a pending unstake.
     #[msg(exec)]
     fn test_commit_unstake(&self, ctx: ExecCtx, tx_id: u64) -> Result<Response, ContractError> {
         #[cfg(any(test, feature = "mt"))]
         {
-            self.commit_unstake(ctx.deps, ctx.env, tx_id)?;
-            Ok(Response::new())
+            if let Some(msg) = self.commit_unstake(ctx.deps, ctx.env, tx_id)? {
+                Ok(Response::new().add_message(msg))
+            } else {
+                Ok(Response::new())
+            }
         }
         #[cfg(not(any(test, feature = "mt")))]
         {