batching delegations/undelegations + stTIA migration

app/app.go

@@ -736,6 +736,8 @@ func NewStrideApp(
 		app.BankKeeper,
 		app.ICAOracleKeeper,
 		app.RatelimitKeeper,
+		app.RecordsKeeper,
+		app.StakeibcKeeper,
 		app.TransferKeeper,
 	)
 	stakeTiaModule := staketia.NewAppModule(appCodec, app.StaketiaKeeper)

app/apptesting/test_helpers.go

@@ -86,7 +86,7 @@ func (s *AppTestHelper) Setup() {
 		GRPCQueryRouter: s.App.GRPCQueryRouter(),
 		Ctx:             s.Ctx,
 	}
-	s.TestAccs = CreateRandomAccounts(3)
+	s.TestAccs = CreateRandomAccounts(4)
 	s.IbcEnabled = false
 	s.IcaAddresses = make(map[string]string)
 

app/upgrades.go

@@ -33,6 +33,7 @@ import (
 	v21 "github.com/Stride-Labs/stride/v23/app/upgrades/v21"
 	v22 "github.com/Stride-Labs/stride/v23/app/upgrades/v22"
 	v23 "github.com/Stride-Labs/stride/v23/app/upgrades/v23"
+	v24 "github.com/Stride-Labs/stride/v23/app/upgrades/v24"
 	v3 "github.com/Stride-Labs/stride/v23/app/upgrades/v3"
 	v4 "github.com/Stride-Labs/stride/v23/app/upgrades/v4"
 	v5 "github.com/Stride-Labs/stride/v23/app/upgrades/v5"
@@ -311,6 +312,19 @@ func (app *StrideApp) setupUpgradeHandlers(appOpts servertypes.AppOptions) {
 		),
 	)
 
+	// v24 upgrade handler
+	app.UpgradeKeeper.SetUpgradeHandler(
+		v24.UpgradeName,
+		v24.CreateUpgradeHandler(
+			app.mm,
+			app.configurator,
+			app.BankKeeper,
+			app.RecordsKeeper,
+			app.StakeibcKeeper,
+			app.StaketiaKeeper,
+		),
+	)
+
 	upgradeInfo, err := app.UpgradeKeeper.ReadUpgradeInfoFromDisk()
 	if err != nil {
 		panic(fmt.Errorf("Failed to read upgrade info from disk: %w", err))

app/upgrades/v10/upgrades_test.go

@@ -165,10 +165,10 @@ func (s *UpgradeTestSuite) TestMigrateCallbackData() {
 	}
 	initialUndelegateCallbackArgs := stakeibctypes.UndelegateCallback{
 		HostZoneId: "host-0",
-		SplitDelegations: []*types.SplitDelegation{{
-			Validator: "val-0",
-			Amount:    sdkmath.NewInt(1),
-		}},
+		// SplitDelegations: []*types.SplitDelegation{{
+		// 	Validator: "val-0",
+		// 	Amount:    sdkmath.NewInt(1),
+		// }},
 	}
 	initialTransferCallbackArgs := recordstypes.TransferCallback{
 		DepositRecordId: 1,

app/upgrades/v24/upgrades.go

@@ -0,0 +1,128 @@
+package v24
+
+import (
+	errorsmod "cosmossdk.io/errors"
+	sdkmath "cosmossdk.io/math"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/types/module"
+	bankkeeper "github.com/cosmos/cosmos-sdk/x/bank/keeper"
+	upgradetypes "github.com/cosmos/cosmos-sdk/x/upgrade/types"
+
+	recordskeeper "github.com/Stride-Labs/stride/v23/x/records/keeper"
+	recordstypes "github.com/Stride-Labs/stride/v23/x/records/types"
+	stakeibckeeper "github.com/Stride-Labs/stride/v23/x/stakeibc/keeper"
+	staketiakeeper "github.com/Stride-Labs/stride/v23/x/staketia/keeper"
+)
+
+var (
+	UpgradeName = "v24"
+)
+
+// CreateUpgradeHandler creates an SDK upgrade handler for v23
+func CreateUpgradeHandler(
+	mm *module.Manager,
+	configurator module.Configurator,
+	bankKeeper bankkeeper.Keeper,
+	recordsKeeper recordskeeper.Keeper,
+	stakeibcKeeper stakeibckeeper.Keeper,
+	staketiaKeeper staketiakeeper.Keeper,
+) upgradetypes.UpgradeHandler {
+	return func(ctx sdk.Context, _ upgradetypes.Plan, vm module.VersionMap) (module.VersionMap, error) {
+		ctx.Logger().Info("Starting upgrade v24...")
+
+		// Migrate data structures
+		MigrateHostZones(ctx, stakeibcKeeper)
+		MigrateEpochUnbondingRecords(ctx, recordsKeeper)
+
+		// Migrate staketia to stakeibc
+		if err := staketiakeeper.InitiateMigration(ctx, staketiaKeeper, bankKeeper, recordsKeeper, stakeibcKeeper); err != nil {
+			return vm, errorsmod.Wrapf(err, "unable to migrate staketia to stakeibc")
+		}
+
+		// TODO: add celestia validator set
+
+		ctx.Logger().Info("Running module migrations...")
+		return mm.RunMigrations(ctx, configurator, vm)
+	}
+}
+
+// Migrate host zones to accomodate the staketia migration changes, adding a
+// redemptions enabled field to each host zone
+func MigrateHostZones(ctx sdk.Context, k stakeibckeeper.Keeper) {
+	ctx.Logger().Info("Migrating host zones...")
+
+	for _, hostZone := range k.GetAllHostZone(ctx) {
+		hostZone.RedemptionsEnabled = true
+		k.SetHostZone(ctx, hostZone)
+	}
+}
+
+// Migrates the deposit records to set the DelegationTxsInProgress field
+// which should be 1 if the status was DELEGATION_IN_PROGRESS, and 0 otherwise
+func MigrateDepositRecords(ctx sdk.Context, k recordskeeper.Keeper) {
+	ctx.Logger().Info("Migrating deposit records...")
+
+	for _, depositRecord := range k.GetAllDepositRecord(ctx) {
+		if depositRecord.Status == recordstypes.DepositRecord_DELEGATION_IN_PROGRESS {
+			depositRecord.DelegationTxsInProgress = 1
+		} else {
+			depositRecord.DelegationTxsInProgress = 0
+		}
+		k.SetDepositRecord(ctx, depositRecord)
+	}
+}
+
+// Migrates a single host zone unbonding record to add the new fields: StTokensToBurn,
+// NativeTokensToUnbond, and ClaimableNativeTokens
+//
+// If the record is in status: UNBONDING_QUEUE, EXIT_TRANSFER_QUEUE, or EXIT_TRANSFER_IN_PROGRESS,
+// set stTokensToBurn, NativeTokensToUnbond, and ClaimableNativeTokens all to 0
+//
+// If the record is in status: UNBONDING_IN_PROGRESS
+// set StTokensToBurn to the value of StTokenAmount, NativeTokensToUnbond to the value of NativeTokenAmount,
+// and ClaimableNativeTokens to 0
+//
+// If the record is in status CLAIMABLE,
+// set StTokensToBurn and NativeTokensToUnbond to 0, and set ClaimableNativeTokens to the value of NativeTokenAmount
+//
+// If the record is in status UNBONDING_IN_PROGRESS, we need to also set UndelegationTxsInProgress to 1;
+// otherwise, it should be set to 0
+func MigrateHostZoneUnbondingRecords(hostZoneUnbonding *recordstypes.HostZoneUnbonding) *recordstypes.HostZoneUnbonding {
+	if hostZoneUnbonding.Status == recordstypes.HostZoneUnbonding_UNBONDING_QUEUE ||
+		hostZoneUnbonding.Status == recordstypes.HostZoneUnbonding_EXIT_TRANSFER_QUEUE ||
+		hostZoneUnbonding.Status == recordstypes.HostZoneUnbonding_EXIT_TRANSFER_IN_PROGRESS {
+
+		hostZoneUnbonding.StTokensToBurn = sdkmath.ZeroInt()
+		hostZoneUnbonding.NativeTokensToUnbond = sdkmath.ZeroInt()
+		hostZoneUnbonding.ClaimableNativeTokens = sdkmath.ZeroInt()
+		hostZoneUnbonding.UndelegationTxsInProgress = 0
+
+	} else if hostZoneUnbonding.Status == recordstypes.HostZoneUnbonding_UNBONDING_IN_PROGRESS {
+		hostZoneUnbonding.StTokensToBurn = hostZoneUnbonding.StTokenAmount
+		hostZoneUnbonding.NativeTokensToUnbond = hostZoneUnbonding.NativeTokenAmount
+		hostZoneUnbonding.ClaimableNativeTokens = sdkmath.ZeroInt()
+		hostZoneUnbonding.UndelegationTxsInProgress = 1
+
+	} else if hostZoneUnbonding.Status == recordstypes.HostZoneUnbonding_CLAIMABLE {
+		hostZoneUnbonding.StTokensToBurn = sdkmath.ZeroInt()
+		hostZoneUnbonding.NativeTokensToUnbond = sdkmath.ZeroInt()
+		hostZoneUnbonding.ClaimableNativeTokens = hostZoneUnbonding.NativeTokenAmount
+		hostZoneUnbonding.UndelegationTxsInProgress = 0
+	}
+
+	return hostZoneUnbonding
+}
+
+// Migrate epoch unbonding records to accomodate the batched undelegations code changes,
+// adding the new accounting fields to the host zone unbonding records
+func MigrateEpochUnbondingRecords(ctx sdk.Context, k recordskeeper.Keeper) {
+	ctx.Logger().Info("Migrating epoch unbonding records...")
+
+	for _, epochUnbondingRecord := range k.GetAllEpochUnbondingRecord(ctx) {
+		for i, oldHostZoneUnbondingRecord := range epochUnbondingRecord.HostZoneUnbondings {
+			updatedHostZoneUnbondingRecord := MigrateHostZoneUnbondingRecords(oldHostZoneUnbondingRecord)
+			epochUnbondingRecord.HostZoneUnbondings[i] = updatedHostZoneUnbondingRecord
+		}
+		k.SetEpochUnbondingRecord(ctx, epochUnbondingRecord)
+	}
+}

app/upgrades/v24/upgrades_test.go

@@ -0,0 +1,205 @@
+package v24_test
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/suite"
+
+	sdkmath "cosmossdk.io/math"
+
+	"github.com/Stride-Labs/stride/v23/app/apptesting"
+	v24 "github.com/Stride-Labs/stride/v23/app/upgrades/v24"
+	recordstypes "github.com/Stride-Labs/stride/v23/x/records/types"
+	stakeibctypes "github.com/Stride-Labs/stride/v23/x/stakeibc/types"
+)
+
+type UpgradeTestSuite struct {
+	apptesting.AppTestHelper
+}
+
+func (s *UpgradeTestSuite) SetupTest() {
+	s.Setup()
+}
+
+func TestKeeperTestSuite(t *testing.T) {
+	suite.Run(t, new(UpgradeTestSuite))
+}
+
+func (s *UpgradeTestSuite) TestUpgrade() {
+}
+
+func (s *UpgradeTestSuite) TestMigrateHostZones() {
+	// Create a host zone with redemptions enabled set to false
+	hostZone := stakeibctypes.HostZone{
+		ChainId:            "chain-0",
+		RedemptionsEnabled: false,
+	}
+	s.App.StakeibcKeeper.SetHostZone(s.Ctx, hostZone)
+
+	// Call migration function
+	v24.MigrateHostZones(s.Ctx, s.App.StakeibcKeeper)
+
+	// Confirm host route was migrated
+	for _, hostZone := range s.App.StakeibcKeeper.GetAllHostZone(s.Ctx) {
+		s.Require().True(hostZone.RedemptionsEnabled)
+	}
+}
+
+func (s *UpgradeTestSuite) TestMigrateDepositRecords() {
+	// Create initial deposit records across each status
+	testCases := []struct {
+		status                          recordstypes.DepositRecord_Status
+		expectedDelegationTxsInProgress uint64
+	}{
+		{
+			status:                          recordstypes.DepositRecord_TRANSFER_QUEUE,
+			expectedDelegationTxsInProgress: 0,
+		},
+		{
+			status:                          recordstypes.DepositRecord_TRANSFER_IN_PROGRESS,
+			expectedDelegationTxsInProgress: 0,
+		},
+		{
+			status:                          recordstypes.DepositRecord_DELEGATION_QUEUE,
+			expectedDelegationTxsInProgress: 0,
+		},
+		{
+			status:                          recordstypes.DepositRecord_DELEGATION_IN_PROGRESS,
+			expectedDelegationTxsInProgress: 1,
+		},
+	}
+
+	for id, tc := range testCases {
+		s.App.RecordsKeeper.SetDepositRecord(s.Ctx, recordstypes.DepositRecord{
+			Id:     uint64(id),
+			Status: tc.status,
+		})
+	}
+
+	// Migrate the records
+	v24.MigrateDepositRecords(s.Ctx, s.App.RecordsKeeper)
+
+	// Confirm the expected status for each
+	for id, tc := range testCases {
+		depositRecord, found := s.App.RecordsKeeper.GetDepositRecord(s.Ctx, uint64(id))
+		s.Require().True(found, "deposit record %d should have been found", id)
+		s.Require().Equal(tc.expectedDelegationTxsInProgress, depositRecord.DelegationTxsInProgress,
+			"delegation txs in progress for record %d", id)
+	}
+}
+
+func (s *UpgradeTestSuite) TestMigrateEpochUnbondingRecords() {
+	recordTestCases := []struct {
+		epochNumber                       uint64
+		chainId                           string
+		status                            recordstypes.HostZoneUnbonding_Status
+		stTokenAmount                     int64
+		nativeTokenAmount                 int64
+		expectedStTokensToBurn            int64
+		expectedNativeTokensToUnbond      int64
+		expectedClaimableNativeTokens     int64
+		expectedUndelegationTxsInProgress uint64
+	}{
+		{
+			epochNumber: 1,
+			chainId:     "chain-1",
+			status:      recordstypes.HostZoneUnbonding_UNBONDING_QUEUE,
+
+			stTokenAmount:     1,
+			nativeTokenAmount: 2,
+
+			expectedStTokensToBurn:            0,
+			expectedNativeTokensToUnbond:      0,
+			expectedClaimableNativeTokens:     0,
+			expectedUndelegationTxsInProgress: 0,
+		},
+		{
+			epochNumber: 1,
+			chainId:     "chain-2",
+			status:      recordstypes.HostZoneUnbonding_UNBONDING_IN_PROGRESS,
+
+			stTokenAmount:     3,
+			nativeTokenAmount: 4,
+
+			expectedStTokensToBurn:            3,
+			expectedNativeTokensToUnbond:      4,
+			expectedClaimableNativeTokens:     0,
+			expectedUndelegationTxsInProgress: 1,
+		},
+		{
+			epochNumber: 2,
+			chainId:     "chain-3",
+			status:      recordstypes.HostZoneUnbonding_EXIT_TRANSFER_QUEUE,
+
+			stTokenAmount:     5,
+			nativeTokenAmount: 6,
+
+			expectedStTokensToBurn:            0,
+			expectedNativeTokensToUnbond:      0,
+			expectedClaimableNativeTokens:     0,
+			expectedUndelegationTxsInProgress: 0,
+		},
+		{
+			epochNumber: 2,
+			chainId:     "chain-4",
+			status:      recordstypes.HostZoneUnbonding_EXIT_TRANSFER_IN_PROGRESS,
+
+			stTokenAmount:     7,
+			nativeTokenAmount: 8,
+
+			expectedStTokensToBurn:            0,
+			expectedNativeTokensToUnbond:      0,
+			expectedClaimableNativeTokens:     0,
+			expectedUndelegationTxsInProgress: 0,
+		},
+		{
+			epochNumber: 4,
+			chainId:     "chain-5",
+			status:      recordstypes.HostZoneUnbonding_CLAIMABLE,
+
+			stTokenAmount:     9,
+			nativeTokenAmount: 10,
+
+			expectedStTokensToBurn:            0,
+			expectedNativeTokensToUnbond:      0,
+			expectedClaimableNativeTokens:     10,
+			expectedUndelegationTxsInProgress: 0,
+		},
+	}
+
+	// Create the initial host zone unbonding and epoch unbonding records
+	for _, tc := range recordTestCases {
+		if _, found := s.App.RecordsKeeper.GetEpochUnbondingRecord(s.Ctx, tc.epochNumber); !found {
+			s.App.RecordsKeeper.SetEpochUnbondingRecord(s.Ctx, recordstypes.EpochUnbondingRecord{
+				EpochNumber: tc.epochNumber,
+			})
+		}
+
+		hostZoneUnbonding := recordstypes.HostZoneUnbonding{
+			HostZoneId:        tc.chainId,
+			Status:            tc.status,
+			StTokenAmount:     sdkmath.NewInt(tc.stTokenAmount),
+			NativeTokenAmount: sdkmath.NewInt(tc.nativeTokenAmount),
+		}
+		err := s.App.RecordsKeeper.SetHostZoneUnbondingRecord(s.Ctx, tc.epochNumber, tc.chainId, hostZoneUnbonding)
+		s.Require().NoError(err, "no error expected when creating epoch unbonding records")
+	}
+
+	// Call migration function
+	v24.MigrateEpochUnbondingRecords(s.Ctx, s.App.RecordsKeeper)
+
+	// Confirm new fields were added
+	for _, tc := range recordTestCases {
+		hostZoneUnbonding, found := s.App.RecordsKeeper.GetHostZoneUnbondingByChainId(s.Ctx, tc.epochNumber, tc.chainId)
+		s.Require().True(found, "host zone unbonding should have been found for %d and %s", tc.epochNumber, tc.chainId)
+
+		s.Require().Equal(tc.expectedStTokensToBurn, hostZoneUnbonding.StTokensToBurn.Int64(),
+			"%s stTokens to burn", tc.chainId)
+		s.Require().Equal(tc.expectedNativeTokensToUnbond, hostZoneUnbonding.NativeTokensToUnbond.Int64(),
+			"%s native to unbond", tc.chainId)
+		s.Require().Equal(tc.expectedClaimableNativeTokens, hostZoneUnbonding.ClaimableNativeTokens.Int64(),
+			"%s claimable native", tc.chainId)
+		s.Require().Equal(tc.expectedUndelegationTxsInProgress, hostZoneUnbonding.UndelegationTxsInProgress,
+			"%s undelegation txs in progress", tc.chainId)
+	}
+}