diff --git a/internal/config/env_source.go b/internal/config/env_source.go index c9b6af8fa1..6a5e0e2258 100644 --- a/internal/config/env_source.go +++ b/internal/config/env_source.go @@ -19,47 +19,49 @@ import ( "fmt" "os" "strings" - "sync" + + "github.com/milvus-io/milvus/internal/util/typeutil" ) type EnvSource struct { - configs sync.Map + configs *typeutil.ConcurrentMap[string, string] KeyFormatter func(string) string } func NewEnvSource(KeyFormatter func(string) string) EnvSource { es := EnvSource{ - configs: sync.Map{}, + configs: typeutil.NewConcurrentMap[string, string](), KeyFormatter: KeyFormatter, } + for _, value := range os.Environ() { rs := []rune(value) in := strings.Index(value, "=") key := string(rs[0:in]) value := string(rs[in+1:]) envKey := KeyFormatter(key) - es.configs.Store(key, value) - es.configs.Store(envKey, value) - + es.configs.Insert(key, value) + es.configs.Insert(envKey, value) } return es } // GetConfigurationByKey implements ConfigSource func (es EnvSource) GetConfigurationByKey(key string) (string, error) { - value, ok := es.configs.Load(key) + value, ok := es.configs.Get(key) + if !ok { return "", fmt.Errorf("key not found: %s", key) } - return value.(string), nil + return value, nil } // GetConfigurations implements ConfigSource func (es EnvSource) GetConfigurations() (map[string]string, error) { configMap := make(map[string]string) - es.configs.Range(func(k, v interface{}) bool { - configMap[k.(string)] = v.(string) + es.configs.Range(func(k, v string) bool { + configMap[k] = v return true }) diff --git a/internal/util/typeutil/map.go b/internal/util/typeutil/map.go index 41a839e222..d96aca47f0 100644 --- a/internal/util/typeutil/map.go +++ b/internal/util/typeutil/map.go @@ -98,3 +98,11 @@ func (m *ConcurrentMap[K, V]) GetAndRemove(key K) (V, bool) { m.len.Dec() return value.(V), true } + +func (m *ConcurrentMap[K, V]) Range(fn func(k K, v V) bool) { + m.inner.Range(func(k, v any) bool { + key := k.(K) + value := v.(V) + return fn(key, value) + }) +} diff --git a/internal/util/typeutil/map_test.go b/internal/util/typeutil/map_test.go index f13c6b251c..06fc82412e 100644 --- a/internal/util/typeutil/map_test.go +++ b/internal/util/typeutil/map_test.go @@ -3,7 +3,6 @@ package typeutil import ( "testing" - "github.com/stretchr/testify/assert" "github.com/stretchr/testify/suite" ) @@ -58,67 +57,80 @@ func (suite *MapUtilSuite) TestGetMapKeys() { func (suite *MapUtilSuite) TestConcurrentMap() { currMap := NewConcurrentMap[int64, string]() - assert.EqualValues(suite.T(), 0, currMap.Len()) + suite.EqualValues(0, currMap.Len()) v, loaded := currMap.GetOrInsert(100, "v-100") - assert.Equal(suite.T(), "v-100", v) - assert.Equal(suite.T(), false, loaded) + suite.Equal("v-100", v) + suite.Equal(false, loaded) v, loaded = currMap.GetOrInsert(100, "v-100") - assert.Equal(suite.T(), "v-100", v) - assert.Equal(suite.T(), true, loaded) + suite.Equal("v-100", v) + suite.Equal(true, loaded) v, loaded = currMap.GetOrInsert(100, "v-100") - assert.Equal(suite.T(), "v-100", v) - assert.Equal(suite.T(), true, loaded) - assert.Equal(suite.T(), uint64(1), currMap.Len()) - assert.EqualValues(suite.T(), 1, currMap.Len()) + suite.Equal("v-100", v) + suite.Equal(true, loaded) + suite.Equal(uint64(1), currMap.Len()) + suite.EqualValues(1, currMap.Len()) currMap.InsertIfNotPresent(100, "v-100-new") currMap.InsertIfNotPresent(200, "v-200") currMap.InsertIfNotPresent(300, "v-300") - assert.Equal(suite.T(), uint64(3), currMap.Len()) - assert.EqualValues(suite.T(), 3, currMap.Len()) + suite.Equal(uint64(3), currMap.Len()) + suite.EqualValues(3, currMap.Len()) var exist bool v, exist = currMap.Get(100) - assert.Equal(suite.T(), "v-100", v) - assert.Equal(suite.T(), true, exist) + suite.Equal("v-100", v) + suite.Equal(true, exist) v, exist = currMap.Get(200) - assert.Equal(suite.T(), "v-200", v) - assert.Equal(suite.T(), true, exist) + suite.Equal("v-200", v) + suite.Equal(true, exist) v, exist = currMap.Get(300) - assert.Equal(suite.T(), "v-300", v) - assert.Equal(suite.T(), true, exist) - assert.EqualValues(suite.T(), 3, currMap.Len()) + suite.Equal("v-300", v) + suite.Equal(true, exist) + suite.EqualValues(3, currMap.Len()) v, exist = currMap.GetOrInsert(100, "new-v") - assert.Equal(suite.T(), "v-100", v) - assert.Equal(suite.T(), true, exist) + suite.Equal("v-100", v) + suite.Equal(true, exist) v, exist = currMap.GetOrInsert(200, "new-v") - assert.Equal(suite.T(), "v-200", v) - assert.Equal(suite.T(), true, exist) + suite.Equal("v-200", v) + suite.Equal(true, exist) v, exist = currMap.GetOrInsert(300, "new-v") - assert.Equal(suite.T(), "v-300", v) - assert.Equal(suite.T(), true, exist) + suite.Equal("v-300", v) + suite.Equal(true, exist) v, exist = currMap.GetOrInsert(400, "new-v") - assert.Equal(suite.T(), "new-v", v) - assert.Equal(suite.T(), false, exist) - assert.EqualValues(suite.T(), 4, currMap.Len()) + suite.Equal("new-v", v) + suite.Equal(false, exist) + suite.EqualValues(4, currMap.Len()) + + currMap.Range(func(k int64, value string) bool { + suite.Contains([]int64{100, 200, 300, 400}, k) + expect, _ := currMap.Get(k) + suite.Equal(expect, value) + + return true + }) v, loaded = currMap.GetAndRemove(100) - assert.Equal(suite.T(), "v-100", v) - assert.Equal(suite.T(), true, loaded) + suite.Equal("v-100", v) + suite.Equal(true, loaded) v, loaded = currMap.GetAndRemove(200) - assert.Equal(suite.T(), "v-200", v) - assert.Equal(suite.T(), true, loaded) + suite.Equal("v-200", v) + suite.Equal(true, loaded) v, loaded = currMap.GetAndRemove(300) - assert.Equal(suite.T(), "v-300", v) - assert.Equal(suite.T(), true, loaded) + suite.Equal("v-300", v) + suite.Equal(true, loaded) v, loaded = currMap.GetAndRemove(400) - assert.Equal(suite.T(), "new-v", v) - assert.Equal(suite.T(), true, loaded) + suite.Equal("new-v", v) + suite.Equal(true, loaded) v, loaded = currMap.GetAndRemove(500) - assert.Equal(suite.T(), "", v) - assert.Equal(suite.T(), false, loaded) - assert.EqualValues(suite.T(), 0, currMap.Len()) + suite.Equal("", v) + suite.Equal(false, loaded) + suite.EqualValues(0, currMap.Len()) + + currMap.Range(func(k int64, value string) bool { + suite.FailNow("empty map range") + return false + }) } func TestMapUtil(t *testing.T) {