use BeeMap instead of a regular map

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
 }