diff --git a/go/vt/vttablet/tabletserver/txthrottler/mock_throttler_test.go b/go/vt/vttablet/tabletserver/txthrottler/mock_throttler_test.go index a3da535037e..75daf37c5a0 100644 --- a/go/vt/vttablet/tabletserver/txthrottler/mock_throttler_test.go +++ b/go/vt/vttablet/tabletserver/txthrottler/mock_throttler_test.go @@ -1,5 +1,5 @@ // Code generated by MockGen. DO NOT EDIT. -// Source: vitess.io/vitess/go/vt/vttablet/tabletserver/txthrottler (interfaces: ThrottlerInterface) +// Source: ./go/vt/vttablet/tabletserver/txthrottler/tx_throttler.go // Package txthrottler is a generated GoMock package. package txthrottler @@ -12,6 +12,7 @@ import ( discovery "vitess.io/vitess/go/vt/discovery" throttlerdata "vitess.io/vitess/go/vt/proto/throttlerdata" + topodata "vitess.io/vitess/go/vt/proto/topodata" ) // MockThrottlerInterface is a mock of ThrottlerInterface interface. @@ -63,6 +64,20 @@ func (mr *MockThrottlerInterfaceMockRecorder) GetConfiguration() *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetConfiguration", reflect.TypeOf((*MockThrottlerInterface)(nil).GetConfiguration)) } +// LastMaxLagNotIgnoredForTabletType mocks base method. +func (m *MockThrottlerInterface) LastMaxLagNotIgnoredForTabletType(tabletType topodata.TabletType) uint32 { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "LastMaxLagNotIgnoredForTabletType", tabletType) + ret0, _ := ret[0].(uint32) + return ret0 +} + +// LastMaxLagNotIgnoredForTabletType indicates an expected call of LastMaxLagNotIgnoredForTabletType. +func (mr *MockThrottlerInterfaceMockRecorder) LastMaxLagNotIgnoredForTabletType(tabletType interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "LastMaxLagNotIgnoredForTabletType", reflect.TypeOf((*MockThrottlerInterface)(nil).LastMaxLagNotIgnoredForTabletType), tabletType) +} + // MaxRate mocks base method. func (m *MockThrottlerInterface) MaxRate() int64 { m.ctrl.T.Helper() @@ -78,15 +93,15 @@ func (mr *MockThrottlerInterfaceMockRecorder) MaxRate() *gomock.Call { } // RecordReplicationLag mocks base method. -func (m *MockThrottlerInterface) RecordReplicationLag(arg0 time.Time, arg1 *discovery.LegacyTabletStats) { +func (m *MockThrottlerInterface) RecordReplicationLag(time time.Time, ts *discovery.LegacyTabletStats) { m.ctrl.T.Helper() - m.ctrl.Call(m, "RecordReplicationLag", arg0, arg1) + m.ctrl.Call(m, "RecordReplicationLag", time, ts) } // RecordReplicationLag indicates an expected call of RecordReplicationLag. -func (mr *MockThrottlerInterfaceMockRecorder) RecordReplicationLag(arg0, arg1 any) *gomock.Call { +func (mr *MockThrottlerInterfaceMockRecorder) RecordReplicationLag(time, ts interface{}) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RecordReplicationLag", reflect.TypeOf((*MockThrottlerInterface)(nil).RecordReplicationLag), arg0, arg1) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RecordReplicationLag", reflect.TypeOf((*MockThrottlerInterface)(nil).RecordReplicationLag), time, ts) } // ResetConfiguration mocks base method. @@ -102,53 +117,53 @@ func (mr *MockThrottlerInterfaceMockRecorder) ResetConfiguration() *gomock.Call } // SetMaxRate mocks base method. -func (m *MockThrottlerInterface) SetMaxRate(arg0 int64) { +func (m *MockThrottlerInterface) SetMaxRate(rate int64) { m.ctrl.T.Helper() - m.ctrl.Call(m, "SetMaxRate", arg0) + m.ctrl.Call(m, "SetMaxRate", rate) } // SetMaxRate indicates an expected call of SetMaxRate. -func (mr *MockThrottlerInterfaceMockRecorder) SetMaxRate(arg0 any) *gomock.Call { +func (mr *MockThrottlerInterfaceMockRecorder) SetMaxRate(rate interface{}) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetMaxRate", reflect.TypeOf((*MockThrottlerInterface)(nil).SetMaxRate), arg0) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetMaxRate", reflect.TypeOf((*MockThrottlerInterface)(nil).SetMaxRate), rate) } // ThreadFinished mocks base method. -func (m *MockThrottlerInterface) ThreadFinished(arg0 int) { +func (m *MockThrottlerInterface) ThreadFinished(threadID int) { m.ctrl.T.Helper() - m.ctrl.Call(m, "ThreadFinished", arg0) + m.ctrl.Call(m, "ThreadFinished", threadID) } // ThreadFinished indicates an expected call of ThreadFinished. -func (mr *MockThrottlerInterfaceMockRecorder) ThreadFinished(arg0 any) *gomock.Call { +func (mr *MockThrottlerInterfaceMockRecorder) ThreadFinished(threadID interface{}) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ThreadFinished", reflect.TypeOf((*MockThrottlerInterface)(nil).ThreadFinished), arg0) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ThreadFinished", reflect.TypeOf((*MockThrottlerInterface)(nil).ThreadFinished), threadID) } // Throttle mocks base method. -func (m *MockThrottlerInterface) Throttle(arg0 int) time.Duration { +func (m *MockThrottlerInterface) Throttle(threadID int) time.Duration { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "Throttle", arg0) + ret := m.ctrl.Call(m, "Throttle", threadID) ret0, _ := ret[0].(time.Duration) return ret0 } // Throttle indicates an expected call of Throttle. -func (mr *MockThrottlerInterfaceMockRecorder) Throttle(arg0 any) *gomock.Call { +func (mr *MockThrottlerInterfaceMockRecorder) Throttle(threadID interface{}) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Throttle", reflect.TypeOf((*MockThrottlerInterface)(nil).Throttle), arg0) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Throttle", reflect.TypeOf((*MockThrottlerInterface)(nil).Throttle), threadID) } // UpdateConfiguration mocks base method. -func (m *MockThrottlerInterface) UpdateConfiguration(arg0 *throttlerdata.Configuration, arg1 bool) error { +func (m *MockThrottlerInterface) UpdateConfiguration(configuration *throttlerdata.Configuration, copyZeroValues bool) error { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "UpdateConfiguration", arg0, arg1) + ret := m.ctrl.Call(m, "UpdateConfiguration", configuration, copyZeroValues) ret0, _ := ret[0].(error) return ret0 } // UpdateConfiguration indicates an expected call of UpdateConfiguration. -func (mr *MockThrottlerInterfaceMockRecorder) UpdateConfiguration(arg0, arg1 any) *gomock.Call { +func (mr *MockThrottlerInterfaceMockRecorder) UpdateConfiguration(configuration, copyZeroValues interface{}) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateConfiguration", reflect.TypeOf((*MockThrottlerInterface)(nil).UpdateConfiguration), arg0, arg1) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateConfiguration", reflect.TypeOf((*MockThrottlerInterface)(nil).UpdateConfiguration), configuration, copyZeroValues) } diff --git a/go/vt/vttablet/tabletserver/txthrottler/tx_throttler_test.go b/go/vt/vttablet/tabletserver/txthrottler/tx_throttler_test.go index 7b9e4bc98bd..65ef64614d1 100644 --- a/go/vt/vttablet/tabletserver/txthrottler/tx_throttler_test.go +++ b/go/vt/vttablet/tabletserver/txthrottler/tx_throttler_test.go @@ -89,27 +89,58 @@ func TestEnabledThrottler(t *testing.T) { return mockThrottler, nil } - call0 := mockThrottler.EXPECT().UpdateConfiguration(gomock.Any(), true /* copyZeroValues */) - call1 := mockThrottler.EXPECT().Throttle(0) - call1.Return(0 * time.Second) + var calls []*gomock.Call + + call := mockThrottler.EXPECT().UpdateConfiguration(gomock.Any(), true /* copyZeroValues */) + calls = append(calls, call) + + call = mockThrottler.EXPECT().LastMaxLagNotIgnoredForTabletType(topodatapb.TabletType_REPLICA) + call.Return(uint32(20)) + calls = append(calls, call) + + call = mockThrottler.EXPECT().Throttle(0) + call.Return(0 * time.Second) + calls = append(calls, call) + tabletStats := &discovery.LegacyTabletStats{ Target: &querypb.Target{ TabletType: topodatapb.TabletType_REPLICA, }, } - call2 := mockThrottler.EXPECT().RecordReplicationLag(gomock.Any(), tabletStats) - call3 := mockThrottler.EXPECT().Throttle(0) - call3.Return(1 * time.Second) - call4 := mockThrottler.EXPECT().Throttle(0) - call4.Return(1 * time.Second) - calllast := mockThrottler.EXPECT().Close() + call = mockThrottler.EXPECT().RecordReplicationLag(gomock.Any(), tabletStats) + calls = append(calls, call) + + call = mockThrottler.EXPECT().LastMaxLagNotIgnoredForTabletType(topodatapb.TabletType_REPLICA) + call.Return(uint32(20)) + calls = append(calls, call) + + call = mockThrottler.EXPECT().Throttle(0) + call.Return(1 * time.Second) + calls = append(calls, call) + + call = mockThrottler.EXPECT().LastMaxLagNotIgnoredForTabletType(topodatapb.TabletType_REPLICA) + call.Return(uint32(20)) + calls = append(calls, call) + + call = mockThrottler.EXPECT().Throttle(0) + call.Return(1 * time.Second) + calls = append(calls, call) - call1.After(call0) - call2.After(call1) - call3.After(call2) - call4.After(call3) - calllast.After(call4) + call = mockThrottler.EXPECT().LastMaxLagNotIgnoredForTabletType(topodatapb.TabletType_REPLICA) + call.Return(uint32(1)) + calls = append(calls, call) + + call = mockThrottler.EXPECT().Throttle(0) + call.Return(1 * time.Second) + calls = append(calls, call) + + call = mockThrottler.EXPECT().Close() + calls = append(calls, call) + + for i := 1; i < len(calls); i++ { + calls[i].After(calls[i-1]) + } config := tabletenv.NewDefaultConfig() config.EnableTxThrottler = true @@ -147,6 +178,12 @@ func TestEnabledThrottler(t *testing.T) { assert.False(t, throttler.Throttle(0, "some-workload")) assert.Equal(t, int64(3), throttler.requestsTotal.Counts()["some-workload"]) assert.Equal(t, int64(1), throttler.requestsThrottled.Counts()["some-workload"]) + + // This call should not throttle despite priority. Check that's the case and counters agree. + assert.False(t, throttler.Throttle(100, "some-workload")) + assert.Equal(t, int64(4), throttler.requestsTotal.Counts()["some-workload"]) + assert.Equal(t, int64(1), throttler.requestsThrottled.Counts()["some-workload"]) + throttler.Close() assert.Zero(t, throttler.throttlerRunning.Get()) }