Add delay for changing reward fee. Return money on failure to withdra…

…w to the original user

staking-farm/src/farm.rs

@@ -21,7 +21,7 @@ pub const GET_OWNER_METHOD: &str = "get_owner_account_id";
 pub struct RewardDistribution {
     pub undistributed: Balance,
     pub unclaimed: Balance,
-    pub reward_per_shard: U256,
+    pub reward_per_share: U256,
     pub reward_round: u64,
 }
 
@@ -66,7 +66,7 @@ impl StakingContract {
             last_distribution: RewardDistribution {
                 undistributed: amount,
                 unclaimed: 0,
-                reward_per_shard: U256::zero(),
+                reward_per_share: U256::zero(),
                 reward_round: 0,
             },
         });
@@ -110,9 +110,9 @@ impl StakingContract {
         distribution.unclaimed += reward_added;
         distribution.undistributed -= reward_added;
         if total_staked == 0 {
-            distribution.reward_per_shard = U256::zero();
+            distribution.reward_per_share = U256::zero();
         } else {
-            distribution.reward_per_shard = farm.last_distribution.reward_per_shard
+            distribution.reward_per_share = farm.last_distribution.reward_per_share
                 + U256::from(reward_added) * U256::from(DENOMINATOR) / U256::from(total_staked);
         }
         Some(distribution)
@@ -122,20 +122,23 @@ impl StakingContract {
         &self,
         account: &Account,
         farm_id: u64,
-        farm: &Farm,
+        farm: &mut Farm,
     ) -> (U256, Balance) {
-        let user_rps = account
-            .last_farm_reward_per_share
-            .get(&farm_id)
-            .cloned()
-            .unwrap_or(U256::zero());
         if let Some(distribution) = self.internal_calculate_distribution(
             &farm,
             self.total_stake_shares - self.total_burn_shares,
         ) {
+            if distribution.reward_round != farm.last_distribution.reward_round {
+                farm.last_distribution = distribution.clone();
+            }
+            let user_rps = account
+                .last_farm_reward_per_share
+                .get(&farm_id)
+                .cloned()
+                .unwrap_or(U256::zero());
             (
-                farm.last_distribution.reward_per_shard,
-                (U256::from(account.stake_shares) * (distribution.reward_per_shard - user_rps)
+                farm.last_distribution.reward_per_share,
+                (U256::from(account.stake_shares) * (distribution.reward_per_share - user_rps)
                     / DENOMINATOR)
                     .as_u128(),
             )
@@ -144,26 +147,14 @@ impl StakingContract {
         }
     }
 
-    fn internal_distribute(&mut self, farm: &mut Farm) {
-        if let Some(distribution) = self.internal_calculate_distribution(
-            &farm,
-            self.total_stake_shares - self.total_burn_shares,
-        ) {
-            if distribution.reward_round != farm.last_distribution.reward_round {
-                farm.last_distribution = distribution;
-            }
-        }
-    }
-
     fn internal_distribute_reward(
         &mut self,
         account: &mut Account,
         farm_id: u64,
         mut farm: &mut Farm,
     ) {
-        self.internal_distribute(&mut farm);
         let (new_user_rps, claim_amount) =
-            self.internal_unclaimed_balance(&account, farm_id, &farm);
+            self.internal_unclaimed_balance(&account, farm_id, &mut farm);
         account
             .last_farm_reward_per_share
             .insert(farm_id, new_user_rps);
@@ -226,7 +217,8 @@ impl StakingContract {
         )
         .then(ext_self::callback_post_withdraw_reward(
             token_id.clone(),
-            send_account_id.clone(),
+            // Return funds to the account that was deducted from vs caller.
+            claim_account_id.clone(),
             U128(amount),
             env::current_account_id(),
             0,

staking-farm/src/internal.rs

@@ -244,7 +244,10 @@ impl StakingContract {
             let burn_fee = self.burn_fee_fraction.multiply(total_reward);
 
             // The validation fee that the contract owner takes.
-            let owners_fee = self.reward_fee_fraction.multiply(total_reward - burn_fee);
+            let owners_fee = self
+                .reward_fee_fraction
+                .current()
+                .multiply(total_reward - burn_fee);
 
             // Distributing the remaining reward to the delegators first.
             let remaining_reward = total_reward - owners_fee - burn_fee;

staking-farm/src/lib.rs

@@ -68,6 +68,44 @@ pub struct BurnInfo {
     pub unstaked_available_epoch_height: EpochHeight,
 }
 
+/// Updatable reward fee only after NUM_EPOCHS_TO_UNLOCK.
+#[derive(BorshDeserialize, BorshSerialize)]
+pub struct UpdatableRewardFee {
+    reward_fee_fraction: Ratio,
+    next_reward_fee_fraction: Ratio,
+    available_epoch_height: EpochHeight,
+}
+
+impl UpdatableRewardFee {
+    pub fn new(reward_fee_fraction: Ratio) -> Self {
+        Self {
+            reward_fee_fraction: reward_fee_fraction.clone(),
+            next_reward_fee_fraction: reward_fee_fraction,
+            available_epoch_height: 0,
+        }
+    }
+
+    pub fn current(&self) -> &Ratio {
+        if env::epoch_height() >= self.available_epoch_height {
+            &self.next_reward_fee_fraction
+        } else {
+            &self.reward_fee_fraction
+        }
+    }
+
+    pub fn next(&self) -> &Ratio {
+        &self.next_reward_fee_fraction
+    }
+
+    pub fn set(&mut self, next_reward_fee_fraction: Ratio) {
+        if env::epoch_height() >= self.available_epoch_height {
+            self.reward_fee_fraction = self.next_reward_fee_fraction.clone();
+        }
+        self.next_reward_fee_fraction = next_reward_fee_fraction;
+        self.available_epoch_height = env::epoch_height() + NUM_EPOCHS_TO_UNLOCK
+    }
+}
+
 #[near_bindgen]
 #[derive(BorshDeserialize, BorshSerialize)]
 pub struct StakingContract {
@@ -85,10 +123,10 @@ pub struct StakingContract {
     pub total_staked_balance: Balance,
     /// The total burn share balance, that will not be accounted in the farming.
     pub total_burn_shares: NumStakeShares,
-    /// The total amount to burn that will be available
+    /// The total amount to burn that will be available.
     /// The fraction of the reward that goes to the owner of the staking pool for running the
     /// validator node.
-    pub reward_fee_fraction: Ratio,
+    pub reward_fee_fraction: UpdatableRewardFee,
     /// The fraction of the reward that gets burnt.
     pub burn_fee_fraction: Ratio,
     /// Persistent map from an account ID to the corresponding account.
@@ -118,7 +156,7 @@ impl Default for StakingContract {
     }
 }
 
-#[derive(BorshDeserialize, BorshSerialize, Serialize, Deserialize, Clone)]
+#[derive(BorshDeserialize, BorshSerialize, Serialize, Deserialize, Clone, PartialEq, Debug)]
 #[serde(crate = "near_sdk::serde")]
 pub struct Ratio {
     pub numerator: u32,
@@ -179,7 +217,7 @@ impl StakingContract {
             total_staked_balance,
             total_stake_shares: NumStakeShares::from(total_staked_balance),
             total_burn_shares: 0,
-            reward_fee_fraction,
+            reward_fee_fraction: UpdatableRewardFee::new(reward_fee_fraction),
             burn_fee_fraction,
             accounts: UnorderedMap::new(StorageKeys::Accounts),
             farms: Vector::new(StorageKeys::Farms),
@@ -676,4 +714,47 @@ mod tests {
         emulator.contract.add_authorized_farm_token(&bob());
         add_farm(&mut emulator, 100);
     }
+
+    #[test]
+    fn test_change_reward_fee() {
+        let mut emulator = Emulator::new(
+            owner(),
+            "KuTCtARNzxZQ3YvXDeLjx83FDqxv2SdQTSbiq876zR7"
+                .parse()
+                .unwrap(),
+            zero_fee(),
+        );
+        assert_eq!(emulator.contract.get_reward_fee_fraction(), zero_fee());
+        emulator.update_context(owner(), 0);
+        let new_fee = Ratio {
+            numerator: 1,
+            denominator: 10,
+        };
+        emulator
+            .contract
+            .update_reward_fee_fraction(new_fee.clone());
+        // The fee is still the old one.
+        assert_eq!(emulator.contract.get_reward_fee_fraction(), zero_fee());
+        assert_eq!(
+            emulator
+                .contract
+                .get_pool_summary()
+                .next_reward_fee_fraction,
+            new_fee
+        );
+        emulator.skip_epochs(1);
+        assert_eq!(emulator.contract.get_reward_fee_fraction(), zero_fee());
+        emulator.skip_epochs(3);
+        assert_eq!(emulator.contract.get_reward_fee_fraction(), new_fee);
+        // Update once again.
+        let new_fee2 = Ratio {
+            numerator: 2,
+            denominator: 10,
+        };
+        emulator.update_context(owner(), 0);
+        emulator
+            .contract
+            .update_reward_fee_fraction(new_fee2.clone());
+        assert_eq!(emulator.contract.get_reward_fee_fraction(), new_fee);
+    }
 }

staking-farm/src/owner.rs

@@ -45,7 +45,7 @@ impl StakingContract {
     }
 
     /// Changes contract owner. Must be called by current owner.
-    pub fn set_owner_id(&self, owner_id: &AccountId) {
+    pub fn set_owner_id(owner_id: &AccountId) {
         let prev_owner = StakingContract::internal_set_owner(owner_id).expect("MUST HAVE OWNER");
         assert_eq!(
             prev_owner,
@@ -71,7 +71,7 @@ impl StakingContract {
         reward_fee_fraction.assert_valid();
 
         let need_to_restake = self.internal_ping();
-        self.reward_fee_fraction = reward_fee_fraction;
+        self.reward_fee_fraction.set(reward_fee_fraction);
         if need_to_restake {
             self.internal_restake();
         }
@@ -211,7 +211,11 @@ pub extern "C" fn update() {
     );
     unsafe {
         // Load code into register 0 result from the promise.
-        sys::promise_result(0, 0);
+        match sys::promise_result(0, 0) {
+            1 => {}
+            // Not ready or failed.
+            _ => env::panic_str("Failed to fetch the new code"),
+        };
         // Update current contract with code from register 0.
         let promise_id = sys::promise_batch_create(
             current_id.as_bytes().len() as _,

staking-farm/src/test_utils.rs

@@ -29,6 +29,7 @@ pub fn yton(yocto_amount: Balance) -> Balance {
 
 /// Checks that two amount are within epsilon
 pub fn almost_equal(left: Balance, right: Balance, epsilon: Balance) -> bool {
+    println!("{} ~= {}", left, right);
     if left > right {
         (left - right) < epsilon
     } else {

staking-farm/src/views.rs

@@ -58,6 +58,8 @@ pub struct PoolSummary {
     /// The fraction of the reward that goes to the owner of the staking pool for running the
     /// validator node.
     pub reward_fee_fraction: Ratio,
+    /// If reward fee fraction is changing, this will be different from current.
+    pub next_reward_fee_fraction: Ratio,
     /// The fraction of the reward that gets burnt.
     pub burn_fee_fraction: Ratio,
     /// Active farms that affect stakers.
@@ -76,7 +78,8 @@ impl StakingContract {
         PoolSummary {
             owner: StakingContract::get_owner_id(),
             total_staked_balance: self.total_staked_balance,
-            reward_fee_fraction: self.reward_fee_fraction.clone(),
+            reward_fee_fraction: self.reward_fee_fraction.current().clone(),
+            next_reward_fee_fraction: self.reward_fee_fraction.next().clone(),
             burn_fee_fraction: self.burn_fee_fraction.clone(),
             farms: self.get_active_farms(),
         }
@@ -143,8 +146,8 @@ impl StakingContract {
             return U128(0);
         }
         let account = self.accounts.get(&account_id).expect("ERR_NO_ACCOUNT");
-        let farm = self.farms.get(farm_id).expect("ERR_NO_FARM");
-        let (_rps, reward) = self.internal_unclaimed_balance(&account, farm_id, &farm);
+        let mut farm = self.farms.get(farm_id).expect("ERR_NO_FARM");
+        let (_rps, reward) = self.internal_unclaimed_balance(&account, farm_id, &mut farm);
         let prev_reward = *account.amounts.get(&farm.token_id).unwrap_or(&0);
         U128(reward + prev_reward)
     }
@@ -183,7 +186,7 @@ impl StakingContract {
 
     /// Returns the current reward fee as a fraction.
     pub fn get_reward_fee_fraction(&self) -> Ratio {
-        self.reward_fee_fraction.clone()
+        self.reward_fee_fraction.current().clone()
     }
 
     /// Returns the staking public key