From 59e479f6b837a7f749a14eea6998ecabb65f54fc Mon Sep 17 00:00:00 2001 From: chowyi Date: Mon, 31 Jul 2023 19:47:29 +0800 Subject: [PATCH] feat: add a new API StopWatching() to stop the groutine created by WatchConfig() manually --- viper.go | 30 +++++++++++++++++++++--------- viper_test.go | 40 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 61 insertions(+), 9 deletions(-) diff --git a/viper.go b/viper.go index b533306..088cd2b 100644 --- a/viper.go +++ b/viper.go @@ -216,7 +216,8 @@ type Viper struct { aliases map[string]string typeByDefValue bool - onConfigChange func(fsnotify.Event) + onConfigChange func(fsnotify.Event) + stopWatchingFunc func() logger Logger @@ -432,13 +433,10 @@ func (v *Viper) OnConfigChange(run func(in fsnotify.Event)) { } // WatchConfig starts watching a config file for changes. -// The function returned for stop watching manually. -func WatchConfig() func() { return v.WatchConfig() } +func WatchConfig() { v.WatchConfig() } // WatchConfig starts watching a config file for changes. -// The function returned for stop watching manually. -func (v *Viper) WatchConfig() func() { - ctx, cancel := context.WithCancel(context.Background()) +func (v *Viper) WatchConfig() { initWG := sync.WaitGroup{} initWG.Add(1) go func() { @@ -460,6 +458,10 @@ func (v *Viper) WatchConfig() func() { configDir, _ := filepath.Split(configFile) realConfigFile, _ := filepath.EvalSymlinks(filename) + // init the stopWatchingFunc + watchingCtx, cancel := context.WithCancel(context.Background()) + v.stopWatchingFunc = cancel + eventsWG := sync.WaitGroup{} eventsWG.Add(1) go func() { @@ -496,8 +498,9 @@ func (v *Viper) WatchConfig() func() { } eventsWG.Done() return - case <-ctx.Done(): // cancel function called - watcher.Close() + case <-watchingCtx.Done(): // StopWatching function called + eventsWG.Done() + return } } }() @@ -506,7 +509,16 @@ func (v *Viper) WatchConfig() func() { eventsWG.Wait() // now, wait for event loop to end in this go-routine... }() initWG.Wait() // make sure that the go routine above fully ended before returning - return cancel +} + +// StopWatching stop watching a config file for changes. +func StopWatching() { v.StopWatching() } + +// StopWatching stop watching a config file for changes. +func (v *Viper) StopWatching() { + if v.stopWatchingFunc != nil { + v.stopWatchingFunc() + } } // SetConfigFile explicitly defines the path, name and extension of the config file. diff --git a/viper_test.go b/viper_test.go index e0bfc57..06a1671 100644 --- a/viper_test.go +++ b/viper_test.go @@ -2545,6 +2545,46 @@ func TestWatchFile(t *testing.T) { }) } +func TestStopWatching(t *testing.T) { + t.Run( + "file content changed after stop watching", 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) + t.Logf("test config file: %s\n", configFile) + + v.WatchConfig() + v.StopWatching() + + // overwriting the file after StopWatching called + err = ioutil.WriteFile(configFile, []byte("foo: baz\n"), 0o640) + time.Sleep(time.Second) // wait for file changed event + // then the config value should not be changed + require.Nil(t, err) + assert.Equal(t, "bar", v.Get("foo")) + + // watch again + wg := sync.WaitGroup{} + wg.Add(1) + var wgDoneOnce sync.Once // OnConfigChange is called twice on Windows + v.OnConfigChange( + func(in fsnotify.Event) { + t.Logf("config file changed again") + wgDoneOnce.Do(func() { wg.Done() }) + }, + ) + v.WatchConfig() + // overwriting the file after StopWatching and Watch again + err = ioutil.WriteFile(configFile, []byte("foo: qux\n"), 0o640) + wg.Wait() + require.Nil(t, err) + assert.Equal(t, "qux", v.Get("foo")) + }, + ) +} + func TestUnmarshal_DotSeparatorBackwardCompatibility(t *testing.T) { flags := pflag.NewFlagSet("test", pflag.ContinueOnError) flags.String("foo.bar", "cobra_flag", "")