Merge #55904

55904: kvserver: replace rate.RateLimiter with quotapool.RateLimiter in consistency_queue r=lunevalex a=lunevalex

Fixes #55652

When running a large restore job, we discovered high CPU usage, which was attributed to the new rate limiter in the consistency_queue introduced in #49763. During the checksum computation the rate limiter ends up doing a lot of nanosecond waits, due to the typically small size of individual key/value pairs. This is inefficient by itself, but since the rate.RateLimiter creates a new Timer for every wait, it ends up being extremely inefficient. The quota pool.RateLimiter pools Timers reducing allocation churn and supports a minimum wait parameter, which converts a lot of small waits into a few larger waits during checksum calculations.

Release note: None

Co-authored-by: Alex Lunev <[email protected]>

pkg/kv/kvserver/BUILD.bazel

@@ -341,7 +341,6 @@ go_test(
         "//pkg/util/hlc",
         "//pkg/util/humanizeutil",
         "//pkg/util/leaktest",
-        "//pkg/util/limit",
         "//pkg/util/log",
         "//pkg/util/metric",
         "//pkg/util/mon",

pkg/kv/kvserver/consistency_queue.go

@@ -41,6 +41,18 @@ var consistencyCheckRate = settings.RegisterPublicValidatedByteSizeSetting(
 	validatePositive,
 )
 
+// consistencyCheckRateBurstFactor we use this to set the burst parameter on the
+// quotapool.RateLimiter. It seems overkill to provide a user setting for this,
+// so we use a factor to scale the burst setting based on the rate defined above.
+const consistencyCheckRateBurstFactor = 8
+
+// consistencyCheckRateMinWait is the minimum time to wait once the rate limit
+// is reached. We check the limit on every key/value pair, which can lead to
+// a lot of nano-second waits because each pair could be very small. Instead we
+// force a larger pause every time the timer is breached to reduce the
+// churn on timers.
+const consistencyCheckRateMinWait = 100 * time.Millisecond
+
 var testingAggressiveConsistencyChecks = envutil.EnvOrDefaultBool("COCKROACH_CONSISTENCY_AGGRESSIVE", false)
 
 type consistencyQueue struct {

pkg/kv/kvserver/replica_consistency.go

@@ -34,9 +34,9 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/bufalloc"
 	"github.com/cockroachdb/cockroach/pkg/util/envutil"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
-	"github.com/cockroachdb/cockroach/pkg/util/limit"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/protoutil"
+	"github.com/cockroachdb/cockroach/pkg/util/quotapool"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/cockroach/pkg/util/uuid"
 	"github.com/cockroachdb/errors"
@@ -570,7 +570,7 @@ func (r *Replica) sha512(
 	snap storage.Reader,
 	snapshot *roachpb.RaftSnapshotData,
 	mode roachpb.ChecksumMode,
-	limiter *limit.LimiterBurstDisabled,
+	limiter *quotapool.RateLimiter,
 ) (*replicaHash, error) {
 	statsOnly := mode == roachpb.ChecksumMode_CHECK_STATS
 
@@ -586,7 +586,7 @@ func (r *Replica) sha512(
 
 	visitor := func(unsafeKey storage.MVCCKey, unsafeValue []byte) error {
 		// Rate Limit the scan through the range
-		if err := limiter.WaitN(ctx, len(unsafeKey.Key)+len(unsafeValue)); err != nil {
+		if err := limiter.WaitN(ctx, int64(len(unsafeKey.Key)+len(unsafeValue))); err != nil {
 			return err
 		}
 

pkg/kv/kvserver/replica_proposal.go

@@ -34,7 +34,6 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/storage/fs"
 	"github.com/cockroachdb/cockroach/pkg/util"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
-	"github.com/cockroachdb/cockroach/pkg/util/limit"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/quotapool"
 	"github.com/cockroachdb/cockroach/pkg/util/sysutil"
@@ -228,8 +227,6 @@ func (r *Replica) computeChecksumPostApply(ctx context.Context, cc kvserverpb.Co
 		}
 	}
 
-	limiter := limit.NewLimiter(rate.Limit(consistencyCheckRate.Get(&r.store.ClusterSettings().SV)))
-
 	// Compute SHA asynchronously and store it in a map by UUID.
 	if err := stopper.RunAsyncTask(ctx, "storage.Replica: computing checksum", func(ctx context.Context) {
 		func() {
@@ -239,7 +236,7 @@ func (r *Replica) computeChecksumPostApply(ctx context.Context, cc kvserverpb.Co
 				snapshot = &roachpb.RaftSnapshotData{}
 			}
 
-			result, err := r.sha512(ctx, desc, snap, snapshot, cc.Mode, limiter)
+			result, err := r.sha512(ctx, desc, snap, snapshot, cc.Mode, r.store.consistencyLimiter)
 			if err != nil {
 				log.Errorf(ctx, "%v", err)
 				result = nil

pkg/kv/kvserver/replica_test.go

@@ -55,10 +55,10 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
-	"github.com/cockroachdb/cockroach/pkg/util/limit"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/metric"
 	"github.com/cockroachdb/cockroach/pkg/util/protoutil"
+	"github.com/cockroachdb/cockroach/pkg/util/quotapool"
 	"github.com/cockroachdb/cockroach/pkg/util/randutil"
 	"github.com/cockroachdb/cockroach/pkg/util/stop"
 	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
@@ -75,7 +75,6 @@ import (
 	"go.etcd.io/etcd/raft/raftpb"
 	"go.etcd.io/etcd/raft/tracker"
 	"golang.org/x/net/trace"
-	"golang.org/x/time/rate"
 )
 
 // allSpans is a SpanSet that covers *everything* for use in tests that don't
@@ -10007,7 +10006,9 @@ func TestReplicaServersideRefreshes(t *testing.T) {
 		// Regression test for #31870.
 		snap := tc.engine.NewSnapshot()
 		defer snap.Close()
-		res, err := tc.repl.sha512(context.Background(), *tc.repl.Desc(), tc.engine, nil /* diff */, roachpb.ChecksumMode_CHECK_FULL, limit.NewLimiter(rate.Inf))
+		res, err := tc.repl.sha512(context.Background(), *tc.repl.Desc(), tc.engine,
+			nil /* diff */, roachpb.ChecksumMode_CHECK_FULL,
+			quotapool.NewRateLimiter("ConsistencyQueue", quotapool.Limit(math.MaxFloat64), math.MaxInt64))
 		if err != nil {
 			return hlc.Timestamp{}, err
 		}

pkg/kv/kvserver/store.go

@@ -416,6 +416,7 @@ type Store struct {
 	tsMaintenanceQueue *timeSeriesMaintenanceQueue // Time series maintenance queue
 	scanner            *replicaScanner             // Replica scanner
 	consistencyQueue   *consistencyQueue           // Replica consistency check queue
+	consistencyLimiter *quotapool.RateLimiter      // Rate limits consistency checks
 	metrics            *StoreMetrics
 	intentResolver     *intentresolver.IntentResolver
 	recoveryMgr        txnrecovery.Manager
@@ -1576,6 +1577,17 @@ func (s *Store) Start(ctx context.Context, stopper *stop.Stopper) error {
 		s.storeRebalancer.Start(ctx, s.stopper)
 	}
 
+	s.consistencyLimiter = quotapool.NewRateLimiter(
+		"ConsistencyQueue",
+		quotapool.Limit(consistencyCheckRate.Get(&s.ClusterSettings().SV)),
+		consistencyCheckRate.Get(&s.ClusterSettings().SV)*consistencyCheckRateBurstFactor,
+		quotapool.WithMinimumWait(consistencyCheckRateMinWait))
+
+	consistencyCheckRate.SetOnChange(&s.ClusterSettings().SV, func() {
+		rate := consistencyCheckRate.Get(&s.ClusterSettings().SV)
+		s.consistencyLimiter.UpdateLimit(quotapool.Limit(rate), rate*consistencyCheckRateBurstFactor)
+	})
+
 	// Storing suggested compactions in the store itself was deprecated with
 	// the removal of the Compactor in 21.1. See discussion in
 	// https://github.com/cockroachdb/cockroach/pull/55893

pkg/util/limit/BUILD.bazel

@@ -2,31 +2,22 @@ load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 
 go_library(
     name = "limit",
-    srcs = [
-        "limiter.go",
-        "rate_limiter.go",
-    ],
+    srcs = ["limiter.go"],
     importpath = "github.com/cockroachdb/cockroach/pkg/util/limit",
     visibility = ["//visibility:public"],
     deps = [
         "//pkg/util/tracing",
         "//vendor/github.com/marusama/semaphore",
-        "//vendor/golang.org/x/time/rate",
     ],
 )
 
 go_test(
     name = "limit_test",
-    srcs = [
-        "limiter_test.go",
-        "rate_limiter_test.go",
-    ],
+    srcs = ["limiter_test.go"],
     embed = [":limit"],
     deps = [
         "//pkg/util/leaktest",
         "//vendor/github.com/cockroachdb/errors",
-        "//vendor/github.com/stretchr/testify/require",
         "//vendor/golang.org/x/sync/errgroup",
-        "//vendor/golang.org/x/time/rate",
     ],
 )