admission: compute cpu time token refill rates

This commit introduces a linear model that computes cpu
time token refill rates. cpuTimeTokenFiller uses this
model to determine how many tokens to add per second to
the buckets in cpuTimeTokenGranter.

Fixes: #154471

Release note: None.

Author: joshimhoff

pkg/util/admission/cpu_time_token_filler.go

@@ -272,7 +272,7 @@ func (m *cpuTimeTokenLinearModel) init() {
 	_, cpuCapacity := m.cpuMetricProvider.GetCPUInfo()
 	m.cpuCapacity = cpuCapacity
 	m.tokenToCPUTimeMultiplier = 1
-	_ = m.calculateRefillRates()
+	_ = m.updateRefillRates()
 }
 
 // fit adjusts tokenToCPUTimeMultiplier based on CPU usage & token usage. fit
@@ -344,28 +344,28 @@ func (m *cpuTimeTokenLinearModel) fit() [numResourceTiers][numBurstQualification
 			alpha*tokenToCPUTimeMultiplier + (1-alpha)*m.tokenToCPUTimeMultiplier
 	}
 
-	return m.calculateRefillRates()
+	return m.updateRefillRates()
 }
 
-// calculateRefillRates computes refill rates from tokenToCPUTimeMultiplier &
-// the admission.cpu_time_tokens.target_util cluster setting. calculateRefillRates
-// returns the delta refill rates. That is calculateRefillRates returns the difference
+// updateRefillRates computes refill rates from tokenToCPUTimeMultiplier &
+// the admission.cpu_time_tokens.target_util cluster setting. updateRefillRates
+// returns the delta refill rates. That is updateRefillRates returns the difference
 // in tokens to add per interval (1s) from previous call to fit to this one.
-func (m *cpuTimeTokenLinearModel) calculateRefillRates() [numResourceTiers][numBurstQualifications]int64 {
-	// Compute goals from cluster setting.
+func (m *cpuTimeTokenLinearModel) updateRefillRates() [numResourceTiers][numBurstQualifications]int64 {
+	// Compute goals from cluster setting. Algorithmically, it is okay if some of
+	// the below goalUtils are greater than 1. This would mean greater risk of
+	// goroutine scheduling latency, but there is no immediate problem -- the
+	// greater some goalUtil is, the more CPU time tokens will be in the corresponding
+	// bucket.
 	var goalUtils [numResourceTiers][numBurstQualifications]float64
 	util := KVCPUTimeUtilGoal.Get(&m.settings.SV)
-	goalUtils[appTenant][noBurst] = util
-	// Algorithmically, it is okay if some of the below goalUtils are greater
-	// than 1. This would mean greater risk of goroutine scheduling latency, but
-	// there is no immediate problem -- the greater some goalUtil is, the more CPU
-	// time tokens will be in the corresponding bucket.
-	goalUtils[appTenant][canBurst] = util + 0.05
-	goalUtils[systemTenant][noBurst] = util + 0.1
-	// The system tenant will never set canBurst = true, so we do not fill
-	// that bucket with tokens.
-	// TODO(): Check that this is fine.
-	goalUtils[systemTenant][canBurst] = 0
+	var iter float64
+	for tier := int(numResourceTiers - 1); tier >= 0; tier-- {
+		for qual := int(numBurstQualifications - 1); qual >= 0; qual-- {
+			goalUtils[tier][qual] = util + 0.05*iter
+			iter++
+		}
+	}
 
 	// Update refill rates. Return change in rates via delta.
 	var delta [numResourceTiers][numBurstQualifications]int64

pkg/util/admission/cpu_time_token_filler_test.go

@@ -17,6 +17,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/datadriven"
+	"github.com/stretchr/testify/require"
 )
 
 func TestCPUTimeTokenFiller(t *testing.T) {
@@ -81,13 +82,13 @@ func (a *testTokenAllocator) allocateTokens(remainingTicks int64) {
 
 type testModel struct {
 	rates [numResourceTiers][numBurstQualifications]int64
+	delta [numResourceTiers][numBurstQualifications]int64
 }
 
 func (m *testModel) init() {}
 
 func (m *testModel) fit() [numResourceTiers][numBurstQualifications]int64 {
-	// TODO(): Test ret value.
-	return [numResourceTiers][numBurstQualifications]int64{}
+	return m.delta
 }
 
 func (m *testModel) getRefillRates() [numResourceTiers][numBurstQualifications]int64 {
@@ -130,10 +131,8 @@ func TestCPUTimeTokenAllocator(t *testing.T) {
 	}
 
 	var buf strings.Builder
-	flushAndReset := func(printGranter bool) string {
-		if printGranter {
-			fmt.Fprint(&buf, granter.String())
-		}
+	flushAndReset := func() string {
+		fmt.Fprint(&buf, granter.String())
 		str := buf.String()
 		buf.Reset()
 		return str
@@ -142,27 +141,40 @@ func TestCPUTimeTokenAllocator(t *testing.T) {
 	datadriven.RunTest(t, datapathutils.TestDataPath(t, "cpu_time_token_allocator"), func(t *testing.T, d *datadriven.TestData) string {
 		switch d.Cmd {
 		case "resetInterval":
-			// TODO(): Test when true also.
-			allocator.resetInterval(false)
-			return flushAndReset(false /* printGranter */)
+			var delta int64
+			d.MaybeScanArgs(t, "delta", &delta)
+			if d.MaybeScanArgs(t, "delta", &delta) {
+				for tier := range model.delta {
+					for qual := range model.delta[tier] {
+						model.delta[tier][qual] = delta
+					}
+				}
+			}
+			skipFit := d.HasArg("skipfit")
+			allocator.resetInterval(skipFit /* skipFittingLinearModel */)
+			for tier := range model.delta {
+				for qual := range model.delta[tier] {
+					model.delta[tier][qual] = 0
+				}
+			}
+			return flushAndReset()
 		case "allocate":
 			var remainingTicks int64
 			d.ScanArgs(t, "remaining", &remainingTicks)
 			allocator.allocateTokens(remainingTicks)
-			return flushAndReset(true /* printGranter */)
+			return flushAndReset()
 		case "clear":
 			granter.mu.buckets[testTier0][canBurst].tokens = 0
 			granter.mu.buckets[testTier0][noBurst].tokens = 0
 			granter.mu.buckets[testTier1][canBurst].tokens = 0
 			granter.mu.buckets[testTier1][noBurst].tokens = 0
-			return flushAndReset(true /* printGranter */)
+			return flushAndReset()
 		default:
 			return fmt.Sprintf("unknown command: %s", d.Cmd)
 		}
 	})
 }
 
-// TODO(): Complete.
 func TestCPUTimeTokenLinearModel(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)
@@ -176,122 +188,143 @@ func TestCPUTimeTokenLinearModel(t *testing.T) {
 		totalCPUTimeMillis:       0,
 		tokenToCPUTimeMultiplier: 1,
 	}
-	dur := 5 * time.Second
-	ttt := &tt{}
-	model.granter = ttt
-	m := &mm{
+	tokenCPUTime := &testTokenUsageTracker{}
+	model.granter = tokenCPUTime
+	actualCPUTime := &testCPUMetricProvider{
 		capacity: 10,
 	}
-	model.cpuMetricProvider = m
-
-	// Up & down.
-	ttt.append(dur.Nanoseconds()/2, 200)
-	ttt.append(dur.Nanoseconds()*2, 100)
-	m.append(dur.Milliseconds(), 100)
-	m.append(dur.Milliseconds()*2, 100)
-	m.append(dur.Milliseconds()*2, 100)
-	for i := 0; i < len(ttt.tokensUsed); i++ {
+	model.cpuMetricProvider = actualCPUTime
+
+	dur := 5 * time.Second
+	actualCPUTime.append(dur.Nanoseconds(), 1) // appended value ignored by init
+	model.init()
+
+	// 2x.
+	tokenCPUTime.append(dur.Nanoseconds()/2, 100)
+	actualCPUTime.append(dur.Milliseconds(), 100)
+	for i := 0; i < 100; i++ {
 		testTime.Advance(time.Second)
 		model.fit()
-		if (i+1)%100 == 0 {
-			fmt.Printf("%f\n", model.tokenToCPUTimeMultiplier)
-		}
 	}
-	// TODO(): Check rates.
-	//rates := model.getRates()
-	//want := 4 * time.Second
-	//require.Equal(t, want.Nanoseconds(), rates[0][canBurst])
-
-	// Cap 20x
-	ttt.append(dur.Nanoseconds(), 100)
-	m.append(dur.Milliseconds()*40, 100)
+	tolerance := 0.01
+	require.InDelta(t, 2, model.tokenToCPUTimeMultiplier, tolerance)
+
+	// 4x.
+	tokenCPUTime.append(dur.Nanoseconds()/2, 100)
+	actualCPUTime.append(dur.Milliseconds()*2, 100)
 	for i := 0; i < 100; i++ {
 		testTime.Advance(time.Second)
 		model.fit()
-		if i == 99 {
-			fmt.Printf("%f\n", model.tokenToCPUTimeMultiplier)
-		}
 	}
+	require.InDelta(t, 4, model.tokenToCPUTimeMultiplier, tolerance)
 
-	// Cap 1x
-	ttt.append(dur.Nanoseconds()*2, 100)
-	m.append(dur.Milliseconds(), 100)
+	// 1x.
+	tokenCPUTime.append(dur.Nanoseconds()*2, 100)
+	actualCPUTime.append(dur.Milliseconds()*2, 100)
 	for i := 0; i < 100; i++ {
 		testTime.Advance(time.Second)
 		model.fit()
-		if i == 99 {
-			fmt.Printf("%f\n", model.tokenToCPUTimeMultiplier)
-		}
 	}
+	require.InDelta(t, 1, model.tokenToCPUTimeMultiplier, tolerance)
 
-	// Low CPU mode
-	ttt = &tt{}
-	model.granter = ttt
-	m = &mm{
-		capacity: 10,
+	// Cap at 20x.
+	tokenCPUTime.append(dur.Nanoseconds(), 100)
+	actualCPUTime.append(dur.Milliseconds()*40, 100)
+	for i := 0; i < 100; i++ {
+		testTime.Advance(time.Second)
+		model.fit()
 	}
-	model.cpuMetricProvider = m
-
-	// Leave 2x as is
-	ttt.append(dur.Nanoseconds(), 100)
-	m.append(dur.Milliseconds()*2, 100)
-	ttt.append(dur.Nanoseconds()/5, 100)
-	m.append(dur.Milliseconds()/5, 100)
-	for i := 0; i < 200; i++ {
+	require.InDelta(t, 20, model.tokenToCPUTimeMultiplier, tolerance)
+
+	// Cap at 1x.
+	tokenCPUTime.append(dur.Nanoseconds()*2, 100)
+	actualCPUTime.append(dur.Milliseconds(), 100)
+	for i := 0; i < 100; i++ {
 		testTime.Advance(time.Second)
 		model.fit()
-		if i == 199 {
-			fmt.Printf("%f\n", model.tokenToCPUTimeMultiplier)
-		}
 	}
+	require.InDelta(t, 1, model.tokenToCPUTimeMultiplier, tolerance)
 
-	// Cut 10x, too high
-	ttt.append(dur.Nanoseconds(), 100)
-	m.append(dur.Milliseconds()*100, 100)
-	ttt.append(dur.Nanoseconds()/5, 100)
-	m.append(dur.Milliseconds()/5, 100)
-	for i := 0; i < 200; i++ {
+	// 2x.
+	tokenCPUTime.append(dur.Nanoseconds(), 100)
+	actualCPUTime.append(dur.Milliseconds()*2, 100)
+	for i := 0; i < 100; i++ {
+		testTime.Advance(time.Second)
+		model.fit()
+	}
+	require.InDelta(t, 2, model.tokenToCPUTimeMultiplier, tolerance)
+
+	// Leave 2x as is, even tho low CPU mode, since multiplier is already low.
+	tokenCPUTime.append(dur.Nanoseconds()/5, 100)
+	actualCPUTime.append(dur.Milliseconds()/5, 100)
+	for i := 0; i < 100; i++ {
+		testTime.Advance(time.Second)
+		model.fit()
+	}
+	require.InDelta(t, 2, model.tokenToCPUTimeMultiplier, tolerance)
+
+	// 20x.
+	tokenCPUTime.append(dur.Nanoseconds(), 100)
+	actualCPUTime.append(dur.Milliseconds()*100, 100)
+	for i := 0; i < 100; i++ {
+		testTime.Advance(time.Second)
+		model.fit()
+	}
+	require.InDelta(t, 20, model.tokenToCPUTimeMultiplier, tolerance)
+
+	// Reduce to 3.6x, since low CPU mode, and multiplier is high.
+	tokenCPUTime.append(dur.Nanoseconds()/5, 100)
+	actualCPUTime.append(dur.Milliseconds()/5, 100)
+	for i := 0; i < 100; i++ {
 		testTime.Advance(time.Second)
 		model.fit()
-		if i == 199 {
-			fmt.Printf("%f\n", model.tokenToCPUTimeMultiplier)
-		}
 	}
+	require.InDelta(t, 3.6, model.tokenToCPUTimeMultiplier, tolerance)
+
+	rates := model.getRefillRates()
+	// Hard-coded to be 0.
+	// 95% -> 10 vCPUs * .95 * 1s = 9.5s, 9.5s / 3.6 ~= 2.63888889
+	require.Equal(t, int64(2638888888), rates[systemTenant][canBurst])
+	// 90% -> 10 vCPUs * .9 * 1s = 9s, 9s / 3.6 ~= 2.5s
+	require.Equal(t, int64(2500000000), rates[systemTenant][noBurst])
+	// 85% -> 10 vCPUs * .85 * 1s = 8.5s, 8.5s / 3.6 ~= 2.36111111s
+	require.Equal(t, int64(2361111111), rates[appTenant][canBurst])
+	// 80% -> 10 vCPUs * .8 * 1s = 8s, 8s / 3.6 ~= 2.22222222s
+	require.Equal(t, int64(2222222222), rates[appTenant][noBurst])
 }
 
-type tt struct {
+type testTokenUsageTracker struct {
 	i          int
 	tokensUsed []int64
 }
 
-func (t *tt) append(tokens int64, count int) {
-	// TODO(): Modernize.
+func (t *testTokenUsageTracker) append(tokens int64, count int) {
 	for i := 0; i < count; i++ {
 		t.tokensUsed = append(t.tokensUsed, tokens)
 	}
 }
 
-func (t *tt) getTokensUsedInInterval() int64 {
+func (t *testTokenUsageTracker) getTokensUsedInInterval() int64 {
 	ret := t.tokensUsed[t.i]
 	t.i++
 	return ret
 }
 
-type mm struct {
+type testCPUMetricProvider struct {
 	i        int
 	cum      int64
 	millis   []int64
 	capacity float64
 }
 
-func (m *mm) GetCPUInfo() (int64, float64) {
+func (m *testCPUMetricProvider) GetCPUInfo() (int64, float64) {
 	cycle := m.millis[m.i]
 	m.i++
 	m.cum += cycle
 	return m.cum, m.capacity
 }
 
-func (t *mm) append(millis int64, count int) {
+func (t *testCPUMetricProvider) append(millis int64, count int) {
 	for i := 0; i < count; i++ {
 		t.millis = append(t.millis, millis)
 	}

pkg/util/admission/cpu_time_token_granter.go

@@ -253,7 +253,6 @@ func (stg *cpuTimeTokenGranter) tryGrantLocked() bool {
 
 // getTokensUsedInInterval returns the net number of tokens deducted from the
 // buckets, since the last call to getTokensUsedInInterval.
-// TODO(): Test.
 func (stg *cpuTimeTokenGranter) getTokensUsedInInterval() int64 {
 	stg.mu.Lock()
 	defer stg.mu.Unlock()

pkg/util/admission/cpu_time_token_granter_test.go

@@ -167,6 +167,11 @@ func TestCPUTimeTokenGranter(t *testing.T) {
 			fmt.Fprintf(&buf, "refill(%v %v)\n", delta, bucketCapacity)
 			return flushAndReset(false /* init */)
 
+		case "get-tokens-used":
+			used := granter.getTokensUsedInInterval()
+			fmt.Fprintf(&buf, "get-tokens-used-in-interval() returned %d\n", used)
+			return flushAndReset(false /* init */)
+
 		// For cpuTimeTokenChildGranter, this is a NOP. Still, it will be
 		// called in production. So best to test it doesn't panic, or similar.
 		case "continue-grant-chain":