diff --git a/cmd/juno/dbcmd.go b/cmd/juno/dbcmd.go index 4fe5cd3a8..0effa2630 100644 --- a/cmd/juno/dbcmd.go +++ b/cmd/juno/dbcmd.go @@ -195,9 +195,10 @@ func dbSize(cmd *cobra.Command, args []string) error { items [][]string ) - for _, b := range db.BucketValues() { + buckets := db.BucketValues() + for _, b := range buckets { 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)}) + bucketItem, err := pebble.CalculatePrefixSize(cmd.Context(), pebbleDB.(*pebble.DB), []byte{byte(b)}, true) if err != nil { return err } @@ -220,6 +221,19 @@ func dbSize(cmd *cobra.Command, args []string) error { } } + // check if there is any data left in the db + lastBucket := buckets[len(buckets)-1] + lastBucketItem, err := pebble.CalculatePrefixSize(cmd.Context(), pebbleDB.(*pebble.DB), []byte{byte(lastBucket + 1)}, false) + if err != nil { + return err + } + + if lastBucketItem.Count > 0 { + items = append(items, []string{"Unknown", lastBucketItem.Size.String(), fmt.Sprintf("%d", lastBucketItem.Count)}) + totalSize += lastBucketItem.Size + totalCount += lastBucketItem.Count + } + table := tablewriter.NewWriter(os.Stdout) table.SetHeader([]string{"Bucket", "Size", "Count"}) table.AppendBulk(items) diff --git a/db/pebble/db.go b/db/pebble/db.go index 5974edf72..935729e61 100644 --- a/db/pebble/db.go +++ b/db/pebble/db.go @@ -17,6 +17,7 @@ const ( // minCache is the minimum amount of memory in megabytes to allocate to pebble read and write caching. // This is also pebble's default value. minCacheSizeMB = 8 + byteLimit = 0xff ) var ( @@ -128,7 +129,7 @@ func (i *Item) add(size utils.DataSize) { i.Size += size } -func CalculatePrefixSize(ctx context.Context, pDB *DB, prefix []byte) (*Item, error) { +func CalculatePrefixSize(ctx context.Context, pDB *DB, prefix []byte, withUpperBound bool) (*Item, error) { var ( err error v []byte @@ -136,9 +137,13 @@ func CalculatePrefixSize(ctx context.Context, pDB *DB, prefix []byte) (*Item, er item = &Item{} ) - const upperBoundofPrefix = 0xff pebbleDB := pDB.Impl().(*pebble.DB) - it, err := pebbleDB.NewIter(&pebble.IterOptions{LowerBound: prefix, UpperBound: append(prefix, upperBoundofPrefix)}) + iterOpt := &pebble.IterOptions{LowerBound: prefix} + if withUpperBound { + iterOpt.UpperBound = upperBound(prefix) + } + + it, err := pebbleDB.NewIter(iterOpt) if err != nil { // No need to call utils.RunAndWrapOnError() since iterator couldn't be created return nil, err @@ -158,3 +163,19 @@ func CalculatePrefixSize(ctx context.Context, pDB *DB, prefix []byte) (*Item, er return item, utils.RunAndWrapOnError(it.Close, err) } + +func upperBound(prefix []byte) []byte { + var ub []byte + + for i := len(prefix) - 1; i >= 0; i-- { + if prefix[i] == byteLimit { + continue + } + ub = make([]byte, i+1) + copy(ub, prefix) + ub[i]++ + return ub + } + + return nil +} diff --git a/db/pebble/db_test.go b/db/pebble/db_test.go index 896d5fac5..d1c5d72f7 100644 --- a/db/pebble/db_test.go +++ b/db/pebble/db_test.go @@ -424,7 +424,7 @@ func TestCalculatePrefixSize(t *testing.T) { t.Run("empty db", func(t *testing.T) { testDB := pebble.NewMemTest(t).(*pebble.DB) - s, err := pebble.CalculatePrefixSize(context.Background(), testDB, []byte("0")) + s, err := pebble.CalculatePrefixSize(context.Background(), testDB, []byte("0"), true) require.NoError(t, err) assert.Zero(t, s.Count) assert.Zero(t, s.Size) @@ -435,7 +435,7 @@ func TestCalculatePrefixSize(t *testing.T) { require.NoError(t, testDB.Update(func(txn db.Transaction) error { return txn.Set(append([]byte("0"), []byte("randomKey")...), []byte("someValue")) })) - s, err := pebble.CalculatePrefixSize(context.Background(), testDB.(*pebble.DB), []byte("1")) + s, err := pebble.CalculatePrefixSize(context.Background(), testDB.(*pebble.DB), []byte("1"), true) require.NoError(t, err) assert.Zero(t, s.Count) assert.Zero(t, s.Size) @@ -455,7 +455,7 @@ func TestCalculatePrefixSize(t *testing.T) { return txn.Set(k3, v3) })) - s, err := pebble.CalculatePrefixSize(context.Background(), testDB.(*pebble.DB), p) + s, err := pebble.CalculatePrefixSize(context.Background(), testDB.(*pebble.DB), p, true) require.NoError(t, err) assert.Equal(t, uint(3), s.Count) assert.Equal(t, utils.DataSize(expectedSize), s.Size) @@ -464,7 +464,7 @@ func TestCalculatePrefixSize(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) cancel() - s, err := pebble.CalculatePrefixSize(ctx, testDB.(*pebble.DB), p) + s, err := pebble.CalculatePrefixSize(ctx, testDB.(*pebble.DB), p, true) assert.EqualError(t, err, context.Canceled.Error()) assert.Zero(t, s.Count) assert.Zero(t, s.Size)