From f6d9cef3dce5aa958db9b736c394dc952c2a1501 Mon Sep 17 00:00:00 2001
From: weiihann <weihan774237@gmail.com>
Date: Thu, 17 Oct 2024 16:50:44 +0800
Subject: [PATCH] check for unknown KV, fix upper bound bug

---
 cmd/juno/dbcmd.go    | 18 ++++++++++++++++--
 db/pebble/db.go      | 27 ++++++++++++++++++++++++---
 db/pebble/db_test.go |  8 ++++----
 3 files changed, 44 insertions(+), 9 deletions(-)

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)