From b52a1521748885a1642f0dd2dd9c0f193b68c30f Mon Sep 17 00:00:00 2001
From: Rahul Aggarwal <rahul.aggarwal@cockroachlabs.com>
Date: Mon, 31 Jul 2023 11:23:27 -0400
Subject: [PATCH] Rate Limit Scan Statistics

This pull request uses a `TokenBucket` to limit the number of keys that
read from `ScanStatistics` within a certain period of time.

Fixes: #2778
---
 db.go | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/db.go b/db.go
index 8c4e35613f..cdbbd32679 100644
--- a/db.go
+++ b/db.go
@@ -29,6 +29,7 @@ import (
 	"github.com/cockroachdb/pebble/sstable"
 	"github.com/cockroachdb/pebble/vfs"
 	"github.com/cockroachdb/pebble/vfs/atomicfs"
+	"github.com/cockroachdb/tokenbucket"
 	"github.com/prometheus/client_golang/prometheus"
 )
 
@@ -2604,9 +2605,21 @@ type LSMKeyStatistics struct {
 func (d *DB) ScanStatistics(ctx context.Context, lower, upper []byte) (LSMKeyStatistics, error) {
 	stats := LSMKeyStatistics{}
 	var prevKey InternalKey
+	tb := tokenbucket.TokenBucket{}
+	tb.Init(50, 50)
+
+	rateLimit := func() {
+		fulfilled, tryAgainAfter := tb.TryToFulfill(1)
+
+		if !fulfilled {
+			time.Sleep(tryAgainAfter)
+		}
+	}
 
 	err := d.ScanInternal(ctx, lower, upper,
 		func(key *InternalKey, value LazyValue, iterInfo iterInfo) error {
+			rateLimit()
+
 			// iterInfo.level == -1 indicates that the key does not come from a valid level.
 			if iterInfo.level == -1 {
 				return nil
@@ -2628,10 +2641,12 @@ func (d *DB) ScanStatistics(ctx context.Context, lower, upper []byte) (LSMKeySta
 			return nil
 		},
 		func(start, end []byte, seqNum uint64) error {
+			rateLimit()
 			stats.accumulated.kindsCount[InternalKeyKindRangeDelete]++
 			return nil
 		},
 		func(start, end []byte, keys []rangekey.Key) error {
+			rateLimit()
 			for _, key := range keys {
 				stats.accumulated.kindsCount[key.Kind()]++
 			}