fix off by one stake subsidy distribution counter

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

@@ -170,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.

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

@@ -2164,18 +2164,22 @@ 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>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  &&
+    <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 && ctx.<a href="sui_system_state_inner.md#0x3_sui_system_state_inner_epoch">epoch</a>() &gt; 560) {
-            <b>let</b> first_safe_mode_epoch = 561;
-            <b>let</b> safe_mode_epoch_count = ctx.<a href="sui_system_state_inner.md#0x3_sui_system_state_inner_epoch">epoch</a>() - first_safe_mode_epoch;
+        <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>());
             });

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.

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

@@ -866,18 +866,22 @@ 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.
-        if (ctx.epoch() >= self.parameters.stake_subsidy_start_epoch  &&
+        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 && ctx.epoch() > 560) {
-                let first_safe_mode_epoch = 561;
-                let safe_mode_epoch_count = ctx.epoch() - first_safe_mode_epoch;
+            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());
                 });
@@ -1141,6 +1145,11 @@ module sui_system::sui_system_state_inner {
         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/tests/rewards_distribution_tests.move

@@ -493,32 +493,103 @@ module sui_system::rewards_distribution_tests {
         test_scenario::return_shared(system_state);
     }
 
-    #[test]
-    fun test_stake_subsidy_with_safe_mode() {
-        use std::unit_test::assert_eq;
+    fun check_distribution_counter_invariant(system: &mut SuiSystemState, ctx: &TxContext) {
+        assert!(ctx.epoch() == system.epoch());
+        // first subsidy distribution was at epoch 20, so counter should always be ahead by 20
+        assert_eq(system.get_stake_subsidy_distribution_counter() + 20, ctx.epoch());
+    }
 
+    #[test]
+    fun test_stake_subsidy_with_safe_mode_epoch_562_to_563() {
         set_up_sui_system_state_with_big_amounts();
 
         let mut test = test_scenario::begin(VALIDATOR_ADDR_1);
         let mut sui_system = test.take_shared<SuiSystemState>();
+        let ctx = test.ctx();
+        // mimic state during epoch 562, if we're in safe mode since the 560 -> 561 epoch change
+        let start_epoch: u64 = 562;
+        let start_distribution_counter = 540;
+        let epoch_start_time = 100000000000;
+        let epoch_duration = sui_system.inner_mut_for_testing().epoch_duration_ms();
+
+        // increment epoch number (safe mode emulation)
+        start_epoch.do!(|_| ctx.increment_epoch_number());
+        sui_system.set_epoch_for_testing(start_epoch);
+        sui_system.set_stake_subsidy_distribution_counter(start_distribution_counter);
+
+        assert!(ctx.epoch() == start_epoch);
+        assert!(ctx.epoch() == sui_system.epoch());
+        assert!(sui_system.get_stake_subsidy_distribution_counter() == start_distribution_counter);
+
+        // perform advance epoch
+        sui_system
+            .inner_mut_for_testing()
+            .advance_epoch(start_epoch + 1, 65, balance::zero(), balance::zero(), 0, 0, 0, 0, epoch_start_time, ctx)
+            .destroy_for_testing(); // balance returned from `advance_epoch`
+        ctx.increment_epoch_number();
+
+        // should distribute 3 epochs worth of subsidies: 560, 561, 562
+        assert_eq(sui_system.get_stake_subsidy_distribution_counter(), start_distribution_counter + 3);
+        check_distribution_counter_invariant(&mut sui_system, ctx);
+
+        // ensure that next epoch change only distributes one epoch's worth
+        sui_system
+            .inner_mut_for_testing()
+            .advance_epoch(start_epoch + 2, 65, balance::zero(), balance::zero(), 0, 0, 0, 0, epoch_start_time + epoch_duration, ctx)
+            .destroy_for_testing(); // balance returned from `advance_epoch`
+        ctx.increment_epoch_number();
 
+        // should distribute 1 epoch's worth of subsidies: 563 only
+        assert_eq(sui_system.get_stake_subsidy_distribution_counter(), start_distribution_counter + 4);
+        check_distribution_counter_invariant(&mut sui_system, ctx);
+
+        test_scenario::return_shared(sui_system);
+        test.end();
+    }
+
+    #[test]
+    fun test_stake_subsidy_with_safe_mode_epoch_563_to_564() {
+        set_up_sui_system_state_with_big_amounts();
+
+        let mut test = test_scenario::begin(VALIDATOR_ADDR_1);
+        let mut sui_system = test.take_shared<SuiSystemState>();
         let ctx = test.ctx();
+        // mimic state during epoch 563, if we're in safe mode since the 560 -> 561 epoch change
+        let start_epoch: u64 = 563;
+        let start_distribution_counter = 540;
+        let epoch_start_time = 100000000000;
+        let epoch_duration = sui_system.inner_mut_for_testing().epoch_duration_ms();
 
         // increment epoch number (safe mode emulation)
-        562u64.do!(|_| ctx.increment_epoch_number());
-        sui_system.set_epoch_for_testing(562);
-        sui_system.set_stake_subsidy_distribution_counter(540);
+        start_epoch.do!(|_| ctx.increment_epoch_number());
+        sui_system.set_epoch_for_testing(start_epoch);
+        sui_system.set_stake_subsidy_distribution_counter(start_distribution_counter);
 
-        assert!(ctx.epoch() == 562);
-        assert!(sui_system.get_stake_subsidy_distribution_counter() == 540);
+        assert!(ctx.epoch() == start_epoch);
+        assert!(ctx.epoch() == sui_system.epoch());
+        assert!(sui_system.get_stake_subsidy_distribution_counter() == start_distribution_counter);
 
         // perform advance epoch
         sui_system
             .inner_mut_for_testing()
-            .advance_epoch(563, 65, balance::zero(), balance::zero(), 0, 0, 0, 0, 100000000000, ctx)
+            .advance_epoch(start_epoch + 1, 65, balance::zero(), balance::zero(), 0, 0, 0, 0, epoch_start_time, ctx)
+            .destroy_for_testing(); // balance returned from `advance_epoch`
+        ctx.increment_epoch_number();
+
+        // should distribute 4 epochs worth of subsidies: 560, 561, 562, 563
+        assert_eq(sui_system.get_stake_subsidy_distribution_counter(), start_distribution_counter + 4);
+        check_distribution_counter_invariant(&mut sui_system, ctx);
+
+        // ensure that next epoch change only distributes one epoch's worth
+        sui_system
+            .inner_mut_for_testing()
+            .advance_epoch(start_epoch + 2, 65, balance::zero(), balance::zero(), 0, 0, 0, 0, epoch_start_time + epoch_duration, ctx)
             .destroy_for_testing(); // balance returned from `advance_epoch`
+        ctx.increment_epoch_number();
 
-        assert_eq!(sui_system.get_stake_subsidy_distribution_counter(), 542);
+        // should distribute 1 epoch's worth of subsidies
+        assert_eq(sui_system.get_stake_subsidy_distribution_counter(), start_distribution_counter + 5);
+        check_distribution_counter_invariant(&mut sui_system, ctx);
 
         test_scenario::return_shared(sui_system);
         test.end();