pebble: Update Tokenbucket package and use WaitCtx

This change updates the cockroach tokenbucket version in pebble as well as uses TokenBucket.WaitCtx inside of DB.ScanStatistics. Fixes: #2784
cockroachdb · Aug 11, 2023 · 77e81e8 · 77e81e8
1 parent 18e6ad4
commit 77e81e8
Show file tree

Hide file tree

Showing 5 changed files with 10 additions and 15 deletions.
diff --git a/db.go b/db.go
@@ -2687,21 +2687,14 @@ func (d *DB) ScanStatistics(
 ) (LSMKeyStatistics, error) {
 	stats := LSMKeyStatistics{}
 	var prevKey InternalKey
-	var rateLimitFunc func(key *InternalKey, val LazyValue)
+	var rateLimitFunc func(key *InternalKey, val LazyValue) error
 	tb := tokenbucket.TokenBucket{}
 
 	if opts.LimitBytesPerSecond != 0 {
 		// Each "token" roughly corresponds to a byte that was read.
 		tb.Init(tokenbucket.TokensPerSecond(opts.LimitBytesPerSecond), tokenbucket.Tokens(1024))
-		rateLimitFunc = func(key *InternalKey, val LazyValue) {
-			for {
-				fulfilled, tryAgainAfter := tb.TryToFulfill(tokenbucket.Tokens(key.Size() + val.Len()))
-
-				if fulfilled {
-					break
-				}
-				time.Sleep(tryAgainAfter)
-			}
+		rateLimitFunc = func(key *InternalKey, val LazyValue) error {
+			return tb.WaitCtx(ctx, tokenbucket.Tokens(key.Size()+val.Len()))
 		}
 	}
 

diff --git a/go.mod b/go.mod
@@ -7,7 +7,7 @@ require (
 	github.com/cockroachdb/datadriven v1.0.3-0.20230413201302-be42291fc80f
 	github.com/cockroachdb/errors v1.8.1
 	github.com/cockroachdb/redact v1.0.8
-	github.com/cockroachdb/tokenbucket v0.0.0-20230613231145-182959a1fad6
+	github.com/cockroachdb/tokenbucket v0.0.0-20230807174530-cc333fc44b06
 	github.com/ghemawat/stream v0.0.0-20171120220530-696b145b53b9
 	github.com/golang/snappy v0.0.4
 	github.com/guptarohit/asciigraph v0.5.5

diff --git a/go.sum b/go.sum
@@ -89,8 +89,8 @@ github.com/cockroachdb/redact v1.0.8 h1:8QG/764wK+vmEYoOlfobpe12EQcS81ukx/a4hdVM
 github.com/cockroachdb/redact v1.0.8/go.mod h1:BVNblN9mBWFyMyqK1k3AAiSxhvhfK2oOZZ2lK+dpvRg=
 github.com/cockroachdb/sentry-go v0.6.1-cockroachdb.2 h1:IKgmqgMQlVJIZj19CdocBeSfSaiCbEBZGKODaixqtHM=
 github.com/cockroachdb/sentry-go v0.6.1-cockroachdb.2/go.mod h1:8BT+cPK6xvFOcRlk0R8eg+OTkcqI6baNH4xAkpiYVvQ=
-github.com/cockroachdb/tokenbucket v0.0.0-20230613231145-182959a1fad6 h1:DJK8W/iB+s/qkTtmXSrHA49lp5O3OsR7E6z4byOLy34=
-github.com/cockroachdb/tokenbucket v0.0.0-20230613231145-182959a1fad6/go.mod h1:7nc4anLGjupUW/PeY5qiNYsdNXj7zopG+eqsS7To5IQ=
+github.com/cockroachdb/tokenbucket v0.0.0-20230807174530-cc333fc44b06 h1:zuQyyAKVxetITBuuhv3BI9cMrmStnpT18zmgmTxunpo=
+github.com/cockroachdb/tokenbucket v0.0.0-20230807174530-cc333fc44b06/go.mod h1:7nc4anLGjupUW/PeY5qiNYsdNXj7zopG+eqsS7To5IQ=
 github.com/codegangsta/inject v0.0.0-20150114235600-33e0aa1cb7c0/go.mod h1:4Zcjuz89kmFXt9morQgcfYZAYZ5n8WHjt81YYWIwtTM=
 github.com/coreos/bbolt v1.3.2/go.mod h1:iRUV2dpdMOn7Bo10OQBFzIJO9kkE559Wcmn+qkEiiKk=
 github.com/coreos/etcd v3.3.10+incompatible/go.mod h1:uF7uidLiAD3TWHmW31ZFd/JWoc32PjwdhPthX9715RE=

diff --git a/options.go b/options.go
@@ -260,7 +260,7 @@ type scanInternalOptions struct {
 	includeObsoleteKeys bool
 
 	// rateLimitFunc is used to limit the amount of bytes read per second.
-	rateLimitFunc func(key *InternalKey, value LazyValue)
+	rateLimitFunc func(key *InternalKey, value LazyValue) error
 }
 
 // RangeKeyMasking configures automatic hiding of point keys by range keys. A

diff --git a/scan_internal.go b/scan_internal.go
@@ -647,7 +647,9 @@ func scanInternalImpl(
 		key := iter.unsafeKey()
 
 		if opts.rateLimitFunc != nil {
-			opts.rateLimitFunc(key, iter.lazyValue())
+			if err := opts.rateLimitFunc(key, iter.lazyValue()); err != nil {
+				return err
+			}
 		}
 
 		switch key.Kind() {