Archive Node Online Migration

app/test_state_store.go

@@ -244,6 +244,26 @@ func (s *InMemoryStateStore) RawImport(ch <-chan types.RawSnapshotNode) error {
 	return nil
 }
 
+func (s *InMemoryStateStore) SetLatestMigratedModule(module string) error {
+	// TODO: Add set call here
+	return nil
+}
+
+func (s *InMemoryStateStore) GetLatestMigratedModule() (string, error) {
+	// TODO: Add get call here
+	return "", nil
+}
+
+func (s *InMemoryStateStore) SetLatestMigratedKey(key []byte) error {
+	// TODO: Add set call here
+	return nil
+}
+
+func (s *InMemoryStateStore) GetLatestMigratedKey() ([]byte, error) {
+	// TODO: Add get call here
+	return nil, nil
+}
+
 func (s *InMemoryStateStore) Prune(version int64) error {
 	s.mu.Lock()
 	defer s.mu.Unlock()

go.mod

@@ -351,7 +351,7 @@ replace (
 	github.com/cosmos/ibc-go/v3 => github.com/sei-protocol/sei-ibc-go/v3 v3.3.2
 	github.com/ethereum/go-ethereum => github.com/sei-protocol/go-ethereum v1.13.5-sei-9.0.20240923025222-815b87dde97b
 	github.com/gogo/protobuf => github.com/regen-network/protobuf v1.3.3-alpha.regen.1
-	github.com/sei-protocol/sei-db => github.com/sei-protocol/sei-db v0.0.44
+	github.com/sei-protocol/sei-db => github.com/sei-protocol/sei-db v0.0.45-0.20240918104613-6c0900823891
 	// Latest goleveldb is broken, we have to stick to this version
 	github.com/syndtr/goleveldb => github.com/syndtr/goleveldb v1.0.1-0.20210819022825-2ae1ddf74ef7
 	github.com/tendermint/tendermint => github.com/sei-protocol/sei-tendermint v0.3.8

tools/migration/ss/migrator.go

@@ -3,6 +3,7 @@ package ss
 import (
 	"bytes"
 	"fmt"
+	"time"
 
 	"github.com/cosmos/iavl"
 	"github.com/sei-protocol/sei-db/config"
@@ -43,36 +44,42 @@ func NewMigrator(homeDir string, db dbm.DB) *Migrator {
 }
 
 func (m *Migrator) Migrate(version int64, homeDir string) error {
-	// TODO: Read in capacity of this buffered channel as param
 	ch := make(chan types.RawSnapshotNode, 1000)
 	errCh := make(chan error, 2)
 
-	fmt.Println("Beginning Migration...")
+	// Get the latest key, if any, to resume from
+	latestKey, err := m.stateStore.GetLatestMigratedKey()
+	if err != nil {
+		return fmt.Errorf("failed to get latest key: %w", err)
+	}
+
+	latestModule, err := m.stateStore.GetLatestMigratedModule()
+	if err != nil {
+		return fmt.Errorf("failed to get latest module: %w", err)
+	}
 
-	// Goroutine to iterate through iavl and export leaf nodes
+	fmt.Println("Starting migration...")
+
+	// Goroutine to iterate through IAVL and export leaf nodes
 	go func() {
 		defer close(ch)
-		errCh <- ExportLeafNodes(m.iavlDB, ch)
+		errCh <- ExportLeafNodesFromKey(m.iavlDB, ch, latestKey, latestModule)
 	}()
 
+	// Import nodes into PebbleDB
 	go func() {
 		errCh <- m.stateStore.RawImport(ch)
 	}()
 
-	// Block on completion of both goroutines
+	// Block until both processes complete
 	for i := 0; i < 2; i++ {
 		if err := <-errCh; err != nil {
 			return err
 		}
 	}
 
-	// Set latest version
-	err := m.stateStore.SetLatestVersion(version)
-	if err != nil {
-		return err
-	}
-
-	return nil
+	// Set latest version in the database
+	return m.stateStore.SetLatestVersion(version)
 }
 
 func (m *Migrator) Verify(version int64) error {
@@ -113,40 +120,52 @@ func (m *Migrator) Verify(version int64) error {
 	return verifyErr
 }
 
-// Export leaf nodes of iavl
-func ExportLeafNodes(db dbm.DB, ch chan<- types.RawSnapshotNode) error {
-	// Module by module, TODO: Potentially parallelize
+func ExportLeafNodesFromKey(db dbm.DB, ch chan<- types.RawSnapshotNode, startKey []byte, startModule string) error {
 	count := 0
 	leafNodeCount := 0
-	fmt.Println("Scanning database and exporting leaf nodes...")
+	fmt.Println("ExportLeafNodesFromKey - Scanning database and exporting leaf nodes...")
+
+	startTimeTotal := time.Now() // Start measuring total time
 
 	for _, module := range modules {
-		fmt.Printf("Iterating through %s module...\n", module)
+		if module != startModule && startModule != "" {
+			continue
+		}
+		startTimeModule := time.Now() // Measure time for each module
+		fmt.Printf("ExportLeafNodesFromKey - Iterating through %s module...\n", module)
 
-		// Can't use the previous, have to create an inner
 		prefixDB := dbm.NewPrefixDB(db, []byte(buildRawPrefix(module)))
-		itr, err := prefixDB.Iterator(nil, nil)
+		var itr dbm.Iterator
+		var err error
+
+		// If there is a starting key, seek to it, otherwise start from the beginning
+		if startKey != nil && bytes.HasPrefix(startKey, []byte(buildRawPrefix(module))) {
+			itr, err = prefixDB.Iterator(startKey, nil) // Start from the latest key
+		} else {
+			itr, err = prefixDB.Iterator(nil, nil) // Start from the beginning
+		}
+
 		if err != nil {
-			fmt.Printf("error Export Leaf Nodes %+v\n", err)
+			fmt.Printf("ExportLeafNodesFromKey - Error creating iterator: %+v\n", err)
 			return fmt.Errorf("failed to create iterator: %w", err)
 		}
 		defer itr.Close()
 
+		startTimeBatch := time.Now() // Measure time for every 10,000 iterations
+
 		for ; itr.Valid(); itr.Next() {
 			value := bytes.Clone(itr.Value())
 
 			node, err := iavl.MakeNode(value)
-
 			if err != nil {
-				fmt.Printf("failed to make node err: %+v\n", err)
+				fmt.Printf("ExportLeafNodesFromKey - Failed to make node: %+v\n", err)
 				return fmt.Errorf("failed to make node: %w", err)
 			}
 
-			// leaf node
+			// Only export leaf nodes
 			if node.GetHeight() == 0 {
 				leafNodeCount++
 				ch <- types.RawSnapshotNode{
-					// TODO: Likely need to clone
 					StoreKey: module,
 					Key:      node.GetNodeKey(),
 					Value:    node.GetValue(),
@@ -156,20 +175,25 @@ func ExportLeafNodes(db dbm.DB, ch chan<- types.RawSnapshotNode) error {
 
 			count++
 			if count%10000 == 0 {
-				fmt.Printf("Total scanned: %d, leaf nodes exported: %d\n", count, leafNodeCount)
+				batchDuration := time.Since(startTimeBatch)
+				fmt.Printf("ExportLeafNodesFromKey - Last 10,000 iterations took: %v. Total scanned: %d, leaf nodes exported: %d\n", batchDuration, count, leafNodeCount)
+
+				startTimeBatch = time.Now() // Reset the start time for the next batch
 			}
 		}
 
-		fmt.Printf("Finished scanning module %s Total scanned: %d, leaf nodes exported: %d\n", module, count, leafNodeCount)
-
 		if err := itr.Error(); err != nil {
-			fmt.Printf("iterator error: %+v\n", err)
+			fmt.Printf("Iterator error: %+v\n", err)
 			return fmt.Errorf("iterator error: %w", err)
 		}
 
+		moduleDuration := time.Since(startTimeModule)
+		fmt.Printf("ExportLeafNodesFromKey - Finished scanning module %s. Time taken: %v. Total scanned: %d, leaf nodes exported: %d\n", module, moduleDuration, count, leafNodeCount)
 	}
 
-	fmt.Printf("DB contains %d entries, exported %d leaf nodes\n", count, leafNodeCount)
+	totalDuration := time.Since(startTimeTotal)
+	fmt.Printf("ExportLeafNodesFromKey - DB scanning completed. Total time taken: %v. Total entries scanned: %d, leaf nodes exported: %d\n", totalDuration, count, leafNodeCount)
+
 	return nil
 }