[merkledb] add support for snapshots

x/merkledb/db.go

@@ -157,6 +157,7 @@ type MerkleDB interface {
 	ChangeProofer
 	RangeProofer
 	Prefetcher
+	NewSnapshot(revisionsLifetime int) (ReadOnlyTrie, error)
 }
 
 type Config struct {
@@ -233,6 +234,9 @@ type merkleDB struct {
 	// Valid children of this trie.
 	childViews []*view
 
+	// Snapshots of this trie.
+	snapshots []*snapshot
+
 	// hashNodesKeyPool controls the number of goroutines that are created
 	// inside [hashChangedNode] at any given time and provides slices for the
 	// keys needed while hashing.
@@ -973,6 +977,10 @@ func (db *merkleDB) commitView(ctx context.Context, trieToCommit *view) error {
 	))
 	defer span.End()
 
+	if err := db.maskChangesInSnapshots(changes); err != nil {
+		return err
+	}
+
 	// invalidate all child views except for the view being committed
 	db.invalidateChildrenExcept(trieToCommit)
 
@@ -1206,6 +1214,68 @@ func (db *merkleDB) VerifyChangeProof(
 	return nil
 }
 
+// NewSnapshot returns a snapshot of this database that can provide the key/values of the database exactly as it was when NewSnapshot was called.
+// After a revisionsLifetime number of new revisions have been committed to the database, the ReadOnlyTrie will become invalid
+//
+// Assumes [db.lock] isn't held.
+func (db *merkleDB) NewSnapshot(revisionsLifetime int) (ReadOnlyTrie, error) {
+	db.lock.Lock()
+	defer db.lock.Unlock()
+
+	baseView, err := newView(db, db, ViewChanges{})
+	if err != nil {
+		return nil, err
+	}
+	snap := &snapshot{
+		innermostParent: baseView,
+		innerView:       baseView,
+		revisionsLeft:   revisionsLifetime,
+	}
+	db.snapshots = append(db.snapshots, snap)
+	return snap, nil
+}
+
+// maskChangesInSnapshots inserts a new view in all snapshots' parent trie chains that masks out the new changes from the db
+// additionally cleans up old snapshots once they have hit their revision limit
+// Assumes [db.lock] is held.
+func (db *merkleDB) maskChangesInSnapshots(changes *changeSummary) error {
+	// clean up all snapshots
+	for i := 0; i < len(db.snapshots); i++ {
+		for db.snapshots[i].revisionsLeft == 0 {
+			db.snapshots[i] = db.snapshots[len(db.snapshots)-1]
+			db.snapshots = db.snapshots[:len(db.snapshots)-1]
+		}
+	}
+
+	// don't construct the reversed changes view if it isn't needed
+	if len(db.snapshots) == 0 {
+		return nil
+	}
+
+	// create a new view that contains the
+	reversedChanges := &changeSummary{
+		rootID: changes.rootID,
+		values: make(map[Key]*change[maybe.Maybe[[]byte]], len(changes.values)),
+		nodes:  make(map[Key]*change[*node], len(changes.nodes)),
+	}
+	reversedChanges.rootChange = change[maybe.Maybe[*node]]{before: changes.rootChange.after, after: changes.rootChange.before}
+	for key, currentChange := range changes.values {
+		reversedChanges.values[key] = &change[maybe.Maybe[[]byte]]{before: currentChange.after, after: currentChange.before}
+	}
+	for key, currentChange := range changes.nodes {
+		reversedChanges.nodes[key] = &change[*node]{before: currentChange.after, after: currentChange.before}
+	}
+	newParentView, err := newViewWithChanges(db, reversedChanges)
+	if err != nil {
+		return err
+	}
+
+	for i := 0; i < len(db.snapshots); i++ {
+		db.snapshots[i].updateParent(newParentView)
+	}
+	return nil
+}
+
 // Invalidates and removes any child views that aren't [exception].
 // Assumes [db.lock] is held.
 func (db *merkleDB) invalidateChildrenExcept(exception *view) {

x/merkledb/snapshot.go

@@ -0,0 +1,33 @@
+// Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package merkledb
+
+import "context"
+
+var _ ReadOnlyTrie = (*snapshot)(nil)
+
+type snapshot struct {
+	revisionsLeft   int
+	innerView       *view
+	innermostParent *view
+}
+
+func (s *snapshot) GetValue(ctx context.Context, key []byte) ([]byte, error) {
+	return s.innerView.GetValue(ctx, key)
+}
+
+func (s *snapshot) GetValues(ctx context.Context, keys [][]byte) ([][]byte, []error) {
+	return s.innerView.GetValues(ctx, keys)
+}
+
+// maskingChanges applies the changes in the changeSummary in reverse so that the snapshot stays consistent even though the underlying db has changed
+func (s *snapshot) updateParent(v *view) {
+	s.revisionsLeft--
+	if s.revisionsLeft == 0 {
+		s.innermostParent.invalidate()
+		return
+	}
+	s.innermostParent.updateParent(v)
+	s.innermostParent = v
+}

x/merkledb/snapshot_test.go

@@ -0,0 +1,103 @@
+// Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package merkledb
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/ava-labs/avalanchego/database"
+)
+
+func Test_Snapshot_Basic(t *testing.T) {
+	require := require.New(t)
+
+	db, err := getBasicDB()
+	require.NoError(err)
+	ctx := context.Background()
+	view, err := db.NewView(
+		ctx,
+		ViewChanges{
+			BatchOps: []database.BatchOp{
+				{Key: []byte{0}, Value: []byte{0}},
+			},
+		},
+	)
+	require.NoError(err)
+	require.NoError(view.CommitToDB(ctx))
+
+	view, err = db.NewView(
+		ctx,
+		ViewChanges{
+			BatchOps: []database.BatchOp{
+				{Key: []byte{1}, Value: []byte{1}},
+			},
+		},
+	)
+	require.NoError(err)
+	require.NoError(view.CommitToDB(ctx))
+
+	snap, err := db.NewSnapshot(3)
+	require.NoError(err)
+
+	// confirm that the snapshot has the expected values
+	val, err := snap.GetValue(ctx, []byte{0})
+	require.NoError(err)
+	require.Equal([]byte{0}, val)
+
+	val, err = snap.GetValue(ctx, []byte{1})
+	require.NoError(err)
+	require.Equal([]byte{1}, val)
+
+	// commit new key/value
+	view, err = db.NewView(
+		ctx,
+		ViewChanges{
+			BatchOps: []database.BatchOp{
+				{Key: []byte{3}, Value: []byte{3}},
+			},
+		},
+	)
+	require.NoError(err)
+	require.NoError(view.CommitToDB(ctx))
+
+	// the snapshot should not contain the new value
+	_, err = snap.GetValue(ctx, []byte{3})
+	require.ErrorIs(err, database.ErrNotFound)
+
+	// write over existing key values
+	view, err = db.NewView(
+		ctx,
+		ViewChanges{
+			BatchOps: []database.BatchOp{
+				{Key: []byte{1}, Value: []byte{4}},
+			},
+		},
+	)
+	require.NoError(err)
+	require.NoError(view.CommitToDB(ctx))
+
+	// should still have the old value
+	val, err = snap.GetValue(ctx, []byte{1})
+	require.NoError(err)
+	require.Equal([]byte{1}, val)
+
+	// commit more to trigger view invalidation
+	view, err = db.NewView(
+		ctx,
+		ViewChanges{
+			BatchOps: []database.BatchOp{
+				{Key: []byte{4}, Value: []byte{4}},
+			},
+		},
+	)
+	require.NoError(err)
+	require.NoError(view.CommitToDB(ctx))
+
+	// too many revisions have passed and now the snapshot is invalid
+	_, err = snap.GetValue(ctx, []byte{1})
+	require.ErrorIs(err, ErrInvalid)
+}

x/merkledb/trie.go

@@ -60,14 +60,7 @@ type Trie interface {
 	MerkleRootGetter
 	ProofGetter
 	database.Iteratee
-
-	// GetValue gets the value associated with the specified key
-	// database.ErrNotFound if the key is not present
-	GetValue(ctx context.Context, key []byte) ([]byte, error)
-
-	// GetValues gets the values associated with the specified keys
-	// database.ErrNotFound if the key is not present
-	GetValues(ctx context.Context, keys [][]byte) ([][]byte, []error)
+	ReadOnlyTrie
 
 	// GetRangeProof returns a proof of up to [maxLength] key-value pairs with
 	// keys in range [start, end].
@@ -84,6 +77,16 @@ type Trie interface {
 	) (View, error)
 }
 
+type ReadOnlyTrie interface {
+	// GetValue gets the value associated with the specified key
+	// database.ErrNotFound if the key is not present
+	GetValue(ctx context.Context, key []byte) ([]byte, error)
+
+	// GetValues gets the values associated with the specified keys
+	// database.ErrNotFound if the key is not present
+	GetValues(ctx context.Context, keys [][]byte) ([][]byte, []error)
+}
+
 type View interface {
 	Trie