admission,kvserver: allow specification of a min-rate for snapshot disk writes

docs/generated/settings/settings.html

@@ -140,6 +140,7 @@
 <tr><td><div id="setting-kv-transaction-write-pipelining-max-batch-size" class="anchored"><code>kv.transaction.write_pipelining.max_batch_size<br />(alias: kv.transaction.write_pipelining_max_batch_size)</code></div></td><td>integer</td><td><code>128</code></td><td>if non-zero, defines that maximum size batch that will be pipelined through Raft consensus</td><td>Basic/Standard/Advanced/Self-Hosted</td></tr>
 <tr><td><div id="setting-kvadmission-store-provisioned-bandwidth" class="anchored"><code>kvadmission.store.provisioned_bandwidth</code></div></td><td>byte size</td><td><code>0 B</code></td><td>if set to a non-zero value, this is used as the provisioned bandwidth (in bytes/s), for each store. It can be overridden on a per-store basis using the --store flag. Note that setting the provisioned bandwidth to a positive value may enable disk bandwidth based admission control, since admission.disk_bandwidth_tokens.elastic.enabled defaults to true</td><td>Advanced/Self-Hosted</td></tr>
 <tr><td><div id="setting-kvadmission-store-snapshot-ingest-bandwidth-control-enabled" class="anchored"><code>kvadmission.store.snapshot_ingest_bandwidth_control.enabled</code></div></td><td>boolean</td><td><code>true</code></td><td>if set to true, snapshot ingests will be subject to disk write control in AC</td><td>Advanced/Self-Hosted</td></tr>
+<tr><td><div id="setting-kvadmission-store-snapshot-ingest-bandwidth-control-min-rate-enabled" class="anchored"><code>kvadmission.store.snapshot_ingest_bandwidth_control.min_rate.enabled</code></div></td><td>boolean</td><td><code>true</code></td><td>if set to true, snapshot ingests will be admitted at a minimum rate when kvadmission.store.provisioned_bandwidth is set to a non-zero value. Disabling this setting can lead to snapshots being starved out by foreground traffic.</td><td>Advanced/Self-Hosted</td></tr>
 <tr><td><div id="setting-log-channel-compatibility-mode-enabled" class="anchored"><code>log.channel_compatibility_mode.enabled</code></div></td><td>boolean</td><td><code>false</code></td><td>when true, logs will to log to their legacy (pre 26.1) logging channels; when false, logs will be logged to new logging channels</td><td>Basic/Standard/Advanced/Self-Hosted</td></tr>
 <tr><td><div id="setting-obs-tablemetadata-automatic-updates-enabled" class="anchored"><code>obs.tablemetadata.automatic_updates.enabled</code></div></td><td>boolean</td><td><code>false</code></td><td>enables automatic updates of the table metadata cache system.table_metadata</td><td>Basic/Standard/Advanced/Self-Hosted</td></tr>
 <tr><td><div id="setting-obs-tablemetadata-data-valid-duration" class="anchored"><code>obs.tablemetadata.data_valid_duration</code></div></td><td>duration</td><td><code>20m0s</code></td><td>the duration for which the data in system.table_metadata is considered valid</td><td>Basic/Standard/Advanced/Self-Hosted</td></tr>

pkg/kv/kvserver/kv_snapshot_strategy.go

@@ -18,6 +18,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/storage/fs"
 	"github.com/cockroachdb/cockroach/pkg/util/admission"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/cockroach/pkg/util/tracing"
 	"github.com/cockroachdb/cockroach/pkg/util/uuid"
 	"github.com/cockroachdb/errors"
@@ -142,7 +143,37 @@ func (kvSS *kvBatchSnapshotStrategy) Receive(
 	// Using a nil pacer is effectively a noop if snapshot control is disabled.
 	var pacer *admission.SnapshotPacer = nil
 	if admission.DiskBandwidthForSnapshotIngest.Get(&s.cfg.Settings.SV) && snapshotQ != nil {
-		pacer = admission.NewSnapshotPacer(snapshotQ)
+		minRate := int64(0)
+		if admission.DiskBandwidthForSnapshotIngestMinRateEnabled.Get(&s.cfg.Settings.SV) {
+			minFractionOfTimeoutForApplyingSnapshot := 1 - snapshotReservationQueueTimeoutFraction.Get(&s.cfg.Settings.SV)
+			// Use a slowdown factor of half the permittedRangeScanSlowdown
+			// factor, which means we allow snapshots to ingest atleast as fast
+			// as would be necessary to complete within half of a snapshots
+			// timeout duration.
+			snapshotApplySlowdownFactor := (permittedRangeScanSlowdown / 2) * minFractionOfTimeoutForApplyingSnapshot
+			if snapshotApplySlowdownFactor < 1 {
+				// Avoid division by 0. A snapshotApplySlowdownFactor between 0
+				// and 1 would cause the minRate to be greater than
+				// rebalanceSnapshotRate, which is the max speed snapshots can
+				// be sent at, so we don't need to ingest snapshots at a faster
+				// rate than that.
+				snapshotApplySlowdownFactor = 1
+			}
+
+			minRate = int64(float64(rebalanceSnapshotRate.Get(&s.cfg.Settings.SV)) / snapshotApplySlowdownFactor)
+			storeBW := s.cfg.KVAdmissionController.GetProvisionedBandwidth(s.StoreID())
+			if storeBW > 0 {
+				if minRate > int64(float64(storeBW)*0.25) {
+					log.KvDistribution.Warningf(ctx,
+						"snapshot ingest minRate is greater than 25%% of the store's provisioned bandwidth, minrate: %d, storeBW: %d",
+						minRate, storeBW,
+					)
+				}
+			}
+		}
+
+		timer := &timeutil.Timer{}
+		pacer = admission.NewSnapshotPacer(snapshotQ, minRate, timer.AsTimerI())
 	}
 
 	for {

pkg/kv/kvserver/kvadmission/kvadmission.go

@@ -193,6 +193,10 @@ type Controller interface {
 	// GetSnapshotQueue returns the SnapshotQueue which is used for ingesting raft
 	// snapshots.
 	GetSnapshotQueue(roachpb.StoreID) *admission.SnapshotQueue
+	// GetProvisionedBandwidth returns the provisioned disk bandwidth for the
+	// given store in bytes/second, or 0 if the store is not found or bandwidth
+	// is not configured.
+	GetProvisionedBandwidth(roachpb.StoreID) int64
 }
 
 // TenantWeightProvider can be periodically asked to provide the tenant
@@ -602,6 +606,10 @@ func (n *controllerImpl) GetSnapshotQueue(storeID roachpb.StoreID) *admission.Sn
 	return sq.(*admission.SnapshotQueue)
 }
 
+func (n *controllerImpl) GetProvisionedBandwidth(storeID roachpb.StoreID) int64 {
+	return n.storeGrantCoords.TryGetProvisionedBandwidthForStore(storeID)
+}
+
 // FollowerStoreWriteBytes captures stats about writes done to a store by a
 // replica that is not the leaseholder. These are used for admission control.
 type FollowerStoreWriteBytes struct {

pkg/util/admission/disk_bandwidth.go

@@ -110,6 +110,11 @@ func newDiskBandwidthLimiter() *diskBandwidthLimiter {
 	}
 }
 
+// getProvisionedBandwidth returns the provisioned disk bandwidth in bytes/second.
+func (d *diskBandwidthLimiter) getProvisionedBandwidth() int64 {
+	return d.state.diskLoad.intProvisionedDiskBytes / adjustmentInterval
+}
+
 // diskTokens tokens represent actual bytes and IO on physical disks. Currently,
 // these are used to impose disk bandwidth limits on elastic traffic, but
 // regular traffic will also deduct from these buckets.

pkg/util/admission/snapshot_queue.go

@@ -63,6 +63,16 @@ var DiskBandwidthForSnapshotIngest = settings.RegisterBoolSetting(
 	settings.WithPublic,
 )
 
+var DiskBandwidthForSnapshotIngestMinRateEnabled = settings.RegisterBoolSetting(
+	settings.SystemOnly,
+	"kvadmission.store.snapshot_ingest_bandwidth_control.min_rate.enabled",
+	"if set to true, snapshot ingests will be admitted at a minimum rate when "+
+		"kvadmission.store.provisioned_bandwidth is set to a non-zero value. Disabling this "+
+		"setting can lead to snapshots being starved out by foreground traffic.",
+	true,
+	settings.WithPublic,
+)
+
 var snapshotWaitDur = metric.Metadata{
 	Name:        "admission.wait_durations.snapshot_ingest",
 	Help:        "Wait time for snapshot ingest requests that waited",
@@ -71,7 +81,8 @@ var snapshotWaitDur = metric.Metadata{
 }
 
 type SnapshotMetrics struct {
-	WaitDurations metric.IHistogram
+	WaitDurations         metric.IHistogram
+	AdmittedSnapshotBytes metric.Counter
 }
 
 func makeSnapshotQueueMetrics(registry *metric.Registry) *SnapshotMetrics {
@@ -82,14 +93,20 @@ func makeSnapshotQueueMetrics(registry *metric.Registry) *SnapshotMetrics {
 			Duration:     base.DefaultHistogramWindowInterval(),
 			BucketConfig: metric.IOLatencyBuckets,
 		}),
+		AdmittedSnapshotBytes: *metric.NewCounter(metric.Metadata{
+			Name:        "admission.admitted_snapshot_bytes",
+			Help:        "Number of bytes admitted for snapshot ingests when provisioned bandwidth AC is enabled",
+			Measurement: "Bytes",
+			Unit:        metric.Unit_BYTES,
+		}),
 	}
 	registry.AddMetricStruct(m)
 	return m
 }
 
 // snapshotRequester is a wrapper used for test purposes.
 type snapshotRequester interface {
-	Admit(ctx context.Context, count int64) error
+	Admit(ctx context.Context, count int64, minRate int64, timerForMinRate timeutil.TimerI) error
 }
 
 // SnapshotQueue implements the requester interface. It is used to request
@@ -165,7 +182,14 @@ func (s *SnapshotQueue) close() {
 // Admit is called whenever a snapshot ingest request needs to update the number
 // of byte tokens it is using. Note that it accepts negative values, in which
 // case it will return the tokens back to the granter.
-func (s *SnapshotQueue) Admit(ctx context.Context, count int64) error {
+func (s *SnapshotQueue) Admit(
+	ctx context.Context, count int64, minRate int64, timerForMinRate timeutil.TimerI,
+) (err error) {
+	defer func() {
+		if err == nil && count > 0 {
+			s.metrics.AdmittedSnapshotBytes.Inc(count)
+		}
+	}()
 	if count == 0 {
 		return nil
 	}
@@ -176,6 +200,8 @@ func (s *SnapshotQueue) Admit(ctx context.Context, count int64) error {
 	if s.snapshotGranter.tryGet(canBurst /*arbitrary*/, count) {
 		return nil
 	}
+	// INVARIANT: count > 0.
+
 	// We were unable to get tokens for admission, so we queue.
 	//
 	// Reminder: there is a race here where a call to hasWaitingRequests after
@@ -198,6 +224,11 @@ func (s *SnapshotQueue) Admit(ctx context.Context, count int64) error {
 		s.addLocked(item)
 	}()
 
+	if minRate != 0 {
+		maxWaitDuration := (time.Second * time.Duration(count)) / time.Duration(minRate)
+		timerForMinRate.Reset(maxWaitDuration)
+		defer timerForMinRate.Stop()
+	}
 	// Start waiting for admission.
 	select {
 	case <-ctx.Done():
@@ -236,6 +267,28 @@ func (s *SnapshotQueue) Admit(ctx context.Context, count int64) error {
 		waitDur := timeutil.Since(item.enqueueingTime).Nanoseconds()
 		s.metrics.WaitDurations.RecordValue(waitDur)
 		return nil
+	case t := <-timerForMinRate.Ch():
+		waitDur := t.Sub(item.enqueueingTime).Nanoseconds()
+		// INVARIANT: tokensToSubtract >= 0, since item.count >= 0.
+		tokensToSubtract := item.count
+		func() {
+			s.mu.Lock()
+			defer s.mu.Unlock()
+			if item.mu.granted {
+				// NB: we must call snapshotGranter.tookWithoutPermission after
+				// releasing the mutex.
+				tokensToSubtract = 0
+			}
+			// NB: See the ctx.Done() case above for why we mark the item as
+			// cancelled instead of removing the item from the queue.
+			item.mu.cancelled = true
+		}()
+		if tokensToSubtract != 0 {
+			s.snapshotGranter.tookWithoutPermission(tokensToSubtract)
+		}
+		shouldRelease = false
+		s.metrics.WaitDurations.RecordValue(waitDur)
+		return nil
 	}
 }
 
@@ -288,14 +341,20 @@ func newSnapshotWorkItem(count int64) *snapshotWorkItem {
 }
 
 type SnapshotPacer struct {
-	snapshotQ     snapshotRequester
-	intWriteBytes int64
+	snapshotQ       snapshotRequester
+	intWriteBytes   int64
+	minRate         int64
+	timerForMinRate timeutil.TimerI
 }
 
-func NewSnapshotPacer(q snapshotRequester) *SnapshotPacer {
+func NewSnapshotPacer(
+	q snapshotRequester, minRate int64, timerForMinRate timeutil.TimerI,
+) *SnapshotPacer {
 	return &SnapshotPacer{
-		snapshotQ:     q,
-		intWriteBytes: 0,
+		snapshotQ:       q,
+		intWriteBytes:   0,
+		minRate:         minRate,
+		timerForMinRate: timerForMinRate,
 	}
 }
 
@@ -308,7 +367,7 @@ func (p *SnapshotPacer) Pace(ctx context.Context, writeBytes int64, final bool)
 	if p.intWriteBytes <= SnapshotBurstSize && !final {
 		return nil
 	}
-	if err := p.snapshotQ.Admit(ctx, p.intWriteBytes); err != nil {
+	if err := p.snapshotQ.Admit(ctx, p.intWriteBytes, p.minRate, p.timerForMinRate); err != nil {
 		return errors.Wrapf(err, "snapshot admission queue")
 	}
 	p.intWriteBytes = 0

pkg/util/admission/snapshot_queue_test.go

@@ -37,9 +37,11 @@ func TestSnapshotQueue(t *testing.T) {
 	var tg *testGranter
 	var buf builderWithMu
 	var wrkMap workMap
+	var minRate int64
 	initialTime := timeutil.FromUnixMicros(int64(0))
 	registry := metric.NewRegistry()
 	metrics := makeSnapshotQueueMetrics(registry)
+	ts := timeutil.NewManualTime(time.Unix(0, 0))
 
 	datadriven.RunTest(t, datapathutils.TestDataPath(t, "snapshot_queue"),
 		func(t *testing.T, d *datadriven.TestData) string {
@@ -66,16 +68,18 @@ func TestSnapshotQueue(t *testing.T) {
 				q.ts.(*timeutil.ManualTime).AdvanceTo(timeutil.FromUnixNanos(int64(createTime) * time.Millisecond.Nanoseconds()))
 				ctx, cancel := context.WithCancel(context.Background())
 				wrkMap.set(id, &testWork{cancel: cancel})
-				go func(ctx context.Context, id int, count int) {
-					err := q.Admit(ctx, int64(count))
+				// Create a fresh timer for each admit call so minRate timing works correctly.
+				timer := ts.NewTimer()
+				go func(ctx context.Context, id int, count int, minRate int64) {
+					err := q.Admit(ctx, int64(count), minRate, timer)
 					if err != nil {
 						buf.printf("id %d: admit failed", id)
 						wrkMap.delete(id)
 					} else {
 						buf.printf("id %d: admit succeeded", id)
 						wrkMap.setAdmitted(id, StoreWorkHandle{})
 					}
-				}(ctx, id, count)
+				}(ctx, id, count, minRate)
 				// Need deterministic output, and this is racing with the goroutine
 				// which is trying to get admitted. Retry to let it get scheduled.
 				maybeRetryWithWait(t, d.Expected, d.Rewrite, buf.String)
@@ -121,6 +125,20 @@ func TestSnapshotQueue(t *testing.T) {
 				})
 				return strconv.FormatBool(q.empty())
 
+			case "set-min-rate":
+				var v int
+				d.ScanArgs(t, "v", &v)
+				minRate = int64(v)
+				return ""
+
+			case "advance-time":
+				var millis int
+				d.ScanArgs(t, "millis", &millis)
+				ts.Advance(time.Duration(millis) * time.Millisecond)
+				// Need deterministic output since advancing time may trigger timer callbacks.
+				maybeRetryWithWait(t, d.Expected, d.Rewrite, buf.String)
+				return buf.stringAndReset()
+
 			default:
 				return fmt.Sprintf("unknown command: %s", d.Cmd)
 			}
@@ -138,7 +156,11 @@ func TestSnapshotPacer(t *testing.T) {
 	require.NoError(t, pacer.Pace(ctx, 1, false))
 
 	q := &testingSnapshotQueue{}
-	pacer = NewSnapshotPacer(q)
+	ts := timeutil.NewManualTime(time.Unix(0, 0))
+	timer := ts.NewTimer()
+
+	var minRate int64 = 10 << 20 // 10 MB/s
+	pacer = NewSnapshotPacer(q, minRate, timer)
 
 	// Should not ask for admission since write bytes = burst size.
 	writeBytes := int64(SnapshotBurstSize)
@@ -150,48 +172,59 @@ func TestSnapshotPacer(t *testing.T) {
 	// Do another write, should go over threshold and seek admission.
 	require.NoError(t, pacer.Pace(ctx, 1, false))
 	require.True(t, q.admitted)
+	require.Equal(t, minRate, q.minRate)
 	require.Equal(t, int64(0), pacer.intWriteBytes)
 	require.Equal(t, writeBytes+1, q.admitCount)
 
 	// Not enough bytes since last admission. Should not ask for admission.
 	q.admitted = false
 	q.admitCount = 0
+	q.minRate = 0
 	require.NoError(t, pacer.Pace(ctx, 5, false))
 	require.False(t, q.admitted)
+	require.Equal(t, int64(0), q.minRate)
 	require.Equal(t, int64(5), pacer.intWriteBytes)
 	require.Equal(t, int64(0), q.admitCount)
 
 	// We now go above the threshold again. Should ask for admission.
 	require.NoError(t, pacer.Pace(ctx, writeBytes, false))
 	require.True(t, q.admitted)
+	require.Equal(t, minRate, q.minRate)
 	require.Equal(t, writeBytes+5, q.admitCount)
 	require.Equal(t, int64(0), pacer.intWriteBytes)
 
 	// Do few more writes.
 	q.admitted = false
 	q.admitCount = 0
+	q.minRate = 0
 	require.NoError(t, pacer.Pace(ctx, 10, false))
 	require.False(t, q.admitted)
 	require.Equal(t, int64(10), pacer.intWriteBytes)
+	require.Equal(t, int64(0), q.minRate)
 	require.Equal(t, int64(0), q.admitCount)
 
 	// If final call to pacer, we should admit regardless of size. It should flush
 	// all intWriteBytes.
 	require.NoError(t, pacer.Pace(ctx, -1, true))
 	require.True(t, q.admitted)
+	require.Equal(t, minRate, q.minRate)
 	require.Equal(t, int64(9), q.admitCount)
 }
 
 // testingSnapshotQueue is used to test SnapshotPacer.
 type testingSnapshotQueue struct {
 	admitted   bool
 	admitCount int64
+	minRate    int64
 }
 
 var _ snapshotRequester = &testingSnapshotQueue{}
 
-func (ts *testingSnapshotQueue) Admit(ctx context.Context, count int64) error {
+func (ts *testingSnapshotQueue) Admit(
+	ctx context.Context, count int64, minRate int64, timerForMinRate timeutil.TimerI,
+) error {
 	ts.admitted = true
 	ts.admitCount = count
+	ts.minRate = minRate
 	return nil
 }

pkg/util/admission/store_grant_coordinator.go

@@ -298,6 +298,18 @@ func (sgc *StoreGrantCoordinators) TryGetSnapshotQueueForStore(storeID roachpb.S
 	return nil
 }
 
+// TryGetProvisionedBandwidthForStore returns the provisioned bandwidth for the
+// given store in bytes/second, or 0 if the store is not found or bandwidth is
+// not configured.
+func (sgc *StoreGrantCoordinators) TryGetProvisionedBandwidthForStore(
+	storeID roachpb.StoreID,
+) int64 {
+	if gc, ok := sgc.gcMap.Load(storeID); ok {
+		return gc.ioLoadListener.diskBandwidthLimiter.getProvisionedBandwidth()
+	}
+	return 0
+}
+
 func (sgc *StoreGrantCoordinators) close() {
 	// closeCh can be nil in tests that never called SetPebbleMetricsProvider.
 	if sgc.closeCh != nil {