ugly, but simple fix for stake subsidies

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>;
@@ -210,4 +211,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,27 @@ 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>();
     // 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> (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)
+    {
+        // 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;
+            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/packages/sui-system/sources/stake_subsidy.move

@@ -83,9 +83,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,27 @@ 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();
         // 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 (ctx.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;
+                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 +1136,11 @@ 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)
+    }
+
     // 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

@@ -3,6 +3,7 @@
 
 #[test_only]
 module sui_system::rewards_distribution_tests {
+    use sui::balance;
     use sui::test_scenario::{Self, Scenario};
     use sui_system::sui_system::SuiSystemState;
     use sui_system::validator_cap::UnverifiedValidatorOperationCap;
@@ -491,4 +492,69 @@ module sui_system::rewards_distribution_tests {
         scenario.return_to_sender(cap);
         test_scenario::return_shared(system_state);
     }
+
+    #[test]
+    fun test_stake_subsidy_with_safe_mode() {
+        use std::unit_test::assert_eq;
+
+        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();
+
+        // 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);
+
+        assert!(ctx.epoch() == 562);
+        assert!(sui_system.get_stake_subsidy_distribution_counter() == 540);
+
+        // perform advance epoch
+        sui_system
+            .inner_mut_for_testing()
+            .advance_epoch(563, 65, balance::zero(), balance::zero(), 0, 0, 0, 0, 100000000000, ctx)
+            .destroy_for_testing(); // balance returned from `advance_epoch`
+
+        assert_eq!(sui_system.get_stake_subsidy_distribution_counter(), 542);
+
+        test_scenario::return_shared(sui_system);
+        test.end();
+    }
+
+    #[test]
+    // Test that the fix for the subsidy distribution doesn't affect testnet,
+    // where the distribution has no epoch delay, and the condition could result
+    // in arithmetic error.
+    fun test_stake_subsidy_with_safe_mode_testnet() {
+        use std::unit_test::assert_eq;
+
+        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();
+
+        // increment epoch number (safe mode emulation)
+        540u64.do!(|_| ctx.increment_epoch_number());
+        sui_system.set_epoch_for_testing(540);
+        sui_system.set_stake_subsidy_distribution_counter(540);
+
+        assert!(ctx.epoch() == 540);
+        assert!(sui_system.get_stake_subsidy_distribution_counter() == 540);
+
+        // perform advance epoch
+        sui_system
+            .inner_mut_for_testing()
+            .advance_epoch(541, 65, balance::zero(), balance::zero(), 0, 0, 0, 0, 100000000000, ctx)
+            .destroy_for_testing(); // balance returned from `advance_epoch`
+
+        assert_eq!(sui_system.get_stake_subsidy_distribution_counter(), 541);
+
+        test_scenario::return_shared(sui_system);
+        test.end();
+    }
 }

crates/sui-framework/published_api.txt

@@ -691,6 +691,9 @@ advance_epoch
 current_epoch_subsidy_amount
 	public fun
 	0x3::stake_subsidy
+get_distribution_counter
+	public(package) fun
+	0x3::stake_subsidy
 SystemParameters
 	public struct
 	0x3::sui_system_state_inner