firewood x-chain

Signed-off-by: Joshua Kim <20001595+joshua-kim@users.noreply.github.com>

Author: joshua-kim

database/memdb/db.go

@@ -5,6 +5,7 @@ package memdb
 
 import (
 	"context"
+	"fmt"
 	"slices"
 	"strings"
 	"sync"
@@ -38,6 +39,21 @@ func New() *Database {
 	return NewWithSize(DefaultSize)
 }
 
+// Copy returns a Database with the same key-value pairs as db
+func Copy(db *Database) (*Database, error) {
+	db.lock.Lock()
+	defer db.lock.Unlock()
+
+	result := New()
+	for k, v := range db.db {
+		if err := result.Put([]byte(k), v); err != nil {
+			return nil, fmt.Errorf("failed to insert key: %w", err)
+		}
+	}
+
+	return result, nil
+}
+
 // NewWithSize returns a map pre-allocated to the provided size with the
 // Database interface methods implemented.
 func NewWithSize(size int) *Database {

vms/avm/block/builder/builder.go

@@ -174,5 +174,5 @@ func (b *builder) BuildBlock(context.Context) (snowman.Block, error) {
 		return nil, err
 	}
 
-	return b.manager.NewBlock(statelessBlk), nil
+	return b.manager.NewBlock(statelessBlk, b.backend.Ctx.SharedMemory), nil
 }

vms/avm/block/builder/builder_test.go

@@ -12,6 +12,7 @@ import (
 	"github.com/stretchr/testify/require"
 	"go.uber.org/mock/gomock"
 
+	"github.com/ava-labs/avalanchego/chains/atomic"
 	"github.com/ava-labs/avalanchego/codec"
 	"github.com/ava-labs/avalanchego/codec/codecmock"
 	"github.com/ava-labs/avalanchego/database/memdb"
@@ -291,8 +292,11 @@ func TestBuilderBuildBlock(t *testing.T) {
 				manager.EXPECT().VerifyUniqueInputs(preferredID, gomock.Any()).Return(nil)
 				// Assert created block has one tx, tx1,
 				// and other fields are set correctly.
-				manager.EXPECT().NewBlock(gomock.Any()).DoAndReturn(
-					func(block *block.StandardBlock) snowman.Block {
+				manager.EXPECT().NewBlock(gomock.Any(), gomock.Any()).DoAndReturn(
+					func(
+						block *block.StandardBlock,
+						_ atomic.SharedMemory,
+					) snowman.Block {
 						require.Len(t, block.Transactions, 1)
 						require.Equal(t, tx1, block.Transactions[0])
 						require.Equal(t, preferredHeight+1, block.Height())
@@ -350,8 +354,11 @@ func TestBuilderBuildBlock(t *testing.T) {
 				manager.EXPECT().GetState(preferredID).Return(preferredState, true)
 				manager.EXPECT().VerifyUniqueInputs(preferredID, gomock.Any()).Return(nil)
 				// Assert that the created block has the right timestamp
-				manager.EXPECT().NewBlock(gomock.Any()).DoAndReturn(
-					func(block *block.StandardBlock) snowman.Block {
+				manager.EXPECT().NewBlock(gomock.Any(), gomock.Any()).DoAndReturn(
+					func(
+						block *block.StandardBlock,
+						_ atomic.SharedMemory,
+					) snowman.Block {
 						require.Equal(t, preferredTimestamp.Unix(), block.Timestamp().Unix())
 						return nil
 					},
@@ -422,8 +429,11 @@ func TestBuilderBuildBlock(t *testing.T) {
 				manager.EXPECT().GetState(preferredID).Return(preferredState, true)
 				manager.EXPECT().VerifyUniqueInputs(preferredID, gomock.Any()).Return(nil)
 				// Assert that the created block has the right timestamp
-				manager.EXPECT().NewBlock(gomock.Any()).DoAndReturn(
-					func(block *block.StandardBlock) snowman.Block {
+				manager.EXPECT().NewBlock(gomock.Any(), gomock.Any()).DoAndReturn(
+					func(
+						block *block.StandardBlock,
+						_ atomic.SharedMemory,
+					) snowman.Block {
 						require.Equal(t, now.Unix(), block.Timestamp().Unix())
 						return nil
 					},
@@ -526,7 +536,7 @@ func TestBlockBuilderAddLocalTx(t *testing.T) {
 	parentBlk, err := block.NewStandardBlock(parentID, 0, parentTimestamp, txs, cm)
 	require.NoError(err)
 	state.AddBlock(parentBlk)
-	state.SetLastAccepted(parentBlk.ID())
+	state.SetLastAccepted(parentBlk.ID(), parentBlk.Height())
 
 	metrics, err := metrics.New(registerer)
 	require.NoError(err)

vms/avm/block/executor/block.go

@@ -36,7 +36,8 @@ var (
 // Exported for testing in avm package.
 type Block struct {
 	block.Block
-	manager *manager
+	manager      *manager
+	sharedMemory atomic.SharedMemory
 }
 
 func (b *Block) Verify(context.Context) error {
@@ -201,15 +202,15 @@ func (b *Block) Verify(context.Context) error {
 
 	// Now that the block has been executed, we can add the block data to the
 	// state diff.
-	stateDiff.SetLastAccepted(blkID)
+	stateDiff.SetLastAccepted(blkID, b.Height())
 	stateDiff.AddBlock(b.Block)
 
 	b.manager.blkIDToState[blkID] = blockState
 	b.manager.mempool.Remove(txs...)
 	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)
 
@@ -240,20 +241,25 @@ func (b *Block) Accept(context.Context) error {
 	}
 
 	// Note that this method writes [batch] to the database.
-	if err := b.manager.backend.Ctx.SharedMemory.Apply(blkState.atomicRequests, batch); err != nil {
+	if err := b.sharedMemory.Apply(blkState.atomicRequests, batch); err != nil {
 		return fmt.Errorf("failed to apply state diff to shared memory: %w", err)
 	}
 
 	if err := b.manager.metrics.MarkBlockAccepted(b); err != nil {
 		return err
 	}
 
+	checksum, err := b.manager.state.Checksum(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to get checksum: %w", err)
+	}
+
 	b.manager.backend.Ctx.Log.Trace(
 		"accepted block",
 		zap.Stringer("blkID", blkID),
 		zap.Uint64("height", b.Height()),
 		zap.Stringer("parentID", b.Parent()),
-		zap.Stringer("checksum", b.manager.state.Checksum()),
+		zap.Stringer("checksum", checksum),
 	)
 	return nil
 }

vms/avm/block/executor/block_test.go

@@ -683,6 +683,7 @@ func TestBlockAccept(t *testing.T) {
 							},
 						},
 					},
+					sharedMemory: mockSharedMemory,
 				}
 			},
 			expectedErr: errTest,
@@ -726,6 +727,7 @@ func TestBlockAccept(t *testing.T) {
 							},
 						},
 					},
+					sharedMemory: mockSharedMemory,
 				}
 			},
 			expectedErr: errTest,
@@ -748,7 +750,7 @@ func TestBlockAccept(t *testing.T) {
 				// because we mock the call to shared memory
 				mockManagerState.EXPECT().CommitBatch().Return(nil, nil)
 				mockManagerState.EXPECT().Abort()
-				mockManagerState.EXPECT().Checksum().Return(ids.Empty)
+				mockManagerState.EXPECT().Checksum(gomock.Any()).Return(ids.Empty, nil)
 
 				mockSharedMemory := atomicmock.NewSharedMemory(ctrl)
 				mockSharedMemory.EXPECT().Apply(gomock.Any(), gomock.Any()).Return(nil)
@@ -772,6 +774,7 @@ func TestBlockAccept(t *testing.T) {
 							},
 						},
 					},
+					sharedMemory: mockSharedMemory,
 				}
 			},
 			expectedErr: nil,

vms/avm/block/executor/executormock/manager.go

@@ -31,13 +31,14 @@ type Manager interface {
 
 	// Returns the ID of the most recently accepted block.
 	LastAccepted() ids.ID
+	SetLastAccepted(id ids.ID)
 
 	SetPreference(blkID ids.ID)
 	Preferred() ids.ID
 
 	GetBlock(blkID ids.ID) (snowman.Block, error)
 	GetStatelessBlock(blkID ids.ID) (block.Block, error)
-	NewBlock(block.Block) snowman.Block
+	NewBlock(blk block.Block, sm atomic.SharedMemory) snowman.Block
 
 	// VerifyTx verifies that the transaction can be issued based on the currently
 	// preferred state. This should *not* be used to verify transactions in a block.
@@ -109,6 +110,10 @@ func (m *manager) LastAccepted() ids.ID {
 	return m.lastAccepted
 }
 
+func (m *manager) SetLastAccepted(blkID ids.ID) {
+	m.lastAccepted = blkID
+}
+
 func (m *manager) SetPreference(blockID ids.ID) {
 	m.preferred = blockID
 }
@@ -122,7 +127,7 @@ func (m *manager) GetBlock(blkID ids.ID) (snowman.Block, error) {
 	if err != nil {
 		return nil, err
 	}
-	return m.NewBlock(blk), nil
+	return m.NewBlock(blk, m.backend.Ctx.SharedMemory), nil
 }
 
 func (m *manager) GetStatelessBlock(blkID ids.ID) (block.Block, error) {
@@ -134,10 +139,14 @@ func (m *manager) GetStatelessBlock(blkID ids.ID) (block.Block, error) {
 	return m.state.GetBlock(blkID)
 }
 
-func (m *manager) NewBlock(blk block.Block) snowman.Block {
+func (m *manager) NewBlock(
+	blk block.Block,
+	sm atomic.SharedMemory,
+) snowman.Block {
 	return &Block{
-		Block:   blk,
-		manager: m,
+		Block:        blk,
+		manager:      m,
+		sharedMemory: sm,
 	}
 }
 

vms/avm/block/executor/manager.go

@@ -25,6 +25,7 @@ import (
 	"github.com/ava-labs/avalanchego/vms/avm/block/executor"
 	"github.com/ava-labs/avalanchego/vms/avm/config"
 	"github.com/ava-labs/avalanchego/vms/avm/fxs"
+	"github.com/ava-labs/avalanchego/vms/avm/state"
 	"github.com/ava-labs/avalanchego/vms/avm/txs"
 	"github.com/ava-labs/avalanchego/vms/avm/txs/txstest"
 	"github.com/ava-labs/avalanchego/vms/components/avax"
@@ -113,7 +114,8 @@ func setup(tb testing.TB, c *envConfig) *environment {
 	}
 
 	vm := &VM{
-		Config: vmStaticConfig,
+		Config:         vmStaticConfig,
+		StateMigration: state.NoMigration{},
 	}
 
 	vmDynamicConfig := DefaultConfig
@@ -408,20 +410,20 @@ func issueAndAccept(
 	require *require.Assertions,
 	vm *VM,
 	tx *txs.Tx,
-) {
+) ids.ID {
 	txID, err := vm.issueTxFromRPC(tx)
 	require.NoError(err)
 	require.Equal(tx.ID(), txID)
 
-	buildAndAccept(require, vm, txID)
+	return buildAndAccept(require, vm, txID)
 }
 
 // buildAndAccept expects the context lock not to be held
 func buildAndAccept(
 	require *require.Assertions,
 	vm *VM,
 	txID ids.ID,
-) {
+) ids.ID {
 	msg, err := vm.WaitForEvent(context.Background())
 	require.NoError(err)
 	require.Equal(common.PendingTxs, msg)
@@ -442,4 +444,6 @@ func buildAndAccept(
 	require.NoError(blk.Verify(context.Background()))
 	require.NoError(vm.SetPreference(context.Background(), blk.ID()))
 	require.NoError(blk.Accept(context.Background()))
+
+	return blk.ID()
 }

vms/avm/environment_test.go

@@ -7,6 +7,7 @@ import (
 	"github.com/ava-labs/avalanchego/utils/logging"
 	"github.com/ava-labs/avalanchego/vms"
 	"github.com/ava-labs/avalanchego/vms/avm/config"
+	"github.com/ava-labs/avalanchego/vms/avm/state"
 )
 
 var _ vms.Factory = (*Factory)(nil)
@@ -16,5 +17,9 @@ type Factory struct {
 }
 
 func (f *Factory) New(logging.Logger) (interface{}, error) {
-	return &VM{Config: f.Config}, nil
+	return &VM{
+			Config:         f.Config,
+			StateMigration: &state.FirewoodMigration{CommitFrequency: 1_000},
+		},
+		nil
 }

vms/avm/factory.go

@@ -38,8 +38,9 @@ type diff struct {
 	addedBlockIDs map[uint64]ids.ID      // map of height -> blockID
 	addedBlocks   map[ids.ID]block.Block // map of blockID -> block
 
-	lastAccepted ids.ID
-	timestamp    time.Time
+	lastAccepted       ids.ID
+	lastAcceptedHeight uint64
+	timestamp          time.Time
 }
 
 func NewDiff(
@@ -149,8 +150,9 @@ func (d *diff) GetLastAccepted() ids.ID {
 	return d.lastAccepted
 }
 
-func (d *diff) SetLastAccepted(lastAccepted ids.ID) {
+func (d *diff) SetLastAccepted(lastAccepted ids.ID, height uint64) {
 	d.lastAccepted = lastAccepted
+	d.lastAcceptedHeight = height
 }
 
 func (d *diff) GetTimestamp() time.Time {
@@ -178,6 +180,6 @@ func (d *diff) Apply(state Chain) {
 		state.AddBlock(blk)
 	}
 
-	state.SetLastAccepted(d.lastAccepted)
+	state.SetLastAccepted(d.lastAccepted, d.lastAcceptedHeight)
 	state.SetTimestamp(d.timestamp)
 }