additions to make viper goroutine save for some actions

This commit is contained in:
Andrey Odintsov 2018-07-27 13:37:05 +03:00
parent e919dff7ac
commit afca17a005

View file

@ -30,21 +30,22 @@ import (
"path/filepath" "path/filepath"
"reflect" "reflect"
"strings" "strings"
"sync"
"time" "time"
yaml "gopkg.in/yaml.v2" "github.com/ghodss/yaml"
"github.com/fsnotify/fsnotify"
"github.com/hashicorp/hcl" "github.com/hashicorp/hcl"
"github.com/hashicorp/hcl/hcl/printer" "github.com/hashicorp/hcl/hcl/printer"
"github.com/magiconair/properties" "github.com/magiconair/properties"
"github.com/mitchellh/mapstructure" "github.com/mitchellh/mapstructure"
toml "github.com/pelletier/go-toml" "github.com/nwingert/fsnotify"
"github.com/pelletier/go-toml"
"github.com/spf13/afero" "github.com/spf13/afero"
"github.com/spf13/cast" "github.com/spf13/cast"
jww "github.com/spf13/jwalterweatherman" jww "github.com/spf13/jwalterweatherman"
"github.com/spf13/pflag" "github.com/spf13/pflag"
) )
// ConfigMarshalError happens when failing to marshal the configuration. // ConfigMarshalError happens when failing to marshal the configuration.
type ConfigMarshalError struct { type ConfigMarshalError struct {
@ -184,6 +185,11 @@ type Viper struct {
properties *properties.Properties properties *properties.Properties
onConfigChange func(fsnotify.Event) onConfigChange func(fsnotify.Event)
wChan chan bool
watched bool
sync.Mutex
} }
// New returns an initialized Viper instance. // New returns an initialized Viper instance.
@ -277,6 +283,12 @@ func (v *Viper) WatchConfig() {
configFile := filepath.Clean(filename) configFile := filepath.Clean(filename)
configDir, _ := filepath.Split(configFile) configDir, _ := filepath.Split(configFile)
wChan := make(chan bool)
v.Lock()
v.watched = true
v.Unlock()
done := make(chan bool) done := make(chan bool)
go func() { go func() {
for { for {
@ -296,6 +308,8 @@ func (v *Viper) WatchConfig() {
} }
case err := <-watcher.Errors: case err := <-watcher.Errors:
log.Println("error:", err) log.Println("error:", err)
case <-wChan:
done <- true
} }
} }
}() }()
@ -305,12 +319,24 @@ func (v *Viper) WatchConfig() {
}() }()
} }
func UnwatchConfig() { v.UnwatchConfig() }
func (v *Viper) UnwatchConfig() {
v.Lock()
defer v.Unlock()
v.watched = false
}
// SetConfigFile explicitly defines the path, name and extension of the config file. // SetConfigFile explicitly defines the path, name and extension of the config file.
// Viper will use this and not check any of the config paths. // Viper will use this and not check any of the config paths.
func SetConfigFile(in string) { v.SetConfigFile(in) } func SetConfigFile(in string) { v.SetConfigFile(in) }
func (v *Viper) SetConfigFile(in string) { func (v *Viper) SetConfigFile(in string) {
if in != "" { if in != "" {
v.configFile = in v.configFile = in
if v.watched {
// We need to recreate watcher in order to utilize hung goroutines
v.UnwatchConfig()
v.WatchConfig()
}
} }
} }
@ -890,6 +916,8 @@ func (v *Viper) BindEnv(input ...string) error {
// Viper will check to see if an alias exists first. // Viper will check to see if an alias exists first.
// Note: this assumes a lower-cased key given. // Note: this assumes a lower-cased key given.
func (v *Viper) find(lcaseKey string) interface{} { func (v *Viper) find(lcaseKey string) interface{} {
v.Lock()
defer v.Unlock()
var ( var (
val interface{} val interface{}
@ -1160,8 +1188,9 @@ func (v *Viper) ReadInConfig() error {
if err != nil { if err != nil {
return err return err
} }
v.Lock()
v.config = config v.config = config
v.Unlock()
return nil return nil
} }
@ -1585,6 +1614,9 @@ func (v *Viper) watchRemoteConfig(provider RemoteProvider) (map[string]interface
// Nested keys are returned with a v.keyDelim (= ".") separator // Nested keys are returned with a v.keyDelim (= ".") separator
func AllKeys() []string { return v.AllKeys() } func AllKeys() []string { return v.AllKeys() }
func (v *Viper) AllKeys() []string { func (v *Viper) AllKeys() []string {
v.Lock()
defer v.Unlock()
m := map[string]bool{} m := map[string]bool{}
// add all paths, by order of descending priority to ensure correct shadowing // add all paths, by order of descending priority to ensure correct shadowing
m = v.flattenAndMergeMap(m, castMapStringToMapInterface(v.aliases), "") m = v.flattenAndMergeMap(m, castMapStringToMapInterface(v.aliases), "")
@ -1594,7 +1626,6 @@ func (v *Viper) AllKeys() []string {
m = v.flattenAndMergeMap(m, v.config, "") m = v.flattenAndMergeMap(m, v.config, "")
m = v.flattenAndMergeMap(m, v.kvstore, "") m = v.flattenAndMergeMap(m, v.kvstore, "")
m = v.flattenAndMergeMap(m, v.defaults, "") m = v.flattenAndMergeMap(m, v.defaults, "")
// convert set of paths to list // convert set of paths to list
a := []string{} a := []string{}
for x := range m { for x := range m {
@ -1637,6 +1668,7 @@ func (v *Viper) flattenAndMergeMap(shadow map[string]bool, m map[string]interfac
// recursively merge to shadow map // recursively merge to shadow map
shadow = v.flattenAndMergeMap(shadow, m2, fullKey) shadow = v.flattenAndMergeMap(shadow, m2, fullKey)
} }
return shadow return shadow
} }
@ -1669,6 +1701,11 @@ func (v *Viper) AllSettings() map[string]interface{} {
// start from the list of keys, and construct the map one value at a time // start from the list of keys, and construct the map one value at a time
for _, k := range v.AllKeys() { for _, k := range v.AllKeys() {
value := v.Get(k) value := v.Get(k)
// ensure that marshalling to json will be supported
switch v := value.(type) {
case map[interface{}]interface{}:
value = castToMapStringInterface(v)
}
if value == nil { if value == nil {
// should not happen, since AllKeys() returns only keys holding a value, // should not happen, since AllKeys() returns only keys holding a value,
// check just in case anything changes // check just in case anything changes