Fix upper bound issue in CalculatePrefixSize (#2223)

* check for unknown KV, fix upper bound bug

* minor chore

* remove upperbound option

* Revert "remove upperbound option"

This reverts commit e4b412b.

* Update db/pebble/db.go

Co-authored-by: Pawel Nowosielski <[email protected]>
Signed-off-by: Ng Wei Han <[email protected]>

* fix

---------

Signed-off-by: Ng Wei Han <[email protected]>
Co-authored-by: Pawel Nowosielski <[email protected]>

Author: weiihann

cmd/juno/dbcmd.go

@@ -204,9 +204,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
 		}
@@ -229,6 +230,20 @@ func dbSize(cmd *cobra.Command, args []string) error {
 		}
 	}
 
+	// check if there is any data left in the db
+	lastBucket := buckets[len(buckets)-1]
+	fmt.Fprintln(cmd.OutOrStdout(), "Calculating remaining data in the db")
+	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)

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
+	maxByte        = ^byte(0)
 )
 
 var (
@@ -128,17 +129,21 @@ 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
 
 		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] == maxByte {
+			continue
+		}
+		ub = make([]byte, i+1)
+		copy(ub, prefix)
+		ub[i]++
+		return ub
+	}
+
+	return nil
+}

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)