Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

refactor(store): remove defer in the loop and refactor code #23351

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
+155 −104
Open

refactor(store): remove defer in the loop and refactor code #23351

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
8d515ec
update
hoank101 Jan 13, 2025
3f222d8
Merge branch 'main' into refactor/remove-defer-in-the-loop
hoank101 Jan 14, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
142 changes: 88 additions & 54 deletions store/rootmulti/store.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -938,85 +938,119 @@ func (rs *Store) Snapshot(height uint64, protoWriter protoio.Writer) error {
}

// Restore implements snapshottypes.Snapshotter.
// returns next snapshot item and error.
// Returns next snapshot item and error.
func (rs *Store) Restore(
height uint64, format uint32, protoReader protoio.Reader,
) (snapshottypes.SnapshotItem, error) {
// Import nodes into stores. The first item is expected to be a SnapshotItem containing
// a SnapshotStoreItem, telling us which store to import into. The following items will contain
// SnapshotNodeItem (i.e. ExportNode) until we reach the next SnapshotStoreItem or EOF.
var importer *iavltree.Importer
var snapshotItem snapshottypes.SnapshotItem
loop:
defer func() {
if importer != nil {
importer.Close()
}
}()

for {
snapshotItem = snapshottypes.SnapshotItem{}
err := protoReader.ReadMsg(&snapshotItem)
if errors.Is(err, io.EOF) {
break
} else if err != nil {
snapshotItem := snapshottypes.SnapshotItem{}
if err := protoReader.ReadMsg(&snapshotItem); err != nil {
if errors.Is(err, io.EOF) {
break
}
return snapshottypes.SnapshotItem{}, errorsmod.Wrap(err, "invalid protobuf message")
}

switch item := snapshotItem.Item.(type) {
case *snapshottypes.SnapshotItem_Store:
if importer != nil {
err = importer.Commit()
if err != nil {
return snapshottypes.SnapshotItem{}, errorsmod.Wrap(err, "IAVL commit failed")
}
importer.Close()
}
store, ok := rs.GetStoreByName(item.Store.Name).(*iavl.Store)
if !ok || store == nil {
return snapshottypes.SnapshotItem{}, errorsmod.Wrapf(types.ErrLogic, "cannot import into non-IAVL store %q", item.Store.Name)
}
importer, err = store.Import(int64(height))
if err != nil {
return snapshottypes.SnapshotItem{}, errorsmod.Wrap(err, "import failed")
if err := rs.handleStoreItem(height, &importer, item.Store); err != nil {
return snapshottypes.SnapshotItem{}, err
}
defer importer.Close()
// Importer height must reflect the node height (which usually matches the block height, but not always)
rs.logger.Debug("restoring snapshot", "store", item.Store.Name)

case *snapshottypes.SnapshotItem_IAVL:
if importer == nil {
rs.logger.Error("failed to restore; received IAVL node item before store item")
return snapshottypes.SnapshotItem{}, errorsmod.Wrap(types.ErrLogic, "received IAVL node item before store item")
}
if item.IAVL.Height > math.MaxInt8 {
return snapshottypes.SnapshotItem{}, errorsmod.Wrapf(types.ErrLogic, "node height %v cannot exceed %v",
item.IAVL.Height, math.MaxInt8)
}
node := &iavltree.ExportNode{
Key: item.IAVL.Key,
Value: item.IAVL.Value,
Height: int8(item.IAVL.Height),
Version: item.IAVL.Version,
}
// Protobuf does not differentiate between []byte{} as 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 && node.Value == nil {
node.Value = []byte{}
}
err := importer.Add(node)
if err != nil {
return snapshottypes.SnapshotItem{}, errorsmod.Wrap(err, "IAVL node import failed")
if err := rs.handleIAVLItem(importer, item.IAVL); err != nil {
return snapshottypes.SnapshotItem{}, err
}

default:
break loop
// Return unhandled item for caller to process
return rs.finalizeRestore(height, importer, snapshotItem)
}
}

return rs.finalizeRestore(height, importer, snapshottypes.SnapshotItem{})
}

// handleStoreItem processes a store snapshot item and sets up the importer
func (rs *Store) handleStoreItem(height uint64, importer **iavltree.Importer, storeItem *snapshottypes.SnapshotStoreItem) error {
if *importer != nil {
if err := (*importer).Commit(); err != nil {
return errorsmod.Wrap(err, "IAVL commit failed")
}
}

store, ok := rs.GetStoreByName(storeItem.Name).(*iavl.Store)
if !ok || store == nil {
return errorsmod.Wrapf(types.ErrLogic, "cannot import into non-IAVL store %q", storeItem.Name)
}

var err error
*importer, err = store.Import(int64(height))
if err != nil {
return errorsmod.Wrap(err, "import failed")
}

rs.logger.Debug("restoring snapshot", "store", storeItem.Name)
return nil
}

// handleIAVLItem processes an IAVL node snapshot item
func (rs *Store) handleIAVLItem(importer *iavltree.Importer, iavlItem *snapshottypes.SnapshotIAVLItem) error {
if importer == nil {
rs.logger.Error("failed to restore; received IAVL node item before store item")
return errorsmod.Wrap(types.ErrLogic, "received IAVL node item before store item")
}

if iavlItem.Height > math.MaxInt8 {
return errorsmod.Wrapf(types.ErrLogic, "node height %v cannot exceed %v",
iavlItem.Height, math.MaxInt8)
}

node := createExportNode(iavlItem)
if err := importer.Add(node); err != nil {
return errorsmod.Wrap(err, "IAVL node import failed")
}

return nil
}

// createExportNode creates an ExportNode from a SnapshotIAVLItem
func createExportNode(item *snapshottypes.SnapshotIAVLItem) *iavltree.ExportNode {
node := &iavltree.ExportNode{
Key: item.Key,
Value: item.Value,
Height: int8(item.Height),
Version: item.Version,
}

// 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 && node.Value == nil {
node.Value = []byte{}
}

return node
}

// finalizeRestore commits any remaining imports and updates store metadata
func (rs *Store) finalizeRestore(height uint64, importer *iavltree.Importer, snapshotItem snapshottypes.SnapshotItem) (snapshottypes.SnapshotItem, error) {
if importer != nil {
err := importer.Commit()
if err != nil {
if err := importer.Commit(); err != nil {
return snapshottypes.SnapshotItem{}, errorsmod.Wrap(err, "IAVL commit failed")
}
importer.Close()
}

rs.flushMetadata(rs.db, int64(height), rs.buildCommitInfo(int64(height)))
Expand Down
117 changes: 67 additions & 50 deletions store/v2/commitment/store.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -509,76 +509,93 @@ func (c *CommitStore) Restore(
format uint32,
protoReader protoio.Reader,
) (snapshotstypes.SnapshotItem, error) {
var (
importer Importer
snapshotItem snapshotstypes.SnapshotItem
)
var importer Importer
defer func() {
if importer != nil {
importer.Close()
}
}()

loop:
for {
snapshotItem = snapshotstypes.SnapshotItem{}
err := protoReader.ReadMsg(&snapshotItem)
if errors.Is(err, io.EOF) {
break
} else if err != nil {
snapshotItem := snapshotstypes.SnapshotItem{}
if err := protoReader.ReadMsg(&snapshotItem); err != nil {
if errors.Is(err, io.EOF) {
break
}
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)
}
if err := importer.Close(); err != nil {
return snapshotstypes.SnapshotItem{}, fmt.Errorf("failed to close importer: %w", err)
}
}
tree := c.multiTrees[item.Store.Name]
if tree == nil {
return snapshotstypes.SnapshotItem{}, fmt.Errorf("store %s not found", item.Store.Name)
if err := handleStoreItem(c, version, &importer, item.Store); err != nil {
return snapshotstypes.SnapshotItem{}, err
}
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{}, errors.New("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{}
}
}
err := importer.Add(node)
if err != nil {
return snapshotstypes.SnapshotItem{}, fmt.Errorf("failed to add node to importer: %w", err)
if err := handleIAVLItem(importer, item.IAVL); err != nil {
return snapshotstypes.SnapshotItem{}, err
}

default:
break loop
// Return the unhandled item for the caller to process
return snapshotItem, c.LoadVersion(version)
}
}

if importer != nil {
if err := importer.Commit(); err != nil {
return snapshotstypes.SnapshotItem{}, fmt.Errorf("failed to commit importer: %w", err)
return snapshotstypes.SnapshotItem{}, fmt.Errorf("failed to commit final importer: %w", err)
}
}

return snapshotItem, c.LoadVersion(version)
return snapshotstypes.SnapshotItem{}, c.LoadVersion(version)
}

func handleStoreItem(c *CommitStore, version uint64, importer *Importer, store *snapshotstypes.SnapshotStoreItem) error {
// Commit and close previous importer if it exists
if *importer != nil {
if err := (*importer).Commit(); err != nil {
return fmt.Errorf("failed to commit importer: %w", err)
}
if err := (*importer).Close(); err != nil {
return fmt.Errorf("failed to close importer: %w", err)
}
}

// Create new importer for the store
tree := c.multiTrees[store.Name]
if tree == nil {
return fmt.Errorf("store %s not found", store.Name)
}

var err error
*importer, err = tree.Import(version)
if err != nil {
return fmt.Errorf("failed to import tree for version %d: %w", version, err)
}

return nil
}

func handleIAVLItem(importer Importer, node *snapshotstypes.SnapshotIAVLItem) error {
if importer == nil {
return errors.New("received IAVL node item before store item")
}

if node.Height > int32(math.MaxInt8) {
return fmt.Errorf("node height %v cannot exceed %v", node.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 && node.Value == nil {
node.Value = []byte{}
}

return importer.Add(node)
}

func (c *CommitStore) GetCommitInfo(version uint64) (*proof.CommitInfo, error) {
Expand Down
Loading