From bbff9193d9c47fd9f5abd6f7cf5cda54cfcfe75e Mon Sep 17 00:00:00 2001 From: aoiasd <45024769+aoiasd@users.noreply.github.com> Date: Tue, 20 Feb 2024 14:16:51 +0800 Subject: [PATCH] enhance: support clean paramtable config event in test (#30534) relate: https://github.com/milvus-io/milvus/issues/30441 Signed-off-by: aoiasd --- internal/proxy/multi_rate_limiter_test.go | 4 ++++ pkg/config/event_dispatcher.go | 7 +++++++ pkg/util/paramtable/component_param.go | 11 +++++++++-- 3 files changed, 20 insertions(+), 2 deletions(-) diff --git a/internal/proxy/multi_rate_limiter_test.go b/internal/proxy/multi_rate_limiter_test.go index b8c4f53f16..489bf2a9a3 100644 --- a/internal/proxy/multi_rate_limiter_test.go +++ b/internal/proxy/multi_rate_limiter_test.go @@ -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( diff --git a/pkg/config/event_dispatcher.go b/pkg/config/event_dispatcher.go index b1697bb61b..197671272c 100644 --- a/pkg/config/event_dispatcher.go +++ b/pkg/config/event_dispatcher.go @@ -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) +} diff --git a/pkg/util/paramtable/component_param.go b/pkg/util/paramtable/component_param.go index 76f7eb8f43..7409c2e95c 100644 --- a/pkg/util/paramtable/component_param.go +++ b/pkg/util/paramtable/component_param.go @@ -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() } // /////////////////////////////////////////////////////////////////////////////