From 716d2ca2b13c6559bd97ea9aa65f083d0164f2b5 Mon Sep 17 00:00:00 2001
From: sumeerbhola <sumeer@cockroachlabs.com>
Date: Mon, 30 Aug 2021 09:37:59 -0400
Subject: [PATCH] admission: fix overflow bug in ioLoadListener

Fixes #69461,#69463

Release justification: High-severity bug fix to new funtionality.

Release note: None
---
 pkg/util/admission/granter.go      |  4 +++-
 pkg/util/admission/granter_test.go | 21 +++++++++++++++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/pkg/util/admission/granter.go b/pkg/util/admission/granter.go
index fcc2d3c1abd6..d502bc4e2622 100644
--- a/pkg/util/admission/granter.go
+++ b/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
diff --git a/pkg/util/admission/granter_test.go b/pkg/util/admission/granter_test.go
index b6c694385b7e..fdc1518543d7 100644
--- a/pkg/util/admission/granter_test.go
+++ b/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