From 3cf58a15fc7189d173f3f4e0af10c98d5bb60149 Mon Sep 17 00:00:00 2001 From: Ng Wei Han <47109095+weiihann@users.noreply.github.com> Date: Fri, 9 Aug 2024 19:20:45 +0800 Subject: [PATCH] Add database CLI tools (#1902) --- cmd/juno/dbcmd.go | 203 +++++++++++++++++++++++++++++++++++++++++ cmd/juno/dbcmd_test.go | 53 +++++++++++ cmd/juno/dbsize.go | 88 ------------------ cmd/juno/juno.go | 3 +- db/pebble/db.go | 30 ++++-- db/pebble/db_test.go | 13 ++- go.mod | 1 + go.sum | 1 + utils/const.go | 5 +- utils/size.go | 21 +++++ 10 files changed, 314 insertions(+), 104 deletions(-) create mode 100644 cmd/juno/dbcmd.go create mode 100644 cmd/juno/dbcmd_test.go delete mode 100644 cmd/juno/dbsize.go create mode 100644 utils/size.go diff --git a/cmd/juno/dbcmd.go b/cmd/juno/dbcmd.go new file mode 100644 index 0000000000..be3d51b64e --- /dev/null +++ b/cmd/juno/dbcmd.go @@ -0,0 +1,203 @@ +package main + +import ( + "encoding/json" + "fmt" + "os" + + "github.com/NethermindEth/juno/blockchain" + "github.com/NethermindEth/juno/core" + "github.com/NethermindEth/juno/core/felt" + "github.com/NethermindEth/juno/db" + "github.com/NethermindEth/juno/db/pebble" + "github.com/NethermindEth/juno/utils" + "github.com/olekukonko/tablewriter" + "github.com/spf13/cobra" +) + +type DBInfo struct { + Network string `json:"network"` + ChainHeight uint64 `json:"chain_height"` + LatestBlockHash *felt.Felt `json:"latest_block_hash"` + LatestStateRoot *felt.Felt `json:"latest_state_root"` + L1Height uint64 `json:"l1_height"` + L1BlockHash *felt.Felt `json:"l1_block_hash"` + L1StateRoot *felt.Felt `json:"l1_state_root"` +} + +func DBCmd(defaultDBPath string) *cobra.Command { + dbCmd := &cobra.Command{ + Use: "db", + Short: "Database related operations", + Long: `This command allows you to perform database operations.`, + } + + dbCmd.PersistentFlags().String(dbPathF, defaultDBPath, dbPathUsage) + dbCmd.AddCommand(DBInfoCmd(), DBSizeCmd()) + return dbCmd +} + +func DBInfoCmd() *cobra.Command { + return &cobra.Command{ + Use: "info", + Short: "Retrieve database information", + Long: `This subcommand retrieves and displays blockchain information stored in the database.`, + RunE: dbInfo, + } +} + +func DBSizeCmd() *cobra.Command { + return &cobra.Command{ + Use: "size", + Short: "Calculate database size information for each data type", + Long: `This subcommand retrieves and displays the storage of each data type stored in the database.`, + RunE: dbSize, + } +} + +func dbInfo(cmd *cobra.Command, args []string) error { + dbPath, err := cmd.Flags().GetString(dbPathF) + if err != nil { + return err + } + + if _, err = os.Stat(dbPath); os.IsNotExist(err) { + fmt.Fprintln(cmd.OutOrStdout(), "Database path does not exist") + return nil + } + + database, err := pebble.New(dbPath) + if err != nil { + return fmt.Errorf("open DB: %w", err) + } + + chain := blockchain.New(database, nil) + info := DBInfo{} + + // Get the latest block information + headBlock, err := chain.Head() + if err != nil { + return fmt.Errorf("failed to get the latest block information: %v", err) + } + + stateUpdate, err := chain.StateUpdateByNumber(headBlock.Number) + if err != nil { + return fmt.Errorf("failed to get the state update: %v", err) + } + + info.Network = getNetwork(headBlock, stateUpdate.StateDiff) + info.ChainHeight = headBlock.Number + info.LatestBlockHash = headBlock.Hash + info.LatestStateRoot = headBlock.GlobalStateRoot + + // Get the latest L1 block information + l1Head, err := chain.L1Head() + if err == nil { + info.L1Height = l1Head.BlockNumber + info.L1BlockHash = l1Head.BlockHash + info.L1StateRoot = l1Head.StateRoot + } else { + fmt.Printf("Failed to get the latest L1 block information: %v\n", err) + } + + jsonData, err := json.MarshalIndent(info, "", "") + if err != nil { + return fmt.Errorf("marshal JSON: %w", err) + } + + fmt.Fprintln(cmd.OutOrStdout(), string(jsonData)) + + return nil +} + +func dbSize(cmd *cobra.Command, args []string) error { + dbPath, err := cmd.Flags().GetString(dbPathF) + if err != nil { + return err + } + + if dbPath == "" { + return fmt.Errorf("--%v cannot be empty", dbPathF) + } + + if _, err = os.Stat(dbPath); os.IsNotExist(err) { + fmt.Fprintln(cmd.OutOrStdout(), "Database path does not exist") + return nil + } + + pebbleDB, err := pebble.New(dbPath) + if err != nil { + return err + } + + var ( + totalSize utils.DataSize + totalCount uint + + withHistorySize utils.DataSize + withoutHistorySize utils.DataSize + + withHistoryCount uint + withoutHistoryCount uint + + items [][]string + ) + + for _, b := range db.BucketValues() { + fmt.Fprintf(cmd.OutOrStdout(), "Calculating size of %s, remaining buckets: %d\n", b, len(db.BucketValues())-int(b)-1) + bucketItem, err := pebble.CalculatePrefixSize(cmd.Context(), pebbleDB.(*pebble.DB), []byte{byte(b)}) + if err != nil { + return err + } + items = append(items, []string{b.String(), bucketItem.Size.String(), fmt.Sprintf("%d", bucketItem.Count)}) + + totalSize += bucketItem.Size + totalCount += bucketItem.Count + + if utils.AnyOf(b, db.StateTrie, db.ContractStorage, db.Class, db.ContractNonce, db.ContractDeploymentHeight) { + withoutHistorySize += bucketItem.Size + withHistorySize += bucketItem.Size + + withoutHistoryCount += bucketItem.Count + withHistoryCount += bucketItem.Count + } + + if utils.AnyOf(b, db.ContractStorageHistory, db.ContractNonceHistory, db.ContractClassHashHistory) { + withHistorySize += bucketItem.Size + withHistoryCount += bucketItem.Count + } + } + + table := tablewriter.NewWriter(os.Stdout) + table.SetHeader([]string{"Bucket", "Size", "Count"}) + table.AppendBulk(items) + table.SetFooter([]string{"Total", totalSize.String(), fmt.Sprintf("%d", totalCount)}) + table.Render() + + tableState := tablewriter.NewWriter(os.Stdout) + tableState.SetHeader([]string{"State", "Size", "Count"}) + tableState.Append([]string{"Without history", withoutHistorySize.String(), fmt.Sprintf("%d", withoutHistoryCount)}) + tableState.Append([]string{"With history", withHistorySize.String(), fmt.Sprintf("%d", withHistoryCount)}) + tableState.Render() + + return nil +} + +func getNetwork(head *core.Block, stateDiff *core.StateDiff) string { + networks := []*utils.Network{ + &utils.Mainnet, + &utils.Sepolia, + &utils.Goerli, + &utils.Goerli2, + &utils.Integration, + &utils.SepoliaIntegration, + } + + for _, network := range networks { + if _, err := core.VerifyBlockHash(head, network, stateDiff); err == nil { + return network.Name + } + } + + return "unknown" +} diff --git a/cmd/juno/dbcmd_test.go b/cmd/juno/dbcmd_test.go new file mode 100644 index 0000000000..3491923962 --- /dev/null +++ b/cmd/juno/dbcmd_test.go @@ -0,0 +1,53 @@ +package main_test + +import ( + "context" + "testing" + + "github.com/NethermindEth/juno/blockchain" + "github.com/NethermindEth/juno/clients/feeder" + juno "github.com/NethermindEth/juno/cmd/juno" + "github.com/NethermindEth/juno/core" + "github.com/NethermindEth/juno/db/pebble" + adaptfeeder "github.com/NethermindEth/juno/starknetdata/feeder" + "github.com/NethermindEth/juno/utils" + "github.com/spf13/cobra" + "github.com/stretchr/testify/require" +) + +var emptyCommitments = core.BlockCommitments{} + +func TestDBCmd(t *testing.T) { + t.Run("retrieve info when db contains block0", func(t *testing.T) { + cmd := juno.DBInfoCmd() + executeCmdInDB(t, cmd) + }) + + t.Run("inspect db when db contains block0", func(t *testing.T) { + cmd := juno.DBSizeCmd() + executeCmdInDB(t, cmd) + }) +} + +func executeCmdInDB(t *testing.T, cmd *cobra.Command) { + cmd.Flags().String("db-path", "", "") + + client := feeder.NewTestClient(t, &utils.Mainnet) + gw := adaptfeeder.New(client) + block0, err := gw.BlockByNumber(context.Background(), 0) + require.NoError(t, err) + + stateUpdate0, err := gw.StateUpdate(context.Background(), 0) + require.NoError(t, err) + + dbPath := t.TempDir() + testDB, err := pebble.New(dbPath) + require.NoError(t, err) + + chain := blockchain.New(testDB, &utils.Mainnet) + require.NoError(t, chain.Store(block0, &emptyCommitments, stateUpdate0, nil)) + testDB.Close() + + require.NoError(t, cmd.Flags().Set("db-path", dbPath)) + require.NoError(t, cmd.Execute()) +} diff --git a/cmd/juno/dbsize.go b/cmd/juno/dbsize.go deleted file mode 100644 index 11f1dbb6f4..0000000000 --- a/cmd/juno/dbsize.go +++ /dev/null @@ -1,88 +0,0 @@ -package main - -import ( - "fmt" - - "github.com/NethermindEth/juno/db" - "github.com/NethermindEth/juno/db/pebble" - "github.com/NethermindEth/juno/utils" - "github.com/spf13/cobra" -) - -func DBSize() *cobra.Command { - dbSizeCmd := &cobra.Command{ - Use: "db-size", - Short: "Calculate's Juno's DB size.", - RunE: func(cmd *cobra.Command, args []string) error { - dbPath, err := cmd.Flags().GetString(dbPathF) - if err != nil { - return err - } - - if dbPath == "" { - return fmt.Errorf("--%v cannot be empty", dbPathF) - } - - pebbleDB, err := pebble.New(dbPath) - if err != nil { - return err - } - - var totalSize, stateSizeWithoutHistory, stateSizeWithHistory uint - - _, err = fmt.Fprintln(cmd.OutOrStdout(), "Total number of DB buckets:", len(db.BucketValues())) - if err != nil { - return err - } - - var bucketSize uint - for _, b := range db.BucketValues() { - bucketSize, err = pebble.CalculatePrefixSize(cmd.Context(), pebbleDB.(*pebble.DB), []byte{byte(b)}) - if err != nil { - return err - } - - _, err = fmt.Fprintln(cmd.OutOrStdout(), uint(b)+1, "Size of", b, "=", bucketSize) - if err != nil { - return err - } - - totalSize += bucketSize - - if utils.AnyOf(b, db.StateTrie, db.ContractStorage, db.Class, db.ContractNonce, - db.ContractDeploymentHeight) { - stateSizeWithoutHistory += bucketSize - stateSizeWithHistory += bucketSize - } - - if utils.AnyOf(b, db.ContractStorageHistory, db.ContractNonceHistory, db.ContractClassHashHistory) { - stateSizeWithHistory += bucketSize - } - } - - _, err = fmt.Fprintln(cmd.OutOrStdout()) - if err != nil { - return err - } - _, err = fmt.Fprintln(cmd.OutOrStdout(), "State size without history =", stateSizeWithoutHistory) - if err != nil { - return err - } - - _, err = fmt.Fprintln(cmd.OutOrStdout(), "State size with history =", stateSizeWithHistory) - if err != nil { - return err - } - - _, err = fmt.Fprintln(cmd.OutOrStdout(), "Total DB size =", totalSize) - return err - }, - } - - // Persistent Flag was not used from the Juno command because GenP2PKeyPair would also inherit it while PersistentPreRun was not used - // because none of the subcommand required access to the node.Config. - defaultDBPath, dbPathShort := "", "p" - dbSizeCmd.Flags().StringP(dbPathF, dbPathShort, defaultDBPath, dbPathUsage) - - return dbSizeCmd -} diff --git a/cmd/juno/juno.go b/cmd/juno/juno.go index 1669143ab8..c61ee40529 100644 --- a/cmd/juno/juno.go +++ b/cmd/juno/juno.go @@ -356,8 +356,7 @@ func NewCmd(config *node.Config, run func(*cobra.Command, []string) error) *cobr junoCmd.Flags().String(versionedConstantsFileF, defaultVersionedConstantsFile, versionedConstantsFileUsage) junoCmd.MarkFlagsMutuallyExclusive(p2pFeederNodeF, p2pPeersF) - junoCmd.AddCommand(GenP2PKeyPair()) - junoCmd.AddCommand(DBSize()) + junoCmd.AddCommand(GenP2PKeyPair(), DBCmd(defaultDBPath)) return junoCmd } diff --git a/db/pebble/db.go b/db/pebble/db.go index 6b4869fb9f..6592a602f0 100644 --- a/db/pebble/db.go +++ b/db/pebble/db.go @@ -118,11 +118,22 @@ func (d *DB) Impl() any { return d.pebble } -func CalculatePrefixSize(ctx context.Context, pDB *DB, prefix []byte) (uint, error) { +type Item struct { + Count uint + Size utils.DataSize +} + +func (i *Item) add(size utils.DataSize) { + i.Count++ + i.Size += size +} + +func CalculatePrefixSize(ctx context.Context, pDB *DB, prefix []byte) (*Item, error) { var ( - err error - size uint - v []byte + err error + v []byte + + item = &Item{} ) const upperBoundofPrefix = 0xff @@ -130,19 +141,20 @@ func CalculatePrefixSize(ctx context.Context, pDB *DB, prefix []byte) (uint, err it, err := pebbleDB.NewIter(&pebble.IterOptions{LowerBound: prefix, UpperBound: append(prefix, upperBoundofPrefix)}) if err != nil { // No need to call utils.RunAndWrapOnError() since iterator couldn't be created - return 0, err + return nil, err } for it.First(); it.Valid(); it.Next() { if ctx.Err() != nil { - return size, utils.RunAndWrapOnError(it.Close, ctx.Err()) + return item, utils.RunAndWrapOnError(it.Close, ctx.Err()) } v, err = it.ValueAndErr() if err != nil { - return 0, utils.RunAndWrapOnError(it.Close, err) + return nil, utils.RunAndWrapOnError(it.Close, err) } - size += uint(len(it.Key()) + len(v)) + + item.add(utils.DataSize(len(it.Key()) + len(v))) } - return size, utils.RunAndWrapOnError(it.Close, err) + return item, utils.RunAndWrapOnError(it.Close, err) } diff --git a/db/pebble/db_test.go b/db/pebble/db_test.go index a51e9e4c24..c239764cee 100644 --- a/db/pebble/db_test.go +++ b/db/pebble/db_test.go @@ -10,6 +10,7 @@ import ( "github.com/NethermindEth/juno/db" "github.com/NethermindEth/juno/db/pebble" + "github.com/NethermindEth/juno/utils" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -425,7 +426,8 @@ func TestCalculatePrefixSize(t *testing.T) { s, err := pebble.CalculatePrefixSize(context.Background(), testDB, []byte("0")) require.NoError(t, err) - assert.Equal(t, uint(0), s) + assert.Zero(t, s.Count) + assert.Zero(t, s.Size) }) t.Run("non empty db but empty prefix", func(t *testing.T) { @@ -435,7 +437,8 @@ func TestCalculatePrefixSize(t *testing.T) { })) s, err := pebble.CalculatePrefixSize(context.Background(), testDB.(*pebble.DB), []byte("1")) require.NoError(t, err) - assert.Equal(t, uint(0), s) + assert.Zero(t, s.Count) + assert.Zero(t, s.Size) }) t.Run("size of all key value pair with the same prefix", func(t *testing.T) { @@ -454,7 +457,8 @@ func TestCalculatePrefixSize(t *testing.T) { s, err := pebble.CalculatePrefixSize(context.Background(), testDB.(*pebble.DB), p) require.NoError(t, err) - assert.Equal(t, expectedSize, s) + assert.Equal(t, uint(3), s.Count) + assert.Equal(t, utils.DataSize(expectedSize), s.Size) t.Run("exit when context is cancelled", func(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) @@ -462,7 +466,8 @@ func TestCalculatePrefixSize(t *testing.T) { s, err := pebble.CalculatePrefixSize(ctx, testDB.(*pebble.DB), p) assert.EqualError(t, err, context.Canceled.Error()) - assert.Equal(t, uint(0), s) + assert.Zero(t, s.Count) + assert.Zero(t, s.Size) }) }) } diff --git a/go.mod b/go.mod index d06543c9c1..3c4064ae33 100644 --- a/go.mod +++ b/go.mod @@ -19,6 +19,7 @@ require ( github.com/libp2p/go-libp2p-pubsub v0.11.0 github.com/mitchellh/mapstructure v1.5.0 github.com/multiformats/go-multiaddr v0.12.4 + github.com/olekukonko/tablewriter v0.0.5 github.com/pkg/errors v0.9.1 github.com/prometheus/client_golang v1.19.1 github.com/rs/cors v1.11.0 diff --git a/go.sum b/go.sum index 191219c24c..9de36b235a 100644 --- a/go.sum +++ b/go.sum @@ -339,6 +339,7 @@ github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovk github.com/mattn/go-isatty v0.0.14/go.mod h1:7GGIvUiUoEMVVmxf/4nioHXj79iQHKdU27kJ6hsGG94= github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= +github.com/mattn/go-runewidth v0.0.9/go.mod h1:H031xJmbD/WCDINGzjvQ9THkh0rPKHF+m2gUSrubnMI= github.com/mattn/go-runewidth v0.0.14 h1:+xnbZSEeDbOIg5/mE6JF0w6n9duR1l3/WmbinWVwUuU= github.com/mattn/go-runewidth v0.0.14/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w= github.com/matttproud/golang_protobuf_extensions v1.0.1/go.mod h1:D8He9yQNgCq6Z5Ld7szi9bcBfOoFv/3dc6xSMkL2PC0= diff --git a/utils/const.go b/utils/const.go index 125ebb6aeb..f3ada59381 100644 --- a/utils/const.go +++ b/utils/const.go @@ -1,5 +1,8 @@ package utils const ( - Megabyte = 1024 * 1024 + Kilobyte = 1024 + Megabyte = 1024 * Kilobyte + Gigabyte = 1024 * Megabyte + Terabyte = 1024 * Gigabyte ) diff --git a/utils/size.go b/utils/size.go new file mode 100644 index 0000000000..610abf85bd --- /dev/null +++ b/utils/size.go @@ -0,0 +1,21 @@ +package utils + +import ( + "fmt" +) + +type DataSize float64 + +func (d DataSize) String() string { + switch { + case d >= Terabyte: + return fmt.Sprintf("%.2f TiB", d/Terabyte) + case d >= Gigabyte: + return fmt.Sprintf("%.2f GiB", d/Gigabyte) + case d >= Megabyte: + return fmt.Sprintf("%.2f MiB", d/Megabyte) + case d >= Kilobyte: + return fmt.Sprintf("%.2f KiB", d/Kilobyte) + } + return fmt.Sprintf("%.2f B", d) +}