Skip to content

Commit

Permalink
add cmd checkpoint trie stats
Browse files Browse the repository at this point in the history
  • Loading branch information
zhangchiqing committed Aug 15, 2023
1 parent 537fdac commit 6aac0fe
Show file tree
Hide file tree
Showing 2 changed files with 133 additions and 0 deletions.
105 changes: 105 additions & 0 deletions cmd/util/cmd/checkpoint-trie-stats/cmd.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,105 @@
package checkpoint_trie_stats

import (
"errors"
"fmt"

"github.com/onflow/flow-go/ledger/complete/mtrie/node"
"github.com/onflow/flow-go/ledger/complete/mtrie/trie"
"github.com/onflow/flow-go/ledger/complete/wal"
"github.com/rs/zerolog"
"github.com/rs/zerolog/log"
"github.com/spf13/cobra"
)

var (
flagCheckpoint string
flagTrieIndex int
)

var Cmd = &cobra.Command{
Use: "checkpoint-trie-stats",
Short: "List the trie node count by types in a checkpoint, show total payload size",
Run: run,
}

func init() {

Cmd.Flags().StringVar(&flagCheckpoint, "checkpoint", "",
"checkpoint file to read")
_ = Cmd.MarkFlagRequired("checkpoint")
Cmd.Flags().IntVar(&flagTrieIndex, "trie-index", 0, "trie index to read, 0 being the first trie, -1 is the last trie")

}

func run(*cobra.Command, []string) {

log.Info().Msgf("loading checkpoint %v, reading %v-th trie", flagCheckpoint, flagTrieIndex)
res, err := scanCheckpoint(flagCheckpoint, flagTrieIndex, log.Logger)
if err != nil {
log.Fatal().Err(err).Msg("fail to scan checkpoint")
}
log.Info().
Str("TrieRootHash", res.trieRootHash).
Int("InterimNodeCount", res.interimNodeCount).
Int("LeafNodeCount", res.leafNodeCount).
Int("TotalPayloadSize", res.totalPayloadSize).
Msgf("successfully scanned checkpoint %v", flagCheckpoint)
}

type result struct {
trieRootHash string
interimNodeCount int
leafNodeCount int
totalPayloadSize int
}

func readTrie(tries []*trie.MTrie, index int) (*trie.MTrie, error) {
if len(tries) == 0 {
return nil, errors.New("No tries available")
}

if index < -len(tries) || index >= len(tries) {
return nil, fmt.Errorf("index %d out of range", index)
}

if index < 0 {
return tries[len(tries)+index], nil
}

return tries[index], nil
}

func scanCheckpoint(checkpoint string, trieIndex int, log zerolog.Logger) (result, error) {
tries, err := wal.LoadCheckpoint(flagCheckpoint, &log)
if err != nil {
return result{}, fmt.Errorf("error while loading checkpoint: %w", err)
}

log.Info().Msgf("checkpoint loaded, total tries: %v", len(tries))

t, err := readTrie(tries, trieIndex)
if err != nil {
return result{}, fmt.Errorf("error while reading trie: %w", err)
}

log.Info().Msgf("trie loaded, root hash: %v", t.RootHash())

var res *result
processNode := func(n *node.Node) error {
if n.IsLeaf() {
res.leafNodeCount++
res.totalPayloadSize += n.Payload().Size()
} else {
res.interimNodeCount++
}
return nil
}

err = trie.TraverseNodes(t, processNode)
if err != nil {
return result{}, fmt.Errorf("fail to traverse the trie: %w", err)
}

return *res, nil
}
28 changes: 28 additions & 0 deletions ledger/complete/mtrie/trie/trie.go
Original file line number Diff line number Diff line change
Expand Up @@ -831,3 +831,31 @@ func minInt(a, b int) int {
}
return b
}

// TraverseNodes traverses all nodes of the trie in DFS order
func TraverseNodes(trie *MTrie, processNode func(*node.Node) error) error {
return traverseRecursive(trie.root, processNode)
}

func traverseRecursive(n *node.Node, processNode func(*node.Node) error) error {
if n == nil {
return nil
}

err := processNode(n)
if err != nil {
return err
}

err = traverseRecursive(n.LeftChild(), processNode)
if err != nil {
return err
}

err = traverseRecursive(n.RightChild(), processNode)
if err != nil {
return err
}

return nil
}

0 comments on commit 6aac0fe

Please sign in to comment.