Add support for dynamically reloading sampling strategies

plugin/sampling/strategystore/static/fixtures/strategies-hot-reload.json

@@ -0,0 +1,11 @@
+{
+  "default_strategy": {
+  },
+  "service_strategies": [
+    {
+      "service": "foo",
+      "type": "probabilistic",
+      "param": 0.5
+    }
+  ]
+}

plugin/sampling/strategystore/static/strategy_store.go

@@ -20,7 +20,10 @@ import (
 	"encoding/json"
 	"fmt"
 	"io/ioutil"
+	"path/filepath"
+	"sync"
 
+	"github.com/fsnotify/fsnotify"
 	"github.com/pkg/errors"
 	"go.uber.org/zap"
 
@@ -29,31 +32,102 @@ import (
 )
 
 type strategyStore struct {
+	// to allow concurrent update of sampling strategies
+	lock sync.RWMutex
+
 	logger *zap.Logger
 
 	defaultStrategy   *sampling.SamplingStrategyResponse
 	serviceStrategies map[string]*sampling.SamplingStrategyResponse
 }
 
+func (h *strategyStore) loadAndParseStrategies(strategiesFile string) error {
+	s, err := loadStrategies(strategiesFile)
+	if err != nil {
+		h.logger.Warn("using the last saved configuration for sampling strategies.", zap.Error(err))
+		return err
+	}
+
+	h.parseStrategies(s)
+	return nil
+}
+
 // NewStrategyStore creates a strategy store that holds static sampling strategies.
 func NewStrategyStore(options Options, logger *zap.Logger) (ss.StrategyStore, error) {
 	h := &strategyStore{
 		logger:            logger,
 		serviceStrategies: make(map[string]*sampling.SamplingStrategyResponse),
 	}
-	strategies, err := loadStrategies(options.StrategiesFile)
-	if err != nil {
+
+	// Read strategies
+	if err := h.loadAndParseStrategies(options.StrategiesFile); err != nil {
 		return nil, err
 	}
-	h.parseStrategies(strategies)
+
+	// Watch strategies file for changes.
+	watcher, err := fsnotify.NewWatcher()
+	if err != nil {
+		logger.Error("failed to create a new watcher for the sampling strategies file", zap.Error(err))
+	}
+
+	go h.runWatcherLoop(watcher, options.StrategiesFile)
+
+	if err := watcher.Add(options.StrategiesFile); err != nil {
+		logger.Error("error adding watcher to file", zap.String("file", options.StrategiesFile), zap.Error(err))
+	} else {
+		logger.Info("watching", zap.String("file", options.StrategiesFile))
+	}
+
+	dir := filepath.Dir(options.StrategiesFile)
+	err = watcher.Add(dir)
+	if err := watcher.Add(dir); err != nil {
+		h.logger.Error("error adding watcher to dir", zap.String("dir", dir), zap.Error(err))
+	} else {
+		h.logger.Info("watching", zap.String("dir", dir))
+	}
+
 	return h, nil
 }
 
+func (h *strategyStore) runWatcherLoop(watcher *fsnotify.Watcher, strategiesFile string) {
+	for {
+		select {
+		case event := <-watcher.Events:
+			if event.Op&fsnotify.Remove == fsnotify.Remove {
+				if event.Name == strategiesFile {
+					h.logger.Warn("the sampling strategies file has been removed")
+
+					// This is a workaround for k8s configmaps. Since k8s loads configmaps as
+					// symlinked files within the containers, changes to the configmap register
+					// as `fsnotify.Remove` events.
+					if err := watcher.Add(strategiesFile); err != nil {
+						h.logger.Warn("cannot watch sampling strategy config file", zap.Error(err))
+					}
+
+					h.loadAndParseStrategies(strategiesFile)
+				}
+				continue
+			}
+			if event.Op&fsnotify.Write == fsnotify.Write {
+				h.loadAndParseStrategies(strategiesFile)
+			}
+		case err, ok := <-watcher.Errors:
+			if !ok {
+				return
+			}
+			h.logger.Error("file watcher error", zap.Error(err))
+		}
+	}
+}
+
 // GetSamplingStrategy implements StrategyStore#GetSamplingStrategy.
 func (h *strategyStore) GetSamplingStrategy(serviceName string) (*sampling.SamplingStrategyResponse, error) {
+	h.lock.RLock()
+	defer h.lock.RUnlock()
 	if strategy, ok := h.serviceStrategies[serviceName]; ok {
 		return strategy, nil
 	}
+
 	return h.defaultStrategy, nil
 }
 
@@ -74,7 +148,11 @@ func loadStrategies(strategiesFile string) (*strategies, error) {
 }
 
 func (h *strategyStore) parseStrategies(strategies *strategies) {
+	h.lock.Lock()
+	defer h.lock.Unlock()
+
 	h.defaultStrategy = defaultStrategyResponse()
+
 	if strategies == nil {
 		h.logger.Info("No sampling strategies provided, using defaults")
 		return

plugin/sampling/strategystore/static/strategy_store_test.go

@@ -16,7 +16,10 @@ package static
 
 import (
 	"fmt"
+	"io/ioutil"
+	"os"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -218,6 +221,54 @@ func TestParseStrategy(t *testing.T) {
 	assert.Contains(t, buf.String(), "Failed to parse sampling strategy")
 }
 
+func TestHotReloadSamplingStrategiesTempFile(t *testing.T) {
+	tmpfile, err := ioutil.TempFile("fixtures", "strategies-hot-reload.*.json")
+	assert.NoError(t, err)
+
+	tmpFileName := tmpfile.Name()
+	defer os.Remove(tmpFileName)
+
+	content, err := ioutil.ReadFile("fixtures/strategies-hot-reload.json")
+	assert.NoError(t, err)
+
+	err = ioutil.WriteFile(tmpFileName, content, 0644)
+	assert.NoError(t, err)
+
+	logger, _ := testutils.NewLogger()
+	store, err := NewStrategyStore(Options{StrategiesFile: tmpFileName}, logger)
+	assert.NoError(t, err)
+
+	s, err := store.GetSamplingStrategy("foo")
+	require.NoError(t, err)
+	assert.EqualValues(t, makeResponse(sampling.SamplingStrategyType_PROBABILISTIC, 0.5), *s)
+
+	newContent, err := ioutil.ReadFile("fixtures/strategies.json")
+	assert.NoError(t, err)
+	err = ioutil.WriteFile(tmpFileName, newContent, 0644)
+	assert.NoError(t, err)
+
+	done := make(chan bool)
+	go func() {
+		for {
+			s, err := store.GetSamplingStrategy("foo")
+			require.NoError(t, err)
+			if s.GetProbabilisticSampling().SamplingRate == 0.8 {
+				done <- true
+			}
+			time.Sleep(10 * time.Millisecond)
+		}
+	}()
+
+	select {
+	case <-done:
+		s, err := store.GetSamplingStrategy("foo")
+		require.NoError(t, err)
+		assert.EqualValues(t, makeResponse(sampling.SamplingStrategyType_PROBABILISTIC, 0.8), *s)
+	case <-time.After(time.Second):
+		assert.Fail(t, "timed out waiting for the hot reload to kick in")
+	}
+}
+
 func makeResponse(samplerType sampling.SamplingStrategyType, param float64) (resp sampling.SamplingStrategyResponse) {
 	resp.StrategyType = samplerType
 	if samplerType == sampling.SamplingStrategyType_PROBABILISTIC {