Fix copylocks for EnvSource (#22948)

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>

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.Insert(key, value)
-		es.configs.Store(envKey, 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 {
+	es.configs.Range(func(k, v string) bool {
-		configMap[k.(string)] = v.(string)
+		configMap[k] = v
 		return true
 	})
 

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)
+	})
+}

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)
+	suite.Equal("v-100", v)
-	assert.Equal(suite.T(), false, loaded)
+	suite.Equal(false, loaded)
 	v, loaded = currMap.GetOrInsert(100, "v-100")
-	assert.Equal(suite.T(), "v-100", v)
+	suite.Equal("v-100", v)
-	assert.Equal(suite.T(), true, loaded)
+	suite.Equal(true, loaded)
 	v, loaded = currMap.GetOrInsert(100, "v-100")
-	assert.Equal(suite.T(), "v-100", v)
+	suite.Equal("v-100", v)
-	assert.Equal(suite.T(), true, loaded)
+	suite.Equal(true, loaded)
-	assert.Equal(suite.T(), uint64(1), currMap.Len())
+	suite.Equal(uint64(1), currMap.Len())
-	assert.EqualValues(suite.T(), 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())
+	suite.Equal(uint64(3), currMap.Len())
-	assert.EqualValues(suite.T(), 3, currMap.Len())
+	suite.EqualValues(3, currMap.Len())
 
 	var exist bool
 	v, exist = currMap.Get(100)
-	assert.Equal(suite.T(), "v-100", v)
+	suite.Equal("v-100", v)
-	assert.Equal(suite.T(), true, exist)
+	suite.Equal(true, exist)
 	v, exist = currMap.Get(200)
-	assert.Equal(suite.T(), "v-200", v)
+	suite.Equal("v-200", v)
-	assert.Equal(suite.T(), true, exist)
+	suite.Equal(true, exist)
 	v, exist = currMap.Get(300)
-	assert.Equal(suite.T(), "v-300", v)
+	suite.Equal("v-300", v)
-	assert.Equal(suite.T(), true, exist)
+	suite.Equal(true, exist)
-	assert.EqualValues(suite.T(), 3, currMap.Len())
+	suite.EqualValues(3, currMap.Len())
 
 	v, exist = currMap.GetOrInsert(100, "new-v")
-	assert.Equal(suite.T(), "v-100", v)
+	suite.Equal("v-100", v)
-	assert.Equal(suite.T(), true, exist)
+	suite.Equal(true, exist)
 	v, exist = currMap.GetOrInsert(200, "new-v")
-	assert.Equal(suite.T(), "v-200", v)
+	suite.Equal("v-200", v)
-	assert.Equal(suite.T(), true, exist)
+	suite.Equal(true, exist)
 	v, exist = currMap.GetOrInsert(300, "new-v")
-	assert.Equal(suite.T(), "v-300", v)
+	suite.Equal("v-300", v)
-	assert.Equal(suite.T(), true, exist)
+	suite.Equal(true, exist)
 	v, exist = currMap.GetOrInsert(400, "new-v")
-	assert.Equal(suite.T(), "new-v", v)
+	suite.Equal("new-v", v)
-	assert.Equal(suite.T(), false, exist)
+	suite.Equal(false, exist)
-	assert.EqualValues(suite.T(), 4, currMap.Len())
+	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)
+	suite.Equal("v-100", v)
-	assert.Equal(suite.T(), true, loaded)
+	suite.Equal(true, loaded)
 	v, loaded = currMap.GetAndRemove(200)
-	assert.Equal(suite.T(), "v-200", v)
+	suite.Equal("v-200", v)
-	assert.Equal(suite.T(), true, loaded)
+	suite.Equal(true, loaded)
 	v, loaded = currMap.GetAndRemove(300)
-	assert.Equal(suite.T(), "v-300", v)
+	suite.Equal("v-300", v)
-	assert.Equal(suite.T(), true, loaded)
+	suite.Equal(true, loaded)
 	v, loaded = currMap.GetAndRemove(400)
-	assert.Equal(suite.T(), "new-v", v)
+	suite.Equal("new-v", v)
-	assert.Equal(suite.T(), true, loaded)
+	suite.Equal(true, loaded)
 	v, loaded = currMap.GetAndRemove(500)
-	assert.Equal(suite.T(), "", v)
+	suite.Equal("", v)
-	assert.Equal(suite.T(), false, loaded)
+	suite.Equal(false, loaded)
-	assert.EqualValues(suite.T(), 0, currMap.Len())
+	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) {