milvus-io
/
milvus
mirror of https://github.com/milvus-io/milvus.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

enhance: support clean paramtable config event in test (#30534) 

relate: https://github.com/milvus-io/milvus/issues/30441

Signed-off-by: aoiasd <zhicheng.yue@zilliz.com>
pull/30689/head
aoiasd 2024-02-20 14:16:51 +08:00 committed by GitHub
parent 77ba3ce3f3
commit bbff9193d9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 20 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
internal/proxy/multi_rate_limiter_test.go
View File

@ -211,6 +211,7 @@ func TestMultiRateLimiter(t *testing.T) {
func TestRateLimiter(t *testing.T) {
t.Run("test limit", func(t *testing.T) {
paramtable.Get().CleanEvent()
limiter := newRateLimiter(false)
for _, rt := range internalpb.RateType_value {
limiter.limiters.Insert(internalpb.RateType(rt), ratelimitutil.NewLimiter(ratelimitutil.Limit(1000), 1))
@ -226,6 +227,7 @@ func TestRateLimiter(t *testing.T) {
})
t.Run("test setRates", func(t *testing.T) {
paramtable.Get().CleanEvent()
limiter := newRateLimiter(false)
for _, rt := range internalpb.RateType_value {
limiter.limiters.Insert(internalpb.RateType(rt), ratelimitutil.NewLimiter(ratelimitutil.Limit(1000), 1))
@ -266,6 +268,7 @@ func TestRateLimiter(t *testing.T) {
})
t.Run("test get error code", func(t *testing.T) {
paramtable.Get().CleanEvent()
limiter := newRateLimiter(false)
for _, rt := range internalpb.RateType_value {
limiter.limiters.Insert(internalpb.RateType(rt), ratelimitutil.NewLimiter(ratelimitutil.Limit(1000), 1))
@ -295,6 +298,7 @@ func TestRateLimiter(t *testing.T) {
})
t.Run("tests refresh rate by config", func(t *testing.T) {
paramtable.Get().CleanEvent()
limiter := newRateLimiter(false)
etcdCli, _ := etcd.GetEtcdClient(

7
pkg/config/event_dispatcher.go
View File

@ -103,3 +103,10 @@ func (ed *EventDispatcher) Unregister(key string, handler EventHandler) {
}
ed.registry[key] = newGroup
}
func (ed *EventDispatcher) Clean() {
ed.mut.Lock()
defer ed.mut.Unlock()
ed.registry = make(map[string][]EventHandler)
}

11
pkg/util/paramtable/component_param.go
View File

@ -154,12 +154,19 @@ func (p *ComponentParam) Watch(key string, watcher config.EventHandler) {
p.baseTable.mgr.Dispatcher.Register(key, watcher)
}
func (p *ComponentParam) WatchKeyPrefix(keyPrefix string, watcher config.EventHandler) {
p.baseTable.mgr.Dispatcher.RegisterForKeyPrefix(keyPrefix, watcher)
}
func (p *ComponentParam) Unwatch(key string, watcher config.EventHandler) {
p.baseTable.mgr.Dispatcher.Unregister(key, watcher)
}
func (p *ComponentParam) WatchKeyPrefix(keyPrefix string, watcher config.EventHandler) {
p.baseTable.mgr.Dispatcher.RegisterForKeyPrefix(keyPrefix, watcher)
// FOR TEST
// clean all config event in dispatcher
func (p *ComponentParam) CleanEvent() {
p.baseTable.mgr.Dispatcher.Clean()
}
// /////////////////////////////////////////////////////////////////////////////