diff --git a/config/env/env.go b/config/env/env.go index 014501f5..a819e51a 100644 --- a/config/env/env.go +++ b/config/env/env.go @@ -18,31 +18,25 @@ import ( "fmt" "os" "strings" - "sync" + + "github.com/astaxie/beego/utils" ) -var env struct { - data map[string]string - lock *sync.RWMutex -} +var env *utils.BeeMap func init() { - env.data = make(map[string]string) - env.lock = &sync.RWMutex{} + env = utils.NewBeeMap() for _, e := range os.Environ() { splits := strings.Split(e, "=") - env.data[splits[0]] = os.Getenv(splits[0]) + env.Set(splits[0], os.Getenv(splits[0])) } } // Get returns a value by key. // If the key does not exist, the default value will be returned. func Get(key string, defVal string) string { - env.lock.RLock() - defer env.lock.RUnlock() - - if val, ok := env.data[key]; ok { - return val + if val := env.Get(key); val != nil { + return val.(string) } return defVal } @@ -50,11 +44,8 @@ func Get(key string, defVal string) string { // MustGet returns a value by key. // If the key does not exist, it will return an error. func MustGet(key string) (string, error) { - env.lock.RLock() - defer env.lock.RUnlock() - - if val, ok := env.data[key]; ok { - return val, nil + if val := env.Get(key); val != nil { + return val.(string), nil } return "", fmt.Errorf("no env variable with %s", key) } @@ -62,28 +53,33 @@ func MustGet(key string) (string, error) { // Set sets a value in the ENV copy. // This does not affect the child process environment. func Set(key string, value string) { - env.lock.Lock() - defer env.lock.Unlock() - env.data[key] = value + env.Set(key, value) } // MustSet sets a value in the ENV copy and the child process environment. // It returns an error in case the set operation failed. func MustSet(key string, value string) error { - env.lock.Lock() - defer env.lock.Unlock() - err := os.Setenv(key, value) if err != nil { return err } - env.data[key] = value + env.Set(key, value) return nil } // GetAll returns all keys/values in the current child process environment. func GetAll() map[string]string { - env.lock.RLock() - defer env.lock.RUnlock() - return env.data + items := env.Items() + envs := make(map[string]string, env.Count()) + + for key, val := range items { + switch key := key.(type) { + case string: + switch val := val.(type) { + case string: + envs[key] = val + } + } + } + return envs }