Merge 3424c8c181 into 727a41c38a

README.md

@@ -176,11 +176,15 @@ Optionally you can provide a function for Viper to run each time a change occurs
 
 ```go
 viper.WatchConfig()
+defer CancelWatchConfig()
 viper.OnConfigChange(func(e fsnotify.Event) {
 	fmt.Println("Config file changed:", e.Name)
 })
 ```
 
+If you wish to stop watching the configPaths, simply call viper.CancelWatchConfig().
+Note: This might be necessary if your tests involve trying out various config files.
+
 ### Reading Config from io.Reader
 
 Viper predefines many configuration sources such as files, environment

viper.go

@@ -216,7 +216,8 @@ type Viper struct {
 	// This will only be used if the configuration read is a properties file.
 	properties *properties.Properties
 
-	onConfigChange func(fsnotify.Event)
+	onConfigChange    func(fsnotify.Event)
+	cancelWatchConfig func()
 }
 
 // New returns an initialized Viper instance.
@@ -333,80 +334,87 @@ var SupportedExts = []string{"json", "toml", "yaml", "yml", "properties", "props
 // SupportedRemoteProviders are universally supported remote providers.
 var SupportedRemoteProviders = []string{"etcd", "consul", "firestore"}
 
+// OnConfigChange is used to set the handler to run for catching the event
+// when the configFile have been externally modified by the filesystem.
 func OnConfigChange(run func(in fsnotify.Event)) { v.OnConfigChange(run) }
 func (v *Viper) OnConfigChange(run func(in fsnotify.Event)) {
 	v.onConfigChange = run
 }
 
+// CancelWatchConfig finishes the current configFile watch if it runs.
+// When this function returns the configFile changes are not more looked after.
+func CancelWatchConfig() { v.cancelWatchConfig() }
+func (v *Viper) CancelWatchConfig() {
+	if v.cancelWatchConfig != nil {
+		v.cancelWatchConfig()
+		v.cancelWatchConfig = nil
+	}
+}
+
+// WatchConfig watches and notifies changes on the file configFile.
 func WatchConfig() { v.WatchConfig() }
-
 func (v *Viper) WatchConfig() {
-	initWG := sync.WaitGroup{}
-	initWG.Add(1)
+	// Make sure not to run twice this routine
+	v.CancelWatchConfig()
+
+	// we have to watch the entire directory to pick up renames/atomic saves in a cross-platform way
+	filename, err := v.getConfigFile()
+	if err != nil {
+		log.Printf("error: %v\n", err)
+		return
+	}
+
+	configFile := filepath.Clean(filename)
+	configDir, _ := filepath.Split(configFile)
+	realConfigFile, _ := filepath.EvalSymlinks(filename)
+
+	watcher, err := fsnotify.NewWatcher()
+	if err != nil {
+		log.Fatal(err)
+	}
+	watcher.Add(configDir)
+
+	// cancellation and waiting point
+	var watcherGroup sync.WaitGroup
+	v.cancelWatchConfig = func() {
+		watcher.Close()
+		watcherGroup.Wait()
+	}
+	// Process watcher events
+	watcherGroup.Add(1)
 	go func() {
-		watcher, err := fsnotify.NewWatcher()
-		if err != nil {
-			log.Fatal(err)
-		}
-		defer watcher.Close()
-		// we have to watch the entire directory to pick up renames/atomic saves in a cross-platform way
-		filename, err := v.getConfigFile()
-		if err != nil {
-			log.Printf("error: %v\n", err)
-			initWG.Done()
-			return
-		}
-
-		configFile := filepath.Clean(filename)
-		configDir, _ := filepath.Split(configFile)
-		realConfigFile, _ := filepath.EvalSymlinks(filename)
-
-		eventsWG := sync.WaitGroup{}
-		eventsWG.Add(1)
-		go func() {
-			for {
-				select {
-				case event, ok := <-watcher.Events:
-					if !ok { // 'Events' channel is closed
-						eventsWG.Done()
-						return
-					}
-					currentConfigFile, _ := filepath.EvalSymlinks(filename)
-					// we only care about the config file with the following cases:
-					// 1 - if the config file was modified or created
-					// 2 - if the real path to the config file changed (eg: k8s ConfigMap replacement)
-					const writeOrCreateMask = fsnotify.Write | fsnotify.Create
-					if (filepath.Clean(event.Name) == configFile &&
-						event.Op&writeOrCreateMask != 0) ||
-						(currentConfigFile != "" && currentConfigFile != realConfigFile) {
-						realConfigFile = currentConfigFile
-						err := v.ReadInConfig()
-						if err != nil {
-							log.Printf("error reading config file: %v\n", err)
-						}
-						if v.onConfigChange != nil {
-							v.onConfigChange(event)
-						}
-					} else if filepath.Clean(event.Name) == configFile &&
-						event.Op&fsnotify.Remove&fsnotify.Remove != 0 {
-						eventsWG.Done()
-						return
-					}
-
-				case err, ok := <-watcher.Errors:
-					if ok { // 'Errors' channel is not closed
-						log.Printf("watcher error: %v\n", err)
-					}
-					eventsWG.Done()
-					return
+		defer watcherGroup.Done()
+		for event := range watcher.Events {
+			currentConfigFile, _ := filepath.EvalSymlinks(filename)
+			// we only care about the config file with the following cases:
+			// 1 - if the config file was modified or created
+			// 2 - if the real path to the config file changed (eg: k8s ConfigMap replacement)
+			const writeOrCreateMask = fsnotify.Write | fsnotify.Create
+			if (filepath.Clean(event.Name) == configFile &&
+				event.Op&writeOrCreateMask != 0) ||
+				(currentConfigFile != "" && currentConfigFile != realConfigFile) {
+				realConfigFile = currentConfigFile
+				err := v.ReadInConfig()
+				if err != nil {
+					log.Printf("error reading config file: %v\n", err)
 				}
+				if v.onConfigChange != nil {
+					v.onConfigChange(event)
+				}
+			} else if filepath.Clean(event.Name) == configFile &&
+				event.Op&fsnotify.Remove&fsnotify.Remove != 0 {
+				return
 			}
-		}()
-		watcher.Add(configDir)
-		initWG.Done()   // done initializing the watch in this go routine, so the parent routine can move on...
-		eventsWG.Wait() // now, wait for event loop to end in this go-routine...
+		}
+	}()
+	// Process watcher errors
+	watcherGroup.Add(1)
+	go func() {
+		defer watcherGroup.Done()
+		for err := range watcher.Errors {
+			log.Printf("watcher error: %v\n", err)
+		}
 	}()
-	initWG.Wait() // make sure that the go routine above fully ended before returning
 }
 
 // SetConfigFile explicitly defines the path, name and extension of the config file.
@@ -1499,6 +1507,10 @@ func (v *Viper) MergeConfigMap(cfg map[string]interface{}) error {
 func WriteConfig() error { return v.WriteConfig() }
 
 func (v *Viper) WriteConfig() error {
+	if v.cancelWatchConfig != nil {
+		v.cancelWatchConfig()
+		defer v.WatchConfig()
+	}
 	filename, err := v.getConfigFile()
 	if err != nil {
 		return err

viper_test.go

@@ -18,7 +18,6 @@ import (
 	"runtime"
 	"sort"
 	"strings"
-	"sync"
 	"testing"
 	"time"
 
@@ -2148,17 +2147,15 @@ func TestWatchFile(t *testing.T) {
 		defer cleanup()
 		_, err := os.Stat(configFile)
 		require.NoError(t, err)
-		t.Logf("test config file: %s\n", configFile)
-		wg := sync.WaitGroup{}
-		wg.Add(1)
+		done := make(chan struct{})
 		v.OnConfigChange(func(in fsnotify.Event) {
 			t.Logf("config file changed")
-			wg.Done()
+			close(done)
 		})
 		v.WatchConfig()
 		// when overwriting the file and waiting for the custom change notification handler to be triggered
 		err = ioutil.WriteFile(configFile, []byte("foo: baz\n"), 0640)
-		wg.Wait()
+		<-done
 		// then the config value should have changed
 		require.Nil(t, err)
 		assert.Equal(t, "baz", v.Get("foo"))
@@ -2171,13 +2168,12 @@ func TestWatchFile(t *testing.T) {
 		}
 		v, watchDir, _, _ := newViperWithSymlinkedConfigFile(t)
 		// defer cleanup()
-		wg := sync.WaitGroup{}
 		v.WatchConfig()
+		done := make(chan struct{})
 		v.OnConfigChange(func(in fsnotify.Event) {
 			t.Logf("config file changed")
-			wg.Done()
+			close(done)
 		})
-		wg.Add(1)
 		// when link to another `config.yaml` file
 		dataDir2 := path.Join(watchDir, "data2")
 		err := os.Mkdir(dataDir2, 0777)
@@ -2188,11 +2184,123 @@ func TestWatchFile(t *testing.T) {
 		// change the symlink using the `ln -sfn` command
 		err = exec.Command("ln", "-sfn", dataDir2, path.Join(watchDir, "data")).Run()
 		require.Nil(t, err)
-		wg.Wait()
+		<-done
 		// then
 		require.Nil(t, err)
 		assert.Equal(t, "baz", v.Get("foo"))
 	})
+
+	t.Run("file content changed after cancel", func(t *testing.T) {
+		// given a `config.yaml` file being watched
+		v, configFile, cleanup := newViperWithConfigFile(t)
+		defer cleanup()
+		_, err := os.Stat(configFile)
+		require.NoError(t, err)
+
+		// first run through with watching enabled
+		done := make(chan struct{})
+		v.OnConfigChange(func(in fsnotify.Event) {
+			done <- struct{}{}
+		})
+		v.WatchConfig()
+		// when overwriting the file and waiting for the custom change notification handler to be triggered
+		err = ioutil.WriteFile(configFile, []byte("foo: baz\n"), 0640)
+		<-done
+		// then the config value should have changed
+		require.Nil(t, err)
+		assert.Equal(t, "baz", v.Get("foo"))
+
+		// cancel and wait for the canceling to finish.
+		v.OnConfigChange(func(in fsnotify.Event) {
+			t.Error("CancelWatchConfig did not prevent second change from being seen.")
+		})
+		v.CancelWatchConfig()
+
+		// use another viper as a signal to wait this invisible write
+		v2 := New()
+		v2.SetConfigFile(configFile)
+		v2.WatchConfig()
+		v2.OnConfigChange(func(in fsnotify.Event) {
+			close(done)
+		})
+		err = ioutil.WriteFile(configFile, []byte("foo: quz\n"), 0640)
+		<-done
+		// the config value should still be the same.
+		require.Nil(t, err)
+		assert.Equal(t, "baz", v.Get("foo"), "CancelWatchConfig did not prevent second change from being seen.")
+	})
+
+	t.Run("do not watchConfig during writeConfig", func(t *testing.T) {
+		// given a `config.yaml` file being watched
+		v, configFile, cleanup := newViperWithConfigFile(t)
+		defer cleanup()
+		_, err := os.Stat(configFile)
+		require.NoError(t, err)
+
+		// first run through with watching enabled
+		done := make(chan struct{})
+		v.OnConfigChange(func(in fsnotify.Event) {
+			close(done)
+		})
+		v.WatchConfig()
+		// when overwriting the file and waiting for the custom change notification handler to be triggered
+		err = ioutil.WriteFile(configFile, []byte("foo: baz\nbar: foo\n"), 0640)
+		<-done
+		// then the config value should have changed
+		require.Nil(t, err)
+		assert.Equal(t, "baz", v.Get("foo"))
+
+		v.Set("foo", "bar")
+		v.WriteConfig()
+		v.Set("baz", "foo")
+		v.WriteConfig()
+		v.Set("aa", "bb")
+		v.WriteConfig()
+		v.Set("aaa", "10s")
+		v.WriteConfig()
+
+		f, err := ioutil.ReadFile(configFile)
+		assert.NoError(t, err)
+		assert.Equal(t, string([]byte("aa: bb\naaa: 10s\nbar: foo\nbaz: foo\nfoo: bar\n")), string(f))
+	})
+	t.Run("still watchConfig after writeConfig", func(t *testing.T) {
+		// given a `config.yaml` file being watched
+		v, configFile, cleanup := newViperWithConfigFile(t)
+		defer cleanup()
+		_, err := os.Stat(configFile)
+		require.NoError(t, err)
+
+		// first run through with watching enabled
+		done := make(chan struct{})
+		v.OnConfigChange(func(in fsnotify.Event) {
+			done <- struct{}{}
+		})
+		v.WatchConfig()
+		// when overwriting the file and waiting for the custom change notification handler to be triggered
+		err = ioutil.WriteFile(configFile, []byte("foo: baz\nbar: foo\n"), 0640)
+		<-done
+		// then the config value should have changed
+		require.Nil(t, err)
+		assert.Equal(t, "baz", v.Get("foo"))
+
+		v.Set("foo", "bar")
+		v.Set("baz", "foo")
+		v.Set("aa", "bb")
+		v.Set("aaa", "10s")
+		v.WriteConfig()
+
+		f, err := ioutil.ReadFile(configFile)
+		assert.NoError(t, err)
+		assert.Equal(t, string([]byte("aa: bb\naaa: 10s\nbar: foo\nbaz: foo\nfoo: bar\n")), string(f))
+		// Check still watching after write
+		v.OnConfigChange(func(in fsnotify.Event) {
+			close(done)
+		})
+		err = ioutil.WriteFile(configFile, []byte("foo: bar\n"), 0640)
+		require.Nil(t, err)
+		<-done // then the config value should have changed
+		assert.Equal(t, "bar", v.Get("foo"))
+	})
 }
 
 func TestUnmarshal_DotSeparatorBackwardCompatibility(t *testing.T) {