@@ -22,6 +22,7 @@ package txthrottler
2222//go:generate mockgen -destination mock_topology_watcher_test.go -package txthrottler vitess.io/vitess/go/vt/vttablet/tabletserver/txthrottler TopologyWatcherInterface
2323
2424import (
25+ "sync/atomic"
2526 "testing"
2627 "time"
2728
@@ -84,32 +85,50 @@ func TestEnabledThrottler(t *testing.T) {
8485 }
8586
8687 mockThrottler := NewMockThrottlerInterface (mockCtrl )
88+
8789 throttlerFactory = func (name , unit string , threadCount int , maxRate int64 , maxReplicationLagConfig throttler.MaxReplicationLagModuleConfig ) (ThrottlerInterface , error ) {
8890 assert .Equal (t , 1 , threadCount )
8991 return mockThrottler , nil
9092 }
9193
92- call0 := mockThrottler .EXPECT ().UpdateConfiguration (gomock .Any (), true /* copyZeroValues */ )
93- call1 := mockThrottler .EXPECT ().Throttle (0 )
94- call1 .Return (0 * time .Second )
94+ var calls []* gomock.Call
95+
96+ call := mockThrottler .EXPECT ().UpdateConfiguration (gomock .Any (), true /* copyZeroValues */ )
97+ calls = append (calls , call )
98+
99+ // 1
100+ call = mockThrottler .EXPECT ().Throttle (0 )
101+ call .Return (0 * time .Second )
102+ calls = append (calls , call )
103+
95104 tabletStats := & discovery.LegacyTabletStats {
96105 Target : & querypb.Target {
97106 TabletType : topodatapb .TabletType_REPLICA ,
98107 },
99108 }
100- call2 := mockThrottler .EXPECT ().RecordReplicationLag (gomock .Any (), tabletStats )
101- call3 := mockThrottler .EXPECT ().Throttle (0 )
102- call3 .Return (1 * time .Second )
103109
104- call4 := mockThrottler .EXPECT ().Throttle (0 )
105- call4 .Return (1 * time .Second )
106- calllast := mockThrottler .EXPECT ().Close ()
110+ call = mockThrottler .EXPECT ().RecordReplicationLag (gomock .Any (), tabletStats )
111+ calls = append (calls , call )
107112
108- call1 .After (call0 )
109- call2 .After (call1 )
110- call3 .After (call2 )
111- call4 .After (call3 )
112- calllast .After (call4 )
113+ // 2
114+ call = mockThrottler .EXPECT ().Throttle (0 )
115+ call .Return (1 * time .Second )
116+ calls = append (calls , call )
117+
118+ // 3
119+ // Nothing gets mocked here because the order of evaluation in txThrottler.Throttle() evaluates first
120+ // whether the priority allows for throttling or not, so no need to mock calls in mockThrottler.Throttle()
121+
122+ // 4
123+ // Nothing gets mocked here because the order of evaluation in txThrottlerStateImpl.Throttle() evaluates first
124+ // whether there is lag or not, so no call to the underlying mockThrottler is issued.
125+
126+ call = mockThrottler .EXPECT ().Close ()
127+ calls = append (calls , call )
128+
129+ for i := 1 ; i < len (calls ); i ++ {
130+ calls [i ].After (calls [i - 1 ])
131+ }
113132
114133 config := tabletenv .NewDefaultConfig ()
115134 config .EnableTxThrottler = true
@@ -126,6 +145,14 @@ func TestEnabledThrottler(t *testing.T) {
126145 assert .Nil (t , throttler .Open ())
127146 assert .Equal (t , int64 (1 ), throttler .throttlerRunning .Get ())
128147
148+ throttlerImpl , ok := throttler .state .(* txThrottlerStateImpl )
149+ assert .True (t , ok )
150+ // Stop the go routine that keeps updating the cached shard's max lag to preventi it from changing the value in a
151+ // way that will interfere with how we manipulate that value in our tests to evaluate different cases:
152+ throttlerImpl .done <- true
153+
154+ // 1 should not throttle due to return value of underlying Throttle(), despite high lag
155+ atomic .StoreInt64 (& throttlerImpl .maxLag , 20 )
129156 assert .False (t , throttler .Throttle (100 , "some-workload" ))
130157 assert .Equal (t , int64 (1 ), throttler .requestsTotal .Counts ()["some-workload" ])
131158 assert .Zero (t , throttler .requestsThrottled .Counts ()["some-workload" ])
@@ -138,15 +165,23 @@ func TestEnabledThrottler(t *testing.T) {
138165 }
139166 // This call should not be forwarded to the go/vt/throttler.Throttler object.
140167 hcListener .StatsUpdate (rdonlyTabletStats )
141- // The second throttle call should reject.
168+
169+ // 2 should throttle due to return value of underlying Throttle(), high lag & priority = 100
142170 assert .True (t , throttler .Throttle (100 , "some-workload" ))
143171 assert .Equal (t , int64 (2 ), throttler .requestsTotal .Counts ()["some-workload" ])
144172 assert .Equal (t , int64 (1 ), throttler .requestsThrottled .Counts ()["some-workload" ])
145173
146- // This call should not throttle due to priority. Check that's the case and counters agree.
174+ // 3 should not throttle despite return value of underlying Throttle() and high lag, due to priority = 0
147175 assert .False (t , throttler .Throttle (0 , "some-workload" ))
148176 assert .Equal (t , int64 (3 ), throttler .requestsTotal .Counts ()["some-workload" ])
149177 assert .Equal (t , int64 (1 ), throttler .requestsThrottled .Counts ()["some-workload" ])
178+
179+ // 4 should not throttle despite return value of underlying Throttle() and priority = 100, due to low lag
180+ atomic .StoreInt64 (& throttlerImpl .maxLag , 1 )
181+ assert .False (t , throttler .Throttle (100 , "some-workload" ))
182+ assert .Equal (t , int64 (4 ), throttler .requestsTotal .Counts ()["some-workload" ])
183+ assert .Equal (t , int64 (1 ), throttler .requestsThrottled .Counts ()["some-workload" ])
184+
150185 throttler .Close ()
151186 assert .Zero (t , throttler .throttlerRunning .Get ())
152187}
0 commit comments