migration impl

Signed-off-by: Joshua Kim <[email protected]>

Author: joshua-kim

vms/avm/block/executor/block.go

@@ -201,7 +201,7 @@ func (b *Block) Verify(context.Context) error {
 	return nil
 }
 
-func (b *Block) Accept(context.Context) error {
+func (b *Block) Accept(ctx context.Context) error {
 	blkID := b.ID()
 	defer b.manager.free(blkID)
 
@@ -246,7 +246,7 @@ func (b *Block) Accept(context.Context) error {
 		return err
 	}
 
-	txChecksum, utxoChecksum := b.manager.state.Checksums()
+	txChecksum, utxoChecksum := b.manager.state.Checksums(ctx)
 	b.manager.backend.Ctx.Log.Trace(
 		"accepted block",
 		zap.Stringer("blkID", blkID),

vms/avm/state/state.go

@@ -4,8 +4,10 @@
 package state
 
 import (
+	"context"
 	"errors"
 	"fmt"
+	"iter"
 	"time"
 
 	"github.com/prometheus/client_golang/prometheus"
@@ -29,6 +31,7 @@ const (
 
 var (
 	utxoPrefix      = []byte("utxo")
+	indexPrefix     = []byte("index")
 	txPrefix        = []byte("tx")
 	blockIDPrefix   = []byte("blockID")
 	blockPrefix     = []byte("block")
@@ -90,7 +93,11 @@ type State interface {
 	CommitBatch() (database.Batch, error)
 
 	// Checksums returns the current TxChecksum and UTXOChecksum.
-	Checksums() (txChecksum ids.ID, utxoChecksum ids.ID)
+	Checksums(ctx context.Context) (txChecksum ids.ID, utxoChecksum ids.ID, err error)
+
+	UTXOs() iter.Seq2[*avax.UTXO, error]
+	Txs() iter.Seq2[*txs.Tx, error]
+	Blocks() iter.Seq2[block.Block, error]
 
 	Close() error
 }
@@ -150,7 +157,30 @@ func New(
 	blockIDDB := prefixdb.New(blockIDPrefix, db)
 	blockDB := prefixdb.New(blockPrefix, db)
 	singletonDB := prefixdb.New(singletonPrefix, db)
+	return newWithDatabases(
+		db,
+		utxoDB,
+		txDB,
+		blockIDDB,
+		blockDB,
+		singletonDB,
+		parser,
+		metrics,
+		trackChecksums,
+	)
+}
 
+func newWithDatabases(
+	db *versiondb.Database,
+	utxoDB database.Database,
+	txDB database.Database,
+	blockIDDB database.Database,
+	blockDB database.Database,
+	singletonDB database.Database,
+	parser block.Parser,
+	metrics prometheus.Registerer,
+	trackChecksums bool,
+) (*state, error) {
 	txCache, err := metercacher.New[ids.ID, *txs.Tx](
 		"tx_cache",
 		metrics,
@@ -178,7 +208,13 @@ func New(
 		return nil, err
 	}
 
-	utxoState, err := avax.NewMeteredUTXOState(utxoDB, parser.Codec(), metrics, trackChecksums)
+	utxoState, err := avax.NewMeteredUTXOState(
+		prefixdb.New(utxoPrefix, utxoDB),
+		prefixdb.New(indexPrefix, utxoDB),
+		parser.Codec(),
+		metrics,
+		trackChecksums,
+	)
 	if err != nil {
 		return nil, err
 	}
@@ -407,6 +443,46 @@ func (s *state) CommitBatch() (database.Batch, error) {
 	return s.db.CommitBatch()
 }
 
+func (s *state) UTXOs() iter.Seq2[*avax.UTXO, error] {
+	return s.utxoState.UTXOs()
+}
+
+func (s *state) Txs() iter.Seq2[*txs.Tx, error] {
+	return func(yield func(*txs.Tx, error) bool) {
+		itr := s.txDB.NewIterator()
+		defer itr.Release()
+
+		for itr.Next() {
+			tx, err := s.parser.ParseGenesisTx(itr.Value())
+			if err != nil {
+				return
+			}
+
+			if !yield(tx, errors.Join(err, itr.Error())) {
+				return
+			}
+		}
+	}
+}
+
+func (s *state) Blocks() iter.Seq2[block.Block, error] {
+	return func(yield func(block.Block, error) bool) {
+		itr := s.blockDB.NewIterator()
+		defer itr.Release()
+
+		for itr.Next() {
+			blk, err := s.parser.ParseBlock(itr.Value())
+			if err != nil {
+				return
+			}
+
+			if !yield(blk, errors.Join(err, itr.Error())) {
+				return
+			}
+		}
+	}
+}
+
 func (s *state) Close() error {
 	return errors.Join(
 		s.utxoDB.Close(),
@@ -500,8 +576,8 @@ func (s *state) writeMetadata() error {
 	return nil
 }
 
-func (s *state) Checksums() (ids.ID, ids.ID) {
-	return s.txChecksum, s.utxoState.Checksum()
+func (s *state) Checksums(context.Context) (ids.ID, ids.ID, error) {
+	return s.txChecksum, s.utxoState.Checksum(), nil
 }
 
 func (s *state) initTxChecksum() error {

vms/avm/state/state_test.go

@@ -4,17 +4,22 @@
 package state
 
 import (
+	"context"
+	"fmt"
 	"testing"
 	"time"
 
+	"github.com/google/go-cmp/cmp/cmpopts"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/stretchr/testify/require"
+	"gotest.tools/v3/assert/cmp"
 
 	"github.com/ava-labs/avalanchego/database"
 	"github.com/ava-labs/avalanchego/database/memdb"
 	"github.com/ava-labs/avalanchego/database/versiondb"
 	"github.com/ava-labs/avalanchego/ids"
 	"github.com/ava-labs/avalanchego/upgrade"
+	"github.com/ava-labs/avalanchego/utils/logging"
 	"github.com/ava-labs/avalanchego/vms/avm/block"
 	"github.com/ava-labs/avalanchego/vms/avm/fxs"
 	"github.com/ava-labs/avalanchego/vms/avm/txs"
@@ -316,3 +321,69 @@ func TestInitializeChainState(t *testing.T) {
 	require.NoError(err)
 	require.Equal(genesis.ID(), lastAccepted.Parent())
 }
+
+func TestStateMigration(t *testing.T) {
+	require := require.New(t)
+
+	prev, err := New(
+		versiondb.New(memdb.New()),
+		parser,
+		prometheus.NewRegistry(),
+		false,
+	)
+	require.NoError(err)
+
+	prev.AddUTXO(populatedUTXO)
+	prev.AddTx(populatedTx)
+	prev.AddBlock(populatedBlk)
+	require.NoError(prev.Commit())
+
+	next, err := NewStateV2(
+		context.Background(),
+		versiondb.New(memdb.New()),
+		parser,
+		prometheus.NewRegistry(),
+		false,
+	)
+	require.NoError(err)
+
+	require.NoError(Migrate(logging.NoLog{}, prev, next))
+
+	gotUTXO, err := next.GetUTXO(populatedUTXOID)
+	require.NoError(err)
+	c := cmp.DeepEqual(
+		populatedUTXO,
+		gotUTXO,
+		cmpopts.EquateEmpty(),
+		cmpopts.IgnoreUnexported(avax.UTXOID{}, secp256k1fx.OutputOwners{}),
+	)()
+	require.Truef(c.Success(), fmt.Sprintf("%s", c))
+
+	gotTx, err := next.GetTx(populatedTxID)
+	require.NoError(err)
+	c = cmp.DeepEqual(
+		populatedTx,
+		gotTx,
+		cmpopts.EquateEmpty(),
+		cmpopts.IgnoreUnexported(txs.Tx{}, txs.BaseTx{}),
+	)()
+	require.Truef(c.Success(), fmt.Sprintf("%s", c))
+
+	gotBlk, err := next.GetBlock(populatedBlkID)
+	require.NoError(err)
+	c = cmp.DeepEqual(
+		populatedBlk,
+		gotBlk,
+		cmpopts.EquateEmpty(),
+		cmpopts.IgnoreUnexported(block.StandardBlock{}, txs.Tx{}, txs.BaseTx{}),
+	)()
+	require.Truef(c.Success(), fmt.Sprintf("%s", c))
+
+	gotBlkID, err := next.GetBlockIDAtHeight(populatedBlkHeight)
+	require.NoError(err)
+	require.Equal(populatedBlkID, gotBlkID)
+
+	require.Equal(prev.GetLastAccepted(), next.GetLastAccepted())
+	require.Equal(prev.GetTimestamp(), next.GetTimestamp())
+	require.True(next.IsInitialized())
+}