Mark Snake timer transitions as parent-locked

AI assisted with development. Every line of code was either written by or carefully reviewed by me :)

Signed-off-by: Brett Wines <bwines@slack-corp.com>

Author: bgwines

go/vt/vttablet/tabletserver/loadshed/codelq.go

@@ -158,16 +158,18 @@ type (
 		cfg               CoDelConfig
 		nowNs             func() int64
 		scheduleDropTimer func(delayNs int64)
+		stopDropTimer     func()
 	}
 )
 
-func newCoDelQueue[T any](cfg CoDelConfig, nowNs func() int64, scheduleDropTimer func(delayNs int64)) *CoDelQueue[T] {
+func newCoDelQueue[T any](cfg CoDelConfig, nowNs func() int64, scheduleDropTimer func(delayNs int64), stopDropTimer func()) *CoDelQueue[T] {
 	q := &CoDelQueue[T]{
 		queue:             list.New(),
 		count:             1,
 		cfg:               cfg,
 		nowNs:             nowNs,
 		scheduleDropTimer: scheduleDropTimer,
+		stopDropTimer:     stopDropTimer,
 	}
 	q.droppable.init()
 	return q

go/vt/vttablet/tabletserver/loadshed/codelq_test.go

@@ -70,6 +70,11 @@ func (r *testDropTimerRecorder) schedule(delayNs int64) {
 	r.delayNs = delayNs
 }
 
+func (r *testDropTimerRecorder) stop() {
+	r.armed = false
+	r.scheduled = false
+}
+
 // reset models a timer fire in tests: clears both the armed flag and the
 // scheduled observability flag, so the next schedule() call is detected.
 func (r *testDropTimerRecorder) reset() {
@@ -79,7 +84,7 @@ func (r *testDropTimerRecorder) reset() {
 
 func newTestQueue(cfg CoDelConfig, clock *testClock) (*testCoDelQueue, *testDropTimerRecorder) {
 	rec := &testDropTimerRecorder{}
-	q := newCoDelQueue[struct{}](cfg, clock.nowFunc, rec.schedule)
+	q := newCoDelQueue[struct{}](cfg, clock.nowFunc, rec.schedule, rec.stop)
 	return q, rec
 }
 

go/vt/vttablet/tabletserver/loadshed/snake.go

@@ -92,7 +92,7 @@ func NewSnake[T any](cfg SnakeConfig) *Snake[T] {
 		timerLag:     stats.NewHistogram("", "", loadshedBucketCutoffs),
 		valveDepth:   stats.NewHistogram("", "", lengthBucketCutoffs),
 	}
-	s.q = newValvedCoDelQueue[T](cfg.CoDel, defaultClock, s.lockedScheduleDropTimer)
+	s.q = newValvedCoDelQueue[T](cfg.CoDel, defaultClock, s.lockedScheduleDropTimer, s.lockedStopDropTimer)
 	return s
 }
 
@@ -252,15 +252,19 @@ func (s *Snake[T]) lockedScheduleDropTimer(delayNs int64) {
 	s.dropTimerExpectedNs = s.clockFunc() + delayNs
 }
 
-func (s *Snake[T]) TimerUpdate() (time.Duration, bool) {
+func (s *Snake[T]) lockedStopDropTimer() {
+	s.dropTimerArmed = false
+}
+
+func (s *Snake[T]) LockedTimerUpdate() (time.Duration, bool) {
 	if !s.dropTimerChanged {
 		return 0, false
 	}
 	s.dropTimerChanged = false
 	return time.Duration(s.dropTimerDelayNs) * time.Nanosecond, true
 }
 
-func (s *Snake[T]) DropTimerFired() []T {
+func (s *Snake[T]) LockedDropTimerFired() []T {
 	if !s.dropTimerArmed {
 		return nil
 	}

go/vt/vttablet/tabletserver/loadshed/valved_codelq.go

@@ -100,12 +100,12 @@ type (
 	}
 )
 
-func newValvedCoDelQueue[T any](cfg CoDelConfig, nowNs func() int64, scheduleDropTimer func(delayNs int64)) *ValvedCoDelQueue[T] {
+func newValvedCoDelQueue[T any](cfg CoDelConfig, nowNs func() int64, scheduleDropTimer func(delayNs int64), stopDropTimer func()) *ValvedCoDelQueue[T] {
 	q := &ValvedCoDelQueue[T]{
 		valves:            make(map[string][]*Request[T]),
 		droppablePerValve: make(map[string]*Request[T]),
 	}
-	q.codelq = newCoDelQueue[T](cfg, nowNs, scheduleDropTimer)
+	q.codelq = newCoDelQueue[T](cfg, nowNs, scheduleDropTimer, stopDropTimer)
 	return q
 }
 

go/vt/vttablet/tabletserver/loadshed/valved_codelq_test.go

@@ -27,7 +27,7 @@ type testValvedCoDelQueue = ValvedCoDelQueue[struct{}]
 
 func newValvedQueue(clock *testClock) (*testValvedCoDelQueue, *testDropTimerRecorder) {
 	rec := &testDropTimerRecorder{}
-	q := newValvedCoDelQueue[struct{}](defaultTestConfig(), clock.nowFunc, rec.schedule)
+	q := newValvedCoDelQueue[struct{}](defaultTestConfig(), clock.nowFunc, rec.schedule, rec.stop)
 	return q, rec
 }