admission: fix overflow bug in ioLoadListener

pkg/util/admission/granter.go

@@ -1463,7 +1463,9 @@ func (io *ioLoadListener) pebbleMetricsTick(m pebble.Metrics) {
 // 1s.
 func (io *ioLoadListener) allocateTokensTick() {
 	var toAllocate int64
-	if io.totalTokens == unlimitedTokens {
+	// unlimitedTokens==MaxInt64, so avoid overflow in the rounding up
+	// calculation.
+	if io.totalTokens >= unlimitedTokens-(adjustmentInterval-1) {
 		toAllocate = io.totalTokens / adjustmentInterval
 	} else {
 		// Round up so that we don't accumulate tokens to give in a burst on the

pkg/util/admission/granter_test.go

@@ -13,6 +13,7 @@ package admission
 import (
 	"context"
 	"fmt"
+	"math"
 	"sort"
 	"strings"
 	"testing"
@@ -376,6 +377,26 @@ func TestIOLoadListener(t *testing.T) {
 		})
 }
 
+func TestIOLoadListenerOverflow(t *testing.T) {
+	req := &testRequesterForIOLL{}
+	kvGranter := &testGranterWithIOTokens{}
+	st := cluster.MakeTestingClusterSettings()
+	ioll := ioLoadListener{
+		settings:    st,
+		kvRequester: req,
+	}
+	ioll.mu.Mutex = &syncutil.Mutex{}
+	ioll.mu.kvGranter = kvGranter
+	for i := int64(0); i < adjustmentInterval; i++ {
+		// Override the totalTokens manually to trigger the overflow bug.
+		ioll.totalTokens = math.MaxInt64 - i
+		ioll.tokensAllocated = 0
+		for j := 0; j < adjustmentInterval; j++ {
+			ioll.allocateTokensTick()
+		}
+	}
+}
+
 // TODO(sumeer):
 // - Test metrics
 // - Test GrantCoordinator with multi-tenant configurations