mirror of
https://gitee.com/milvus-io/milvus.git
synced 2024-11-30 10:59:32 +08:00
Fix copylocks for EnvSource (#22948)
Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
This commit is contained in:
parent
e8f8c1b445
commit
1b4e29e1ee
@ -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
|
||||
})
|
||||
|
||||
|
@ -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)
|
||||
})
|
||||
}
|
||||
|
@ -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) {
|
||||
|
Loading…
Reference in New Issue
Block a user