sui-system: fix next epoch stake book-keeping

crates/sui-framework-snapshot/manifest.json

@@ -486,5 +486,25 @@
       "0x000000000000000000000000000000000000000000000000000000000000dee9",
       "0x000000000000000000000000000000000000000000000000000000000000000b"
     ]
+  },
+  "66": {
+    "git_revision": "86fa6e86b62a",
+    "package_ids": [
+      "0x0000000000000000000000000000000000000000000000000000000000000001",
+      "0x0000000000000000000000000000000000000000000000000000000000000002",
+      "0x0000000000000000000000000000000000000000000000000000000000000003",
+      "0x000000000000000000000000000000000000000000000000000000000000dee9",
+      "0x000000000000000000000000000000000000000000000000000000000000000b"
+    ]
+  },
+  "67": {
+    "git_revision": "3ada97c109cc",
+    "package_ids": [
+      "0x0000000000000000000000000000000000000000000000000000000000000001",
+      "0x0000000000000000000000000000000000000000000000000000000000000002",
+      "0x0000000000000000000000000000000000000000000000000000000000000003",
+      "0x000000000000000000000000000000000000000000000000000000000000dee9",
+      "0x000000000000000000000000000000000000000000000000000000000000000b"
+    ]
   }
 }

crates/sui-framework/docs/sui-system/stake_subsidy.md

@@ -9,6 +9,7 @@ title: Module `0x3::stake_subsidy`
 -  [Function `create`](#0x3_stake_subsidy_create)
 -  [Function `advance_epoch`](#0x3_stake_subsidy_advance_epoch)
 -  [Function `current_epoch_subsidy_amount`](#0x3_stake_subsidy_current_epoch_subsidy_amount)
+-  [Function `get_distribution_counter`](#0x3_stake_subsidy_get_distribution_counter)
 
 
 <pre><code><b>use</b> <a href="../move-stdlib/u64.md#0x1_u64">0x1::u64</a>;
@@ -169,7 +170,6 @@ Advance the epoch counter and draw down the subsidy for the epoch.
 
     // Drawn down the subsidy for this epoch.
     <b>let</b> <a href="stake_subsidy.md#0x3_stake_subsidy">stake_subsidy</a> = self.<a href="../sui-framework/balance.md#0x2_balance">balance</a>.split(to_withdraw);
-
     self.distribution_counter = self.distribution_counter + 1;
 
     // Decrease the subsidy amount only when the current period ends.
@@ -210,4 +210,29 @@ Returns the amount of stake subsidy to be added at the end of the current epoch.
 
 
 
+</details>
+
+<a name="0x3_stake_subsidy_get_distribution_counter"></a>
+
+## Function `get_distribution_counter`
+
+Returns the number of distributions that have occurred.
+
+
+<pre><code><b>public</b>(<b>friend</b>) <b>fun</b> <a href="stake_subsidy.md#0x3_stake_subsidy_get_distribution_counter">get_distribution_counter</a>(self: &<a href="stake_subsidy.md#0x3_stake_subsidy_StakeSubsidy">stake_subsidy::StakeSubsidy</a>): <a href="../move-stdlib/u64.md#0x1_u64">u64</a>
+</code></pre>
+
+
+
+<details>
+<summary>Implementation</summary>
+
+
+<pre><code><b>public</b>(package) <b>fun</b> <a href="stake_subsidy.md#0x3_stake_subsidy_get_distribution_counter">get_distribution_counter</a>(self: &<a href="stake_subsidy.md#0x3_stake_subsidy_StakeSubsidy">StakeSubsidy</a>): <a href="../move-stdlib/u64.md#0x1_u64">u64</a> {
+    self.distribution_counter
+}
+</code></pre>
+
+
+
 </details>

crates/sui-framework/docs/sui-system/sui_system_state_inner.md

@@ -2163,18 +2163,31 @@ gas coins.
 
     <b>let</b> storage_charge = storage_reward.value();
     <b>let</b> computation_charge = computation_reward.value();
+    <b>let</b> <b>mut</b> <a href="stake_subsidy.md#0x3_stake_subsidy">stake_subsidy</a> = <a href="../sui-framework/balance.md#0x2_balance_zero">balance::zero</a>();
 
+    // during the transition from epoch N <b>to</b> epoch N + 1, ctx.<a href="sui_system_state_inner.md#0x3_sui_system_state_inner_epoch">epoch</a>() will <b>return</b> N
+    <b>let</b> old_epoch = ctx.<a href="sui_system_state_inner.md#0x3_sui_system_state_inner_epoch">epoch</a>();
     // Include stake subsidy in the rewards given out <b>to</b> validators and stakers.
     // Delay distributing any stake subsidies until after `stake_subsidy_start_epoch`.
     // And <b>if</b> this epoch is shorter than the regular epoch duration, don't distribute any stake subsidy.
-    <b>let</b> <a href="stake_subsidy.md#0x3_stake_subsidy">stake_subsidy</a> =
-        <b>if</b> (ctx.<a href="sui_system_state_inner.md#0x3_sui_system_state_inner_epoch">epoch</a>() &gt;= self.parameters.stake_subsidy_start_epoch  &&
-            epoch_start_timestamp_ms &gt;= prev_epoch_start_timestamp + self.parameters.epoch_duration_ms)
-        {
-            self.<a href="stake_subsidy.md#0x3_stake_subsidy">stake_subsidy</a>.<a href="sui_system_state_inner.md#0x3_sui_system_state_inner_advance_epoch">advance_epoch</a>()
-        } <b>else</b> {
-            <a href="../sui-framework/balance.md#0x2_balance_zero">balance::zero</a>()
+    <b>if</b> (old_epoch &gt;= self.parameters.stake_subsidy_start_epoch  &&
+        epoch_start_timestamp_ms &gt;= prev_epoch_start_timestamp + self.parameters.epoch_duration_ms)
+    {
+        // special case for epoch 560 -&gt; 561 change bug. add extra subsidies for "safe mode"
+        // <b>where</b> reward distribution was skipped. <b>use</b> distribution counter and epoch check <b>to</b>
+        // avoiding affecting devnet and testnet
+        <b>if</b> (self.<a href="stake_subsidy.md#0x3_stake_subsidy">stake_subsidy</a>.get_distribution_counter() == 540 && old_epoch &gt; 560) {
+            // safe mode was entered on the change from 560 <b>to</b> 561. so 560 was the first epoch without proper subsidy distribution
+            <b>let</b> first_safe_mode_epoch = 560;
+            <b>let</b> safe_mode_epoch_count = old_epoch - first_safe_mode_epoch;
+            safe_mode_epoch_count.do!(|_| {
+                <a href="stake_subsidy.md#0x3_stake_subsidy">stake_subsidy</a>.join(self.<a href="stake_subsidy.md#0x3_stake_subsidy">stake_subsidy</a>.<a href="sui_system_state_inner.md#0x3_sui_system_state_inner_advance_epoch">advance_epoch</a>());
+            });
+            // done <b>with</b> catchup for safe mode epochs. distribution counter is now &gt;540, we won't hit this again
+            // fall through <b>to</b> the normal logic, which will add subsidies for the current epoch
         };
+        <a href="stake_subsidy.md#0x3_stake_subsidy">stake_subsidy</a>.join(self.<a href="stake_subsidy.md#0x3_stake_subsidy">stake_subsidy</a>.<a href="sui_system_state_inner.md#0x3_sui_system_state_inner_advance_epoch">advance_epoch</a>());
+    };
 
     <b>let</b> stake_subsidy_amount = <a href="stake_subsidy.md#0x3_stake_subsidy">stake_subsidy</a>.value();
     computation_reward.join(<a href="stake_subsidy.md#0x3_stake_subsidy">stake_subsidy</a>);

crates/sui-framework/docs/sui-system/validator.md

@@ -1076,6 +1076,8 @@ Request to add stake to the validator's staking pool, processed at the end of th
 
     <b>let</b> <a href="../sui-framework/sui.md#0x2_sui">sui</a> = self.<a href="staking_pool.md#0x3_staking_pool">staking_pool</a>.<a href="validator.md#0x3_validator_redeem_fungible_staked_sui">redeem_fungible_staked_sui</a>(fungible_staked_sui, ctx);
 
+    self.next_epoch_stake = self.next_epoch_stake - <a href="../sui-framework/sui.md#0x2_sui">sui</a>.value();
+
     <a href="../sui-framework/event.md#0x2_event_emit">event::emit</a>(
         <a href="validator.md#0x3_validator_RedeemingFungibleStakedSuiEvent">RedeemingFungibleStakedSuiEvent</a> {
             pool_id: self.<a href="validator.md#0x3_validator_staking_pool_id">staking_pool_id</a>(),
@@ -1346,7 +1348,8 @@ Process pending stakes and withdraws, called at the end of the epoch.
 
 <pre><code><b>public</b>(package) <b>fun</b> <a href="validator.md#0x3_validator_process_pending_stakes_and_withdraws">process_pending_stakes_and_withdraws</a>(self: &<b>mut</b> <a href="validator.md#0x3_validator_Validator">Validator</a>, ctx: &TxContext) {
     self.<a href="staking_pool.md#0x3_staking_pool">staking_pool</a>.<a href="validator.md#0x3_validator_process_pending_stakes_and_withdraws">process_pending_stakes_and_withdraws</a>(ctx);
-    <b>assert</b>!(<a href="validator.md#0x3_validator_stake_amount">stake_amount</a>(self) == self.next_epoch_stake, <a href="validator.md#0x3_validator_EInvalidStakeAmount">EInvalidStakeAmount</a>);
+    // TODO: bring this assertion back when we are ready.
+    // <b>assert</b>!(<a href="validator.md#0x3_validator_stake_amount">stake_amount</a>(self) == self.next_epoch_stake, <a href="validator.md#0x3_validator_EInvalidStakeAmount">EInvalidStakeAmount</a>);
 }
 </code></pre>
 

crates/sui-framework/packages/sui-system/sources/stake_subsidy.move

@@ -65,7 +65,6 @@ module sui_system::stake_subsidy {
 
         // Drawn down the subsidy for this epoch.
         let stake_subsidy = self.balance.split(to_withdraw);
-
         self.distribution_counter = self.distribution_counter + 1;
 
         // Decrease the subsidy amount only when the current period ends.
@@ -83,9 +82,13 @@ module sui_system::stake_subsidy {
         self.current_distribution_amount.min(self.balance.value())
     }
 
-    #[test_only]
     /// Returns the number of distributions that have occurred.
     public(package) fun get_distribution_counter(self: &StakeSubsidy): u64 {
         self.distribution_counter
     }
+
+    #[test_only]
+    public(package) fun set_distribution_counter(self: &mut StakeSubsidy, distribution_counter: u64) {
+        self.distribution_counter = distribution_counter;
+    }
 }

crates/sui-framework/packages/sui-system/sources/sui_system.move

@@ -726,6 +726,17 @@ module sui_system::sui_system {
         self.get_stake_subsidy_distribution_counter()
     }
 
+    #[test_only]
+    public fun set_stake_subsidy_distribution_counter(wrapper: &mut SuiSystemState, counter: u64) {
+        let self = load_system_state_mut(wrapper);
+        self.set_stake_subsidy_distribution_counter(counter)
+    }
+
+    #[test_only]
+    public fun inner_mut_for_testing(wrapper: &mut SuiSystemState): &mut SuiSystemStateInnerV2 {
+        wrapper.load_system_state_mut()
+    }
+
     // CAUTION: THIS CODE IS ONLY FOR TESTING AND THIS MACRO MUST NEVER EVER BE REMOVED.  Creates a
     // candidate validator - bypassing the proof of possession check and other metadata validation
     // in the process.

crates/sui-framework/packages/sui-system/sources/sui_system_state_inner.move

@@ -865,18 +865,31 @@ module sui_system::sui_system_state_inner {
 
         let storage_charge = storage_reward.value();
         let computation_charge = computation_reward.value();
+        let mut stake_subsidy = balance::zero();
 
+        // during the transition from epoch N to epoch N + 1, ctx.epoch() will return N
+        let old_epoch = ctx.epoch();
         // Include stake subsidy in the rewards given out to validators and stakers.
         // Delay distributing any stake subsidies until after `stake_subsidy_start_epoch`.
         // And if this epoch is shorter than the regular epoch duration, don't distribute any stake subsidy.
-        let stake_subsidy =
-            if (ctx.epoch() >= self.parameters.stake_subsidy_start_epoch  &&
-                epoch_start_timestamp_ms >= prev_epoch_start_timestamp + self.parameters.epoch_duration_ms)
-            {
-                self.stake_subsidy.advance_epoch()
-            } else {
-                balance::zero()
+        if (old_epoch >= self.parameters.stake_subsidy_start_epoch  &&
+            epoch_start_timestamp_ms >= prev_epoch_start_timestamp + self.parameters.epoch_duration_ms)
+        {
+            // special case for epoch 560 -> 561 change bug. add extra subsidies for "safe mode"
+            // where reward distribution was skipped. use distribution counter and epoch check to
+            // avoiding affecting devnet and testnet
+            if (self.stake_subsidy.get_distribution_counter() == 540 && old_epoch > 560) {
+                // safe mode was entered on the change from 560 to 561. so 560 was the first epoch without proper subsidy distribution
+                let first_safe_mode_epoch = 560;
+                let safe_mode_epoch_count = old_epoch - first_safe_mode_epoch;
+                safe_mode_epoch_count.do!(|_| {
+                    stake_subsidy.join(self.stake_subsidy.advance_epoch());
+                });
+                // done with catchup for safe mode epochs. distribution counter is now >540, we won't hit this again
+                // fall through to the normal logic, which will add subsidies for the current epoch
             };
+            stake_subsidy.join(self.stake_subsidy.advance_epoch());
+        };
 
         let stake_subsidy_amount = stake_subsidy.value();
         computation_reward.join(stake_subsidy);
@@ -1127,6 +1140,16 @@ module sui_system::sui_system_state_inner {
         self.validators.request_add_validator(min_joining_stake_for_testing, ctx);
     }
 
+    #[test_only]
+    public(package) fun set_stake_subsidy_distribution_counter(self: &mut SuiSystemStateInnerV2, counter: u64) {
+        self.stake_subsidy.set_distribution_counter(counter)
+    }
+
+    #[test_only]
+    public(package) fun epoch_duration_ms(self: &SuiSystemStateInnerV2): u64 {
+        self.parameters.epoch_duration_ms
+    }
+
     // CAUTION: THIS CODE IS ONLY FOR TESTING AND THIS MACRO MUST NEVER EVER BE REMOVED.  Creates a
     // candidate validator - bypassing the proof of possession check and other metadata validation
     // in the process.

crates/sui-framework/packages/sui-system/sources/validator.move

@@ -352,6 +352,8 @@ module sui_system::validator {
 
         let sui = self.staking_pool.redeem_fungible_staked_sui(fungible_staked_sui, ctx);
 
+        self.next_epoch_stake = self.next_epoch_stake - sui.value();
+
         event::emit(
             RedeemingFungibleStakedSuiEvent {
                 pool_id: self.staking_pool_id(),
@@ -462,7 +464,8 @@ module sui_system::validator {
     /// Process pending stakes and withdraws, called at the end of the epoch.
     public(package) fun process_pending_stakes_and_withdraws(self: &mut Validator, ctx: &TxContext) {
         self.staking_pool.process_pending_stakes_and_withdraws(ctx);
-        assert!(stake_amount(self) == self.next_epoch_stake, EInvalidStakeAmount);
+        // TODO: bring this assertion back when we are ready.
+        // assert!(stake_amount(self) == self.next_epoch_stake, EInvalidStakeAmount);
     }
 
     /// Returns true if the validator is preactive.