feat(store/v2): snapshot manager

store/changeset.go

@@ -4,8 +4,9 @@ package store
 // track writes. Deletion can be denoted by a nil value or explicitly by the
 // Delete field.
 type KVPair struct {
-	Key   []byte
-	Value []byte
+	Key      []byte
+	Value    []byte
+	StoreKey string // Optional for snapshot restore
 }
 
 type KVPairs []KVPair

store/commitment/iavl/tree.go

@@ -1,6 +1,7 @@
 package iavl
 
 import (
+	"errors"
 	"fmt"
 
 	dbm "github.com/cosmos/cosmos-db"
@@ -9,6 +10,7 @@ import (
 
 	log "cosmossdk.io/log"
 	"cosmossdk.io/store/v2/commitment"
+	snapshotstypes "cosmossdk.io/store/v2/snapshots/types"
 )
 
 var _ commitment.Tree = (*IavlTree)(nil)
@@ -77,7 +79,84 @@ func (t *IavlTree) Prune(version uint64) error {
 	return t.tree.DeleteVersionsTo(int64(version))
 }
 
+// Export exports the tree exporter at the given version.
+func (t *IavlTree) Export(version uint64) (commitment.Exporter, error) {
+	tree, err := t.tree.GetImmutable(int64(version))
+	if err != nil {
+		return nil, err
+	}
+	exporter, err := tree.Export()
+	return &Exporter{
+		exporter: exporter,
+	}, err
+}
+
+// Import imports the tree importer at the given version.
+func (t *IavlTree) Import(version uint64) (commitment.Importer, error) {
+	importer, err := t.tree.Import(int64(version))
+	return &Importer{
+		importer: importer,
+	}, err
+}
+
 // Close closes the iavl tree.
 func (t *IavlTree) Close() error {
 	return nil
 }
+
+// Exporter is a wrapper around iavl.Exporter.
+type Exporter struct {
+	exporter *iavl.Exporter
+}
+
+// Next returns the next item in the exporter.
+func (e *Exporter) Next() (*snapshotstypes.SnapshotIAVLItem, error) {
+	item, err := e.exporter.Next()
+	if err != nil {
+		if errors.Is(err, iavl.ErrorExportDone) {
+			return nil, commitment.ErrorExportDone
+		}
+		return nil, err
+	}
+
+	return &snapshotstypes.SnapshotIAVLItem{
+		Key:     item.Key,
+		Value:   item.Value,
+		Version: item.Version,
+		Height:  int32(item.Height),
+	}, nil
+}
+
+// Close closes the exporter.
+func (e *Exporter) Close() error {
+	e.exporter.Close()
+
+	return nil
+}
+
+// Importer is a wrapper around iavl.Importer.
+type Importer struct {
+	importer *iavl.Importer
+}
+
+// Add adds the given item to the importer.
+func (i *Importer) Add(item *snapshotstypes.SnapshotIAVLItem) error {
+	return i.importer.Add(&iavl.ExportNode{
+		Key:     item.Key,
+		Value:   item.Value,
+		Version: item.Version,
+		Height:  int8(item.Height),
+	})
+}
+
+// Commit commits the importer.
+func (i *Importer) Commit() error {
+	return i.importer.Commit()
+}
+
+// Close closes the importer.
+func (i *Importer) Close() error {
+	i.importer.Close()
+
+	return nil
+}

store/commitment/iavl/tree_test.go

@@ -9,15 +9,15 @@ import (
 	"cosmossdk.io/log"
 )
 
-func generateTree(treeType string) *IavlTree {
+func generateTree() *IavlTree {
 	cfg := DefaultConfig()
 	db := dbm.NewMemDB()
 	return NewIavlTree(db, log.NewNopLogger(), cfg)
 }
 
 func TestIavlTree(t *testing.T) {
 	// generate a new tree
-	tree := generateTree("iavl")
+	tree := generateTree()
 	require.NotNil(t, tree)
 
 	initVersion := tree.GetLatestVersion()

store/commitment/store.go

@@ -3,14 +3,19 @@ package commitment
 import (
 	"errors"
 	"fmt"
+	"io"
+	"math"
 
+	protoio "github.com/cosmos/gogoproto/io"
 	ics23 "github.com/cosmos/ics23/go"
 
 	"cosmossdk.io/log"
 	"cosmossdk.io/store/v2"
+	snapshotstypes "cosmossdk.io/store/v2/snapshots/types"
 )
 
 var _ store.Committer = (*CommitStore)(nil)
+var _ snapshotstypes.CommitSnapshotter = (*CommitStore)(nil)
 
 // CommitStore is a wrapper around multiple Tree objects mapped by a unique store
 // key. Each store key reflects dedicated and unique usage within a module. A caller
@@ -127,6 +132,140 @@ func (c *CommitStore) Prune(version uint64) (ferr error) {
 	return ferr
 }
 
+// Snapshot implements snapshotstypes.CommitSnapshotter.
+func (c *CommitStore) Snapshot(version uint64, protoWriter protoio.Writer) error {
+	if version == 0 {
+		return fmt.Errorf("the snapshot version must be greater than 0")
+	}
+
+	latestVersion, err := c.GetLatestVersion()
+	if err != nil {
+		return err
+	}
+	if version > latestVersion {
+		return fmt.Errorf("the snapshot version %d is greater than the latest version %d", version, latestVersion)
+	}
+
+	for storeKey, tree := range c.multiTrees {
+		exporter, err := tree.Export(version)
+		if err != nil {
+			return fmt.Errorf("failed to export tree for version %d: %w", version, err)
+		}
+
+		defer exporter.Close()
+
+		err = protoWriter.WriteMsg(&snapshotstypes.SnapshotItem{
+			Item: &snapshotstypes.SnapshotItem_Store{
+				Store: &snapshotstypes.SnapshotStoreItem{
+					Name: storeKey,
+				},
+			},
+		})
+		if err != nil {
+			return fmt.Errorf("failed to write store name: %w", err)
+		}
+
+		for {
+			item, err := exporter.Next()
+			if errors.Is(err, ErrorExportDone) {
+				break
+			} else if err != nil {
+				return fmt.Errorf("failed to get the next export node: %w", err)
+			}
+
+			if err = protoWriter.WriteMsg(&snapshotstypes.SnapshotItem{
+				Item: &snapshotstypes.SnapshotItem_IAVL{
+					IAVL: item,
+				},
+			}); err != nil {
+				return fmt.Errorf("failed to write iavl node: %w", err)
+			}
+		}
+	}
+
+	return nil
+}
+
+// Restore implements snapshotstypes.CommitSnapshotter.
+func (c *CommitStore) Restore(version uint64, format uint32, protoReader protoio.Reader, chStorage chan<- *store.KVPair) (snapshotstypes.SnapshotItem, error) {
+	var (
+		importer     Importer
+		snapshotItem snapshotstypes.SnapshotItem
+		storeKey     string
+	)
+
+loop:
+	for {
+		snapshotItem = snapshotstypes.SnapshotItem{}
+		err := protoReader.ReadMsg(&snapshotItem)
+		if errors.Is(err, io.EOF) {
+			break
+		} else if err != nil {
+			return snapshotstypes.SnapshotItem{}, fmt.Errorf("invalid protobuf message: %w", err)
+		}
+
+		switch item := snapshotItem.Item.(type) {
+		case *snapshotstypes.SnapshotItem_Store:
+			if importer != nil {
+				if err := importer.Commit(); err != nil {
+					return snapshotstypes.SnapshotItem{}, fmt.Errorf("failed to commit importer: %w", err)
+				}
+				importer.Close()
+			}
+			storeKey = item.Store.Name
+			tree := c.multiTrees[storeKey]
+			if tree == nil {
+				return snapshotstypes.SnapshotItem{}, fmt.Errorf("store %s not found", storeKey)
+			}
+			importer, err = tree.Import(version)
+			if err != nil {
+				return snapshotstypes.SnapshotItem{}, fmt.Errorf("failed to import tree for version %d: %w", version, err)
+			}
+			defer importer.Close()
+
+		case *snapshotstypes.SnapshotItem_IAVL:
+			if importer == nil {
+				return snapshotstypes.SnapshotItem{}, fmt.Errorf("received IAVL node item before store item")
+			}
+			node := item.IAVL
+			if node.Height > int32(math.MaxInt8) {
+				return snapshotstypes.SnapshotItem{}, fmt.Errorf("node height %v cannot exceed %v",
+					item.IAVL.Height, math.MaxInt8)
+			}
+			// Protobuf does not differentiate between []byte{} and nil, but fortunately IAVL does
+			// not allow nil keys nor nil values for leaf nodes, so we can always set them to empty.
+			if node.Key == nil {
+				node.Key = []byte{}
+			}
+			if node.Height == 0 {
+				if node.Value == nil {
+					node.Value = []byte{}
+				}
+				// If the node is a leaf node, it will be written to the storage.
+				chStorage <- &store.KVPair{
+					Key:      node.Key,
+					Value:    node.Value,
+					StoreKey: storeKey,
+				}
+			}
+			err := importer.Add(node)
+			if err != nil {
+				return snapshotstypes.SnapshotItem{}, fmt.Errorf("failed to add node to importer: %w", err)
+			}
+		default:
+			break loop
+		}
+	}
+
+	if importer != nil {
+		if err := importer.Commit(); err != nil {
+			return snapshotstypes.SnapshotItem{}, fmt.Errorf("failed to commit importer: %w", err)
+		}
+	}
+
+	return snapshotItem, c.LoadVersion(version)
+}
+
 func (c *CommitStore) Close() (ferr error) {
 	for _, tree := range c.multiTrees {
 		if err := tree.Close(); err != nil {

store/commitment/store_test.go

@@ -0,0 +1,65 @@
+package commitment
+
+import (
+	"testing"
+)
+
+func TestSnapshotter(t *testing.T) {
+	// generate a new tree
+	// storeKey := "store"
+	// tree := generateTree(storeKey)
+	// require.NotNil(t, tree)
+
+	// latestVersion := uint64(10)
+	// kvCount := 10
+	// for i := uint64(1); i <= latestVersion; i++ {
+	// 	cs := store.NewChangeset()
+	// 	for j := 0; j < kvCount; j++ {
+	// 		key := []byte(fmt.Sprintf("key-%d-%d", i, j))
+	// 		value := []byte(fmt.Sprintf("value-%d-%d", i, j))
+	// 		cs.Add(key, value)
+	// 	}
+	// 	err := tree.WriteBatch(cs)
+	// 	require.NoError(t, err)
+
+	// 	_, err = tree.Commit()
+	// 	require.NoError(t, err)
+	// }
+
+	// latestHash := tree.WorkingHash()
+
+	// // create a snapshot
+	// dummyExtensionItem := snapshotstypes.SnapshotItem{
+	// 	Item: &snapshotstypes.SnapshotItem_Extension{
+	// 		Extension: &snapshotstypes.SnapshotExtensionMeta{
+	// 			Name:   "test",
+	// 			Format: 1,
+	// 		},
+	// 	},
+	// }
+	// target := generateTree("")
+	// chunks := make(chan io.ReadCloser, kvCount*int(latestVersion))
+	// go func() {
+	// 	streamWriter := snapshots.NewStreamWriter(chunks)
+	// 	require.NotNil(t, streamWriter)
+	// 	defer streamWriter.Close()
+	// 	err := tree.Snapshot(latestVersion, streamWriter)
+	// 	require.NoError(t, err)
+	// 	// write an extension metadata
+	// 	err = streamWriter.WriteMsg(&dummyExtensionItem)
+	// 	require.NoError(t, err)
+	// }()
+
+	// streamReader, err := snapshots.NewStreamReader(chunks)
+	// chStorage := make(chan *store.KVPair, 100)
+	// require.NoError(t, err)
+	// nextItem, err := target.Restore(latestVersion, snapshotstypes.CurrentFormat, streamReader, chStorage)
+	// require.NoError(t, err)
+	// require.Equal(t, *dummyExtensionItem.GetExtension(), *nextItem.GetExtension())
+
+	// // check the store key
+	// require.Equal(t, storeKey, target.storeKey)
+
+	// // check the restored tree hash
+	// require.Equal(t, latestHash, target.WorkingHash())
+}

store/commitment/tree.go

@@ -1,11 +1,17 @@
 package commitment
 
 import (
+	"errors"
 	"io"
 
 	ics23 "github.com/cosmos/ics23/go"
+
+	snapshotstypes "cosmossdk.io/store/v2/snapshots/types"
 )
 
+// ErrorExportDone is returned by Exporter.Next() when all items have been exported.
+var ErrorExportDone = errors.New("export is complete")
+
 // Tree is the interface that wraps the basic Tree methods.
 type Tree interface {
 	Set(key, value []byte) error
@@ -16,6 +22,23 @@ type Tree interface {
 	Commit() ([]byte, error)
 	GetProof(version uint64, key []byte) (*ics23.CommitmentProof, error)
 	Prune(version uint64) error
+	Export(version uint64) (Exporter, error)
+	Import(version uint64) (Importer, error)
+
+	io.Closer
+}
+
+// Exporter is the interface that wraps the basic Export methods.
+type Exporter interface {
+	Next() (*snapshotstypes.SnapshotIAVLItem, error)
+
+	io.Closer
+}
+
+// Importer is the interface that wraps the basic Import methods.
+type Importer interface {
+	Add(*snapshotstypes.SnapshotIAVLItem) error
+	Commit() error
 
 	io.Closer
 }