feat: merkle_root bundles migration

app/app.go

@@ -409,6 +409,7 @@ func New(
 			app.StakersKeeper,
 			app.StakingKeeper,
 			app.BankKeeper,
+			app.BundlesKeeper,
 		),
 	)
 

app/upgrades/v2_0/upgrade.go

@@ -3,6 +3,7 @@ package v2_0
 import (
 	"context"
 	"fmt"
+	bundleskeeper "github.com/KYVENetwork/chain/x/bundles/keeper"
 
 	"github.com/KYVENetwork/chain/x/stakers/types"
 	bankkeeper "github.com/cosmos/cosmos-sdk/x/bank/keeper"
@@ -33,6 +34,7 @@ func CreateUpgradeHandler(
 	stakersKeeper *stakerskeeper.Keeper,
 	stakingKeeper *stakingkeeper.Keeper,
 	bankKeeper bankkeeper.Keeper,
+	bundlesKeeper bundleskeeper.Keeper,
 ) upgradetypes.UpgradeHandler {
 	return func(ctx context.Context, plan upgradetypes.Plan, fromVM module.VersionMap) (module.VersionMap, error) {
 		sdkCtx := sdk.UnwrapSDKContext(ctx)
@@ -45,6 +47,9 @@ func CreateUpgradeHandler(
 		// Run KYVE migrations
 		migrateProtocolStakers(sdkCtx, delegationKeeper, stakersKeeper, stakingKeeper, bankKeeper)
 
+		// Run Bundles Merkle Roots migrations
+		bundlesKeeper.SetBundlesMigrationUpgradeHeight(sdkCtx, uint64(sdkCtx.BlockHeight()))
+
 		logger.Info(fmt.Sprintf("finished upgrade %v", UpgradeName))
 
 		return migratedVersionMap, err

x/bundles/keeper/getters_migration.go

@@ -0,0 +1,29 @@
+package keeper
+
+import (
+	"encoding/binary"
+	"errors"
+	"github.com/KYVENetwork/chain/x/bundles/types"
+	"github.com/cosmos/cosmos-sdk/runtime"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+)
+
+func (k Keeper) GetBundlesMigrationUpgradeHeight(ctx sdk.Context) (uint64, error) {
+	storeAdapter := runtime.KVStoreAdapter(k.storeService.OpenKVStore(ctx))
+
+	if storeAdapter.Has(types.BundlesMigrationHeightKey) {
+		return binary.BigEndian.Uint64(storeAdapter.Get(types.BundlesMigrationHeightKey)), nil
+	}
+	return 0, errors.New("upgrade height can't be zero")
+}
+
+// SetBundlesMigrationUpgradeHeight stores the upgrade height of the v2.0 bundles migration
+// upgrade in the KV-Store.
+func (k Keeper) SetBundlesMigrationUpgradeHeight(ctx sdk.Context, upgradeHeight uint64) {
+	storeAdapter := runtime.KVStoreAdapter(k.storeService.OpenKVStore(ctx))
+
+	bz := make([]byte, 8)
+	binary.BigEndian.PutUint64(bz, upgradeHeight)
+
+	storeAdapter.Set(types.BundlesMigrationHeightKey, bz)
+}

x/bundles/keeper/keeper.go

@@ -1,13 +1,11 @@
 package keeper
 
 import (
-	"fmt"
-
 	"cosmossdk.io/collections"
 	"cosmossdk.io/core/store"
 	"cosmossdk.io/log"
-
 	storetypes "cosmossdk.io/store/types"
+	"fmt"
 	"github.com/KYVENetwork/chain/util"
 	"github.com/KYVENetwork/chain/x/bundles/types"
 	"github.com/cosmos/cosmos-sdk/codec"
@@ -92,3 +90,13 @@ func (k Keeper) InitMemStore(gasCtx sdk.Context) {
 		memStoreInitialized = true
 	}
 }
+
+// TODO: remove after migration
+func (k Keeper) Migration_GetStoreService() store.KVStoreService {
+	return k.storeService
+}
+
+// TODO: remove after migration
+func (k Keeper) Migration_GetCodec() codec.BinaryCodec {
+	return k.cdc
+}

x/bundles/migration/migration.go

@@ -0,0 +1,121 @@
+package migration
+
+import (
+	"cosmossdk.io/log"
+	"cosmossdk.io/store/prefix"
+	storeTypes "cosmossdk.io/store/types"
+	"embed"
+	_ "embed"
+	"encoding/hex"
+	"fmt"
+	"github.com/KYVENetwork/chain/util"
+	bundleskeeper "github.com/KYVENetwork/chain/x/bundles/keeper"
+	"github.com/KYVENetwork/chain/x/bundles/types"
+	"github.com/cosmos/cosmos-sdk/runtime"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"strconv"
+	"strings"
+)
+
+var logger log.Logger
+
+const (
+	BundlesMigrationStepSizePerPool uint64 = 100
+	WaitingBlockPeriod                     = 1
+)
+
+//go:embed files/*
+var merkelRoots embed.FS
+
+type BundlesMigrationEntry struct {
+	merkleRoots []byte
+	poolId      uint64
+	maxBundleId uint64
+}
+
+// bundlesMigration includes the poolId and maxBundleId (exclusive) to determine which bundles are migrated
+var bundlesMigration []BundlesMigrationEntry
+
+func init() {
+	dir, err := merkelRoots.ReadDir("files")
+	if err != nil {
+		panic(err)
+	}
+
+	for _, file := range dir {
+		readFile, err := merkelRoots.ReadFile(fmt.Sprintf("files/%s", file.Name()))
+		if err != nil {
+			panic(err)
+		}
+
+		poolId, err := strconv.ParseUint(strings.ReplaceAll(file.Name(), "merkle_roots_pool_", ""), 10, 64)
+		if err != nil {
+			panic(err)
+		}
+
+		bundlesMigration = append(bundlesMigration, BundlesMigrationEntry{
+			merkleRoots: readFile,
+			poolId:      poolId,
+			maxBundleId: uint64(len(readFile)) / 32,
+		})
+	}
+}
+
+// MigrateBundlesModule migrates the bundles by adding the missing Merkle Roots to the bundle summary.
+func MigrateBundlesModule(sdkCtx sdk.Context, bundlesKeeper bundleskeeper.Keeper, upgradeHeight int64) {
+	logger = sdkCtx.Logger().With("upgrade", "bundles-migration")
+
+	if sdkCtx.BlockHeight()-upgradeHeight < WaitingBlockPeriod {
+		logger.Info("sdkCtx.BlockHeight()-upgradeHeight < WaitingBlockPeriod > return")
+		return
+	}
+
+	for _, bundlesMigrationEntry := range bundlesMigration {
+		step := uint64(sdkCtx.BlockHeight()-upgradeHeight) - WaitingBlockPeriod
+		offset := step * BundlesMigrationStepSizePerPool
+
+		// Skip if all bundles have already been migrated
+		if offset > bundlesMigrationEntry.maxBundleId+BundlesMigrationStepSizePerPool {
+			continue
+		}
+
+		if err := migrateFinalizedBundles(sdkCtx, bundlesKeeper, offset, bundlesMigrationEntry); err != nil {
+			panic(err)
+		}
+	}
+}
+
+// migrateFinalizedBundles sets the updated bundles for a certain range.
+func migrateFinalizedBundles(ctx sdk.Context, bundlesKeeper bundleskeeper.Keeper, offset uint64, bundlesMigrationEntry BundlesMigrationEntry) error {
+	// Init Bundles Store
+	storeAdapter := runtime.KVStoreAdapter(bundlesKeeper.Migration_GetStoreService().OpenKVStore(ctx))
+	store := prefix.NewStore(storeAdapter, util.GetByteKey(types.FinalizedBundlePrefix, bundlesMigrationEntry.poolId))
+
+	var iterator storeTypes.Iterator
+	iterator = store.Iterator(util.GetByteKey(offset), util.GetByteKey(offset+BundlesMigrationStepSizePerPool))
+
+	var migratedBundles []types.FinalizedBundle
+
+	for ; iterator.Valid(); iterator.Next() {
+		var rawFinalizedBundle types.FinalizedBundle
+		if err := bundlesKeeper.Migration_GetCodec().Unmarshal(iterator.Value(), &rawFinalizedBundle); err != nil {
+			return err
+		}
+
+		if rawFinalizedBundle.Id >= bundlesMigrationEntry.maxBundleId {
+			break
+		}
+
+		merkleRoot := bundlesMigrationEntry.merkleRoots[rawFinalizedBundle.Id*32 : rawFinalizedBundle.Id*32+32]
+
+		rawFinalizedBundle.BundleSummary = fmt.Sprintf("{\"merkle_root\":\"%v\"}", hex.EncodeToString(merkleRoot))
+
+		migratedBundles = append(migratedBundles, rawFinalizedBundle)
+	}
+	iterator.Close()
+
+	for _, migratedBundle := range migratedBundles {
+		bundlesKeeper.SetFinalizedBundle(ctx, migratedBundle)
+	}
+	return nil
+}

x/bundles/module.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"github.com/KYVENetwork/chain/x/bundles/migration"
 
 	"cosmossdk.io/core/appmodule"
 	"cosmossdk.io/core/store"
@@ -175,6 +176,12 @@ func (am AppModule) BeginBlock(ctx context.Context) error {
 	sdkCtx := sdk.UnwrapSDKContext(ctx)
 	am.keeper.InitMemStore(sdkCtx)
 	SplitInflation(sdkCtx, am.keeper, am.bankKeeper, am.mintKeeper, am.poolKeeper, am.teamKeeper, am.upgradeKeeper)
+
+	upgradeHeight, err := am.keeper.GetBundlesMigrationUpgradeHeight(sdkCtx)
+	if err == nil {
+		migration.MigrateBundlesModule(sdkCtx, am.keeper, int64(upgradeHeight))
+	}
+
 	return nil
 }
 

x/bundles/types/keys.go

@@ -27,6 +27,8 @@ var (
 	FinalizedBundleVersionMapKey = []byte{3}
 	// RoundRobinProgressPrefix ...
 	RoundRobinProgressPrefix = []byte{4}
+	// BundlesMigrationHeightKey ...
+	BundlesMigrationHeightKey = []byte{5}
 
 	FinalizedBundleByIndexPrefix = []byte{11}
 )