From 384383a7f0daf21cd2fe5d2ef34a65b6e66dcfe7 Mon Sep 17 00:00:00 2001
From: Jeremy Edwards <1312331+jeremyje@users.noreply.github.com>
Date: Thu, 25 Apr 2019 22:17:31 -0700
Subject: [PATCH] Synchronize read and writes to the Viper configuration.

---
 .travis.yml   |   2 +-
 go.sum        |   9 +++++
 util.go       |  59 ++++++++++++++++++++++++++++
 viper.go      | 104 ++++++++++++++++----------------------------------
 viper_test.go |  65 +++++++++++++++++++++++++++++++
 5 files changed, 167 insertions(+), 72 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index bb83057..108d486 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -22,7 +22,7 @@ matrix:
 script:
   - go install ./...
   - diff -u <(echo -n) <(gofmt -d .)
-  - go test -v ./...
+  - go test -v ./... -race -cover
 
 after_success:
   - go get -u -d github.com/spf13/hugo
diff --git a/go.sum b/go.sum
index 5c9fb7d..5b1ee32 100644
--- a/go.sum
+++ b/go.sum
@@ -1,7 +1,11 @@
+github.com/armon/consul-api v0.0.0-20180202201655-eb2c6b5be1b6 h1:G1bPvciwNyF7IUmKXNt9Ak3m6u9DE1rF+RmtIkBpVdA=
 github.com/armon/consul-api v0.0.0-20180202201655-eb2c6b5be1b6/go.mod h1:grANhF5doyWs3UAsr3K4I6qtAmlQcZDesFNEHPZAzj8=
+github.com/coreos/etcd v3.3.10+incompatible h1:jFneRYjIvLMLhDLCzuTuU4rSJUjRplcJQ7pD7MnhC04=
 github.com/coreos/etcd v3.3.10+incompatible/go.mod h1:uF7uidLiAD3TWHmW31ZFd/JWoc32PjwdhPthX9715RE=
 github.com/coreos/go-etcd v2.0.0+incompatible/go.mod h1:Jez6KQU2B/sWsbdaef3ED8NzMklzPG4d5KIOhIy30Tk=
+github.com/coreos/go-semver v0.2.0 h1:3Jm3tLmsgAYcjC+4Up7hJrFBPr+n7rAqYeSw/SZazuY=
 github.com/coreos/go-semver v0.2.0/go.mod h1:nnelYz7RCh+5ahJtPPxZlU+153eP4D4r3EedlOD2RNk=
+github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/fsnotify/fsnotify v1.4.7 h1:IXs+QLmnXW2CcXuY+8Mzv/fWEsPGWxqefPtCP5CnV9I=
 github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo=
@@ -13,6 +17,7 @@ github.com/mitchellh/mapstructure v1.1.2 h1:fmNYVwqnSfB9mZU6OS2O6GsXM+wcskZDuKQz
 github.com/mitchellh/mapstructure v1.1.2/go.mod h1:FVVH3fgwuzCH5S8UJGiWEs2h04kUh9fWfEaFds41c1Y=
 github.com/pelletier/go-toml v1.2.0 h1:T5zMGML61Wp+FlcbWjRDT7yAxhJNAiPPLOFECq181zc=
 github.com/pelletier/go-toml v1.2.0/go.mod h1:5z9KED0ma1S8pY6P1sdut58dfprrGBbd/94hg7ilaic=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/spf13/afero v1.1.2 h1:m8/z1t7/fwjysjQRYbP0RD+bUIF/8tJwPdEZsI83ACI=
 github.com/spf13/afero v1.1.2/go.mod h1:j4pytiNVoe2o6bmDsKpLACNPDBIoEAkihy7loJ1B0CQ=
@@ -22,9 +27,13 @@ github.com/spf13/jwalterweatherman v1.0.0 h1:XHEdyB+EcvlqZamSM4ZOMGlc93t6AcsBEu9
 github.com/spf13/jwalterweatherman v1.0.0/go.mod h1:cQK4TGJAtQXfYWX+Ddv3mKDzgVb68N+wFjFa4jdeBTo=
 github.com/spf13/pflag v1.0.3 h1:zPAT6CGy6wXeQ7NtTnaTerfKOsV6V6F8agHXFiazDkg=
 github.com/spf13/pflag v1.0.3/go.mod h1:DYY7MBk1bdzusC3SYhjObp+wFpr4gzcvqqNjLnInEg4=
+github.com/stretchr/testify v1.2.2 h1:bSDNvY7ZPG5RlJ8otE/7V6gMiyenm9RtJ7IUVIAoJ1w=
 github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
+github.com/ugorji/go/codec v0.0.0-20181204163529-d75b2dcb6bc8 h1:3SVOIvH7Ae1KRYyQWRjXWJEA9sS/c/pjvH++55Gr648=
 github.com/ugorji/go/codec v0.0.0-20181204163529-d75b2dcb6bc8/go.mod h1:VFNgLljTbGfSG7qAOspJ7OScBnGdDN/yBr0sguwnwf0=
+github.com/xordataexchange/crypt v0.0.3-0.20170626215501-b2862e3d0a77 h1:ESFSdwYZvkeru3RtdrYueztKhOBCSAAzS4Gf+k0tEow=
 github.com/xordataexchange/crypt v0.0.3-0.20170626215501-b2862e3d0a77/go.mod h1:aYKd//L2LvnjZzWKhF00oedf4jCCReLcmhLdhm1A27Q=
+golang.org/x/crypto v0.0.0-20181203042331-505ab145d0a9 h1:mKdxBk7AujPs8kU4m80U72y/zjbZ3UcXC7dClwKbUI0=
 golang.org/x/crypto v0.0.0-20181203042331-505ab145d0a9/go.mod h1:6SG95UA2DQfeDnfUPMdvaQW0Q7yPrPDi9nlGo2tz2b4=
 golang.org/x/sys v0.0.0-20181205085412-a5c9d58dba9a h1:1n5lsVfiQW3yfsRGu98756EH1YthsFqr/5mxHduZW2A=
 golang.org/x/sys v0.0.0-20181205085412-a5c9d58dba9a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
diff --git a/util.go b/util.go
index 952cad4..e0703e6 100644
--- a/util.go
+++ b/util.go
@@ -219,3 +219,62 @@ func deepSearch(m map[string]interface{}, path []string) map[string]interface{}
 	}
 	return m
 }
+
+// mergeFlatMap merges the given maps, excluding values of the second map
+// shadowed by values from the first map.
+func mergeFlatMap(shadow map[string]bool, m map[string]interface{}, keyDelim string) map[string]bool {
+	// scan keys
+outer:
+	for k, _ := range m {
+		path := strings.Split(k, keyDelim)
+		// scan intermediate paths
+		var parentKey string
+		for i := 1; i < len(path); i++ {
+			parentKey = strings.Join(path[0:i], v.keyDelim)
+			if shadow[parentKey] {
+				// path is shadowed, continue
+				continue outer
+			}
+		}
+		// add key
+		shadow[strings.ToLower(k)] = true
+	}
+	return shadow
+}
+
+// flattenAndMergeMap recursively flattens the given map into a map[string]bool
+// of key paths (used as a set, easier to manipulate than a []string):
+// - each path is merged into a single key string, delimited with v.keyDelim (= ".")
+// - if a path is shadowed by an earlier value in the initial shadow map,
+//   it is skipped.
+// The resulting set of paths is merged to the given shadow set at the same time.
+func flattenAndMergeMap(shadow map[string]bool, m map[string]interface{}, prefix string, keyDelim string) map[string]bool {
+	if shadow != nil && prefix != "" && shadow[prefix] {
+		// prefix is shadowed => nothing more to flatten
+		return shadow
+	}
+	if shadow == nil {
+		shadow = make(map[string]bool)
+	}
+
+	var m2 map[string]interface{}
+	if prefix != "" {
+		prefix += keyDelim
+	}
+	for k, val := range m {
+		fullKey := prefix + k
+		switch val.(type) {
+		case map[string]interface{}:
+			m2 = val.(map[string]interface{})
+		case map[interface{}]interface{}:
+			m2 = cast.ToStringMap(val)
+		default:
+			// immediate value
+			shadow[strings.ToLower(fullKey)] = true
+			continue
+		}
+		// recursively merge to shadow map
+		shadow = flattenAndMergeMap(shadow, m2, fullKey, keyDelim)
+	}
+	return shadow
+}
diff --git a/viper.go b/viper.go
index a3d37f8..6c4d666 100644
--- a/viper.go
+++ b/viper.go
@@ -204,6 +204,7 @@ type Viper struct {
 	properties *properties.Properties
 
 	onConfigChange func(fsnotify.Event)
+	configMutex    sync.RWMutex
 }
 
 // New returns an initialized Viper instance.
@@ -972,7 +973,8 @@ func (v *Viper) BindEnv(input ...string) error {
 // Viper will check to see if an alias exists first.
 // Note: this assumes a lower-cased key given.
 func (v *Viper) find(lcaseKey string) interface{} {
-
+	v.configMutex.RLock()
+	defer v.configMutex.RUnlock()
 	var (
 		val    interface{}
 		exists bool
@@ -1212,6 +1214,8 @@ func (v *Viper) Set(key string, value interface{}) {
 	lastKey := strings.ToLower(path[len(path)-1])
 	deepestMap := deepSearch(v.override, path[0:len(path)-1])
 
+	v.configMutex.Lock()
+	defer v.configMutex.Unlock()
 	// set innermost value
 	deepestMap[lastKey] = value
 }
@@ -1243,6 +1247,8 @@ func (v *Viper) ReadInConfig() error {
 		return err
 	}
 
+	v.configMutex.Lock()
+	defer v.configMutex.Unlock()
 	v.config = config
 	return nil
 }
@@ -1272,8 +1278,12 @@ func (v *Viper) MergeInConfig() error {
 // key does not exist in the file.
 func ReadConfig(in io.Reader) error { return v.ReadConfig(in) }
 func (v *Viper) ReadConfig(in io.Reader) error {
-	v.config = make(map[string]interface{})
-	return v.unmarshalReader(in, v.config)
+	config := make(map[string]interface{})
+	err := v.unmarshalReader(in, config)
+	v.configMutex.Lock()
+	defer v.configMutex.Unlock()
+	v.config = config
+	return err
 }
 
 // MergeConfig merges a new configuration with an existing config.
@@ -1290,11 +1300,18 @@ func (v *Viper) MergeConfig(in io.Reader) error {
 // Note that the map given may be modified.
 func MergeConfigMap(cfg map[string]interface{}) error { return v.MergeConfigMap(cfg) }
 func (v *Viper) MergeConfigMap(cfg map[string]interface{}) error {
-	if v.config == nil {
-		v.config = make(map[string]interface{})
+	v.configMutex.RLock()
+	config := v.config
+	v.configMutex.RUnlock()
+	if config == nil {
+		config = make(map[string]interface{})
 	}
 	insensitiviseMap(cfg)
 	mergeMaps(cfg, v.config, nil)
+
+	v.configMutex.Lock()
+	defer v.configMutex.Unlock()
+	v.config = config
 	return nil
 }
 
@@ -1367,6 +1384,8 @@ func unmarshalReader(in io.Reader, c map[string]interface{}) error {
 	return v.unmarshalReader(in, c)
 }
 func (v *Viper) unmarshalReader(in io.Reader, c map[string]interface{}) error {
+	v.configMutex.RLock()
+	defer v.configMutex.RUnlock()
 	buf := new(bytes.Buffer)
 	buf.ReadFrom(in)
 
@@ -1668,15 +1687,17 @@ func (v *Viper) watchRemoteConfig(provider RemoteProvider) (map[string]interface
 // Nested keys are returned with a v.keyDelim (= ".") separator
 func AllKeys() []string { return v.AllKeys() }
 func (v *Viper) AllKeys() []string {
+	v.configMutex.RLock()
+	defer v.configMutex.RUnlock()
 	m := map[string]bool{}
 	// add all paths, by order of descending priority to ensure correct shadowing
-	m = v.flattenAndMergeMap(m, castMapStringToMapInterface(v.aliases), "")
-	m = v.flattenAndMergeMap(m, v.override, "")
-	m = v.mergeFlatMap(m, castMapFlagToMapInterface(v.pflags))
-	m = v.mergeFlatMap(m, castMapStringToMapInterface(v.env))
-	m = v.flattenAndMergeMap(m, v.config, "")
-	m = v.flattenAndMergeMap(m, v.kvstore, "")
-	m = v.flattenAndMergeMap(m, v.defaults, "")
+	m = flattenAndMergeMap(m, castMapStringToMapInterface(v.aliases), "", v.keyDelim)
+	m = flattenAndMergeMap(m, v.override, "", v.keyDelim)
+	m = mergeFlatMap(m, castMapFlagToMapInterface(v.pflags), v.keyDelim)
+	m = mergeFlatMap(m, castMapStringToMapInterface(v.env), v.keyDelim)
+	m = flattenAndMergeMap(m, v.config, "", v.keyDelim)
+	m = flattenAndMergeMap(m, v.kvstore, "", v.keyDelim)
+	m = flattenAndMergeMap(m, v.defaults, "", v.keyDelim)
 
 	// convert set of paths to list
 	a := []string{}
@@ -1686,65 +1707,6 @@ func (v *Viper) AllKeys() []string {
 	return a
 }
 
-// flattenAndMergeMap recursively flattens the given map into a map[string]bool
-// of key paths (used as a set, easier to manipulate than a []string):
-// - each path is merged into a single key string, delimited with v.keyDelim (= ".")
-// - if a path is shadowed by an earlier value in the initial shadow map,
-//   it is skipped.
-// The resulting set of paths is merged to the given shadow set at the same time.
-func (v *Viper) flattenAndMergeMap(shadow map[string]bool, m map[string]interface{}, prefix string) map[string]bool {
-	if shadow != nil && prefix != "" && shadow[prefix] {
-		// prefix is shadowed => nothing more to flatten
-		return shadow
-	}
-	if shadow == nil {
-		shadow = make(map[string]bool)
-	}
-
-	var m2 map[string]interface{}
-	if prefix != "" {
-		prefix += v.keyDelim
-	}
-	for k, val := range m {
-		fullKey := prefix + k
-		switch val.(type) {
-		case map[string]interface{}:
-			m2 = val.(map[string]interface{})
-		case map[interface{}]interface{}:
-			m2 = cast.ToStringMap(val)
-		default:
-			// immediate value
-			shadow[strings.ToLower(fullKey)] = true
-			continue
-		}
-		// recursively merge to shadow map
-		shadow = v.flattenAndMergeMap(shadow, m2, fullKey)
-	}
-	return shadow
-}
-
-// mergeFlatMap merges the given maps, excluding values of the second map
-// shadowed by values from the first map.
-func (v *Viper) mergeFlatMap(shadow map[string]bool, m map[string]interface{}) map[string]bool {
-	// scan keys
-outer:
-	for k, _ := range m {
-		path := strings.Split(k, v.keyDelim)
-		// scan intermediate paths
-		var parentKey string
-		for i := 1; i < len(path); i++ {
-			parentKey = strings.Join(path[0:i], v.keyDelim)
-			if shadow[parentKey] {
-				// path is shadowed, continue
-				continue outer
-			}
-		}
-		// add key
-		shadow[strings.ToLower(k)] = true
-	}
-	return shadow
-}
-
 // AllSettings merges all settings and returns them as a map[string]interface{}.
 func AllSettings() map[string]interface{} { return v.AllSettings() }
 func (v *Viper) AllSettings() map[string]interface{} {
diff --git a/viper_test.go b/viper_test.go
index f43a38b..0b1b38b 100644
--- a/viper_test.go
+++ b/viper_test.go
@@ -1609,6 +1609,71 @@ func TestWatchFile(t *testing.T) {
 
 }
 
+func TestHighMutationRateAcrossGoRoutines(t *testing.T) {
+	var wg sync.WaitGroup
+	end := 25
+	v := New()
+	// given a `config.yaml` file being watched
+	v, configFile, cleanup := newViperWithConfigFile(t)
+	defer cleanup()
+	_, err := os.Stat(configFile)
+	require.NoError(t, err)
+	t.Logf("test config file: %s\n", configFile)
+
+	v.WatchConfig()
+
+	for i := 0; i < end; i++ {
+		// Set and Get
+		wg.Add(1)
+		go func(i int) {
+			defer wg.Done()
+			v.Set("key", i)
+			val := v.GetInt("key")
+			if !v.IsSet("key") {
+				t.Errorf("'key' should be set.")
+			}
+			if !(0 <= val && val < end) {
+				t.Errorf("Expected range 0 <= %d <= %d", val, end)
+			}
+		}(i)
+
+		// Trigger WatchConfig reload
+		wg.Add(1)
+		go func(i int) {
+			defer wg.Done()
+			if err := ioutil.WriteFile(configFile, []byte(fmt.Sprintf("foo: %d\n", i)), 0640); err != nil {
+				t.Fatal(err)
+			}
+		}(i)
+		// ReadConfig
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			if err := v.ReadConfig(bytes.NewBuffer(yamlMergeExampleTgt)); err != nil {
+				t.Fatal(err)
+			}
+		}()
+
+		// MergeConfig
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			if err := v.MergeConfig(bytes.NewBuffer(yamlMergeExampleSrc)); err != nil {
+				t.Fatal(err)
+			}
+		}()
+
+		// AllSettings and AllKeys
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			v.AllSettings()
+			v.AllKeys()
+		}()
+	}
+	wg.Wait()
+}
+
 func BenchmarkGetBool(b *testing.B) {
 	key := "BenchmarkGetBool"
 	v = New()