Implement ApplyConfig for Loki

Makefile

@@ -170,7 +170,7 @@ lint:
 # for packages that have known race detection issues
 test:
 	GOGC=10 go test $(MOD_FLAG) -race -cover -coverprofile=cover.out -p=4 ./...
-	GOGC=10 go test $(MOD_FLAG) -cover -coverprofile=cover-norace.out -p=4 ./pkg/integrations/node_exporter
+	GOGC=10 go test $(MOD_FLAG) -cover -coverprofile=cover-norace.out -p=4 ./pkg/integrations/node_exporter ./pkg/loki
 
 clean:
 	rm -rf cmd/agent/agent

go.mod

@@ -12,6 +12,7 @@ require (
 	github.com/gogo/protobuf v1.3.2
 	github.com/golang/protobuf v1.4.3
 	github.com/google/dnsmasq_exporter v0.0.0-00010101000000-000000000000
+	github.com/google/go-cmp v0.5.4
 	github.com/gorilla/mux v1.8.0
 	github.com/grafana/loki v1.6.2-0.20210205130758-59a34f9867ce
 	github.com/hashicorp/consul/api v1.8.1

pkg/loki/loki.go

@@ -4,9 +4,12 @@ package loki
 import (
 	"fmt"
 	"os"
+	"sync"
 
 	"github.com/go-kit/kit/log"
 	"github.com/go-kit/kit/log/level"
+	"github.com/google/go-cmp/cmp"
+	"github.com/grafana/agent/pkg/util"
 	"github.com/grafana/loki/pkg/promtail"
 	"github.com/grafana/loki/pkg/promtail/client"
 	"github.com/grafana/loki/pkg/promtail/config"
@@ -20,50 +23,160 @@ func init() {
 }
 
 type Loki struct {
-	promtails []*promtail.Promtail
+	mut sync.Mutex
+
+	reg       prometheus.Registerer
+	l         log.Logger
+	instances map[string]*Instance
 }
 
 // New creates and starts Loki log collection.
 func New(reg prometheus.Registerer, c Config, l log.Logger) (*Loki, error) {
 	l = log.With(l, "component", "loki")
 
+	loki := &Loki{
+		instances: make(map[string]*Instance),
+		reg:       reg,
+		l:         log.With(l, "component", "loki"),
+	}
+	if err := loki.ApplyConfig(c); err != nil {
+		return nil, err
+	}
+	return loki, nil
+}
+
+// ApplyConfig updates Loki with a new Config.
+func (l *Loki) ApplyConfig(c Config) error {
+	l.mut.Lock()
+	defer l.mut.Unlock()
+
 	if c.PositionsDirectory != "" {
 		err := os.MkdirAll(c.PositionsDirectory, 0700)
 		if err != nil {
-			level.Warn(l).Log("msg", "failed to create the positions directory. logs may be unable to save their position", "path", c.PositionsDirectory, "err", err)
+			level.Warn(l.l).Log("msg", "failed to create the positions directory. logs may be unable to save their position", "path", c.PositionsDirectory, "err", err)
 		}
 	}
 
-	promtails := make([]*promtail.Promtail, 0, len(c.Configs))
+	newInstances := make(map[string]*Instance, len(c.Configs))
 
 	for _, ic := range c.Configs {
-		if len(ic.ClientConfigs) == 0 {
-			level.Info(l).Log("msg", "skipping creation of a promtail because no client_configs are present", "config", ic.Name)
+		// If an old instance existed, update it and move it to the new map.
+		if old, ok := l.instances[ic.Name]; ok {
+			err := old.ApplyConfig(ic)
+			if err != nil {
+				return err
+			}
+
+			newInstances[ic.Name] = old
 			continue
 		}
 
-		r := prometheus.WrapRegistererWith(prometheus.Labels{"loki_config": ic.Name}, reg)
-		il := log.With(l, "loki_config", ic.Name)
-
-		p, err := promtail.New(config.Config{
-			ServerConfig:    server.Config{Disable: true},
-			ClientConfigs:   ic.ClientConfigs,
-			PositionsConfig: ic.PositionsConfig,
-			ScrapeConfig:    ic.ScrapeConfig,
-			TargetConfig:    ic.TargetConfig,
-		}, false, promtail.WithLogger(il), promtail.WithRegisterer(r))
+		inst, err := NewInstance(l.reg, ic, l.l)
 		if err != nil {
-			return nil, err
+			return fmt.Errorf("unable to apply config for %s: %w", ic.Name, err)
 		}
+		newInstances[ic.Name] = inst
+	}
 
-		promtails = append(promtails, p)
+	// Any promtail in l.instances that isn't in newInstances has been removed
+	// from the config. Stop them before replacing the map.
+	for key, i := range l.instances {
+		if _, exist := newInstances[key]; exist {
+			continue
+		}
+		i.Stop()
 	}
+	l.instances = newInstances
 
-	return &Loki{promtails: promtails}, nil
+	return nil
 }
 
 func (l *Loki) Stop() {
-	for _, p := range l.promtails {
-		p.Shutdown()
+	l.mut.Lock()
+	defer l.mut.Unlock()
+
+	for _, i := range l.instances {
+		i.Stop()
+	}
+}
+
+// Instance is an individual Loki instance.
+type Instance struct {
+	mut sync.Mutex
+
+	cfg *InstanceConfig
+	log log.Logger
+	reg *util.Unregisterer
+
+	promtail *promtail.Promtail
+}
+
+// NewInstance creates and starts a Loki instance.
+func NewInstance(reg prometheus.Registerer, c *InstanceConfig, l log.Logger) (*Instance, error) {
+	instReg := prometheus.WrapRegistererWith(prometheus.Labels{"loki_config": c.Name}, reg)
+
+	inst := Instance{
+		reg: util.WrapWithUnregisterer(instReg),
+		log: log.With(l, "loki_config", c.Name),
+	}
+	if err := inst.ApplyConfig(c); err != nil {
+		return nil, err
+	}
+	return &inst, nil
+}
+
+// ApplyConfig will apply a new InstanceConfig. If the config hasn't changed,
+// then nothing will happen, otherwise the old Promtail will be stopped and
+// then replaced with a new one.
+func (i *Instance) ApplyConfig(c *InstanceConfig) error {
+	i.mut.Lock()
+	defer i.mut.Unlock()
+
+	// No-op if the configs haven't changed.
+	if cmp.Equal(c, i.cfg) {
+		level.Debug(i.log).Log("msg", "instance config hasn't changed, not recreating Promtail")
+		return nil
+	}
+	i.cfg = c
+
+	if i.promtail != nil {
+		i.promtail.Shutdown()
+		i.promtail = nil
+	}
+
+	// Unregister all existing metrics before trying to create a new instance.
+	if !i.reg.UnregisterAll() {
+		// If UnregisterAll fails, we need to abort, otherwise the new promtail
+		// would try to re-register an existing metric and might panic.
+		return fmt.Errorf("failed to unregister all metrics from previous promtail. THIS IS A BUG")
+	}
+
+	if len(c.ClientConfigs) == 0 {
+		level.Debug(i.log).Log("msg", "skipping creation of a promtail because no client_configs are present")
+		return nil
+	}
+
+	p, err := promtail.New(config.Config{
+		ServerConfig:    server.Config{Disable: true},
+		ClientConfigs:   c.ClientConfigs,
+		PositionsConfig: c.PositionsConfig,
+		ScrapeConfig:    c.ScrapeConfig,
+		TargetConfig:    c.TargetConfig,
+	}, false, promtail.WithLogger(i.log), promtail.WithRegisterer(i.reg))
+	if err != nil {
+		return fmt.Errorf("unable to create Loki logging instance: %w", err)
+	}
+
+	i.promtail = p
+	return nil
+}
+
+func (i *Instance) Stop() {
+	i.mut.Lock()
+	defer i.mut.Unlock()
+
+	if i.promtail != nil {
+		i.promtail.Shutdown()
+		i.promtail = nil
 	}
 }

pkg/loki/loki_test.go

@@ -0,0 +1,135 @@
+//+build !race
+
+package loki
+
+import (
+	"fmt"
+	"io/ioutil"
+	"net"
+	"net/http"
+	"os"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/go-kit/kit/log"
+	"github.com/grafana/agent/pkg/util"
+	"github.com/grafana/loki/pkg/distributor"
+	"github.com/grafana/loki/pkg/logproto"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/stretchr/testify/require"
+	"gopkg.in/yaml.v2"
+)
+
+func TestLoki(t *testing.T) {
+	//
+	// Create a temporary file to tail
+	//
+	positionsDir, err := ioutil.TempDir(os.TempDir(), "positions-*")
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		_ = os.RemoveAll(positionsDir)
+	})
+
+	tmpFile, err := ioutil.TempFile(os.TempDir(), "*.log")
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		_ = os.RemoveAll(tmpFile.Name())
+	})
+
+	//
+	// Listen for push requests and pass them through to a channel
+	//
+	pushes := make(chan *logproto.PushRequest)
+
+	lis, err := net.Listen("tcp", "127.0.0.1:0")
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		require.NoError(t, lis.Close())
+	})
+	go func() {
+		_ = http.Serve(lis, http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
+			req, err := distributor.ParseRequest(r)
+			require.NoError(t, err)
+
+			pushes <- req
+			_, _ = rw.Write(nil)
+		}))
+	}()
+
+	//
+	// Launch Loki so it starts tailing the file and writes to our server.
+	//
+	cfgText := util.Untab(fmt.Sprintf(`
+positions_directory: %s
+configs:
+- name: default
+  clients:
+  - url: http://%s/loki/api/v1/push
+		batchwait: 50ms
+		batchsize: 1
+  scrape_configs:
+  - job_name: system
+    static_configs:
+    - targets: [localhost]
+      labels:
+        job: test
+        __path__: %s
+	`, positionsDir, lis.Addr().String(), tmpFile.Name()))
+
+	var cfg Config
+	dec := yaml.NewDecoder(strings.NewReader(cfgText))
+	dec.SetStrict(true)
+	require.NoError(t, dec.Decode(&cfg))
+
+	logger := log.NewSyncLogger(log.NewNopLogger())
+	l, err := New(prometheus.NewRegistry(), cfg, logger)
+	require.NoError(t, err)
+	defer l.Stop()
+
+	//
+	// Write a log line and wait for it to come through.
+	//
+	fmt.Fprintf(tmpFile, "Hello, world!\n")
+	select {
+	case <-time.After(time.Second * 30):
+		require.FailNow(t, "timed out waiting for data to be pushed")
+	case req := <-pushes:
+		require.Equal(t, "Hello, world!", req.Streams[0].Entries[0].Line)
+	}
+
+	//
+	// Apply a new config and write a new line.
+	//
+	cfgText = util.Untab(fmt.Sprintf(`
+positions_directory: %s
+configs:
+- name: default
+  clients:
+  - url: http://%s/loki/api/v1/push
+		batchwait: 50ms
+		batchsize: 5
+  scrape_configs:
+  - job_name: system
+    static_configs:
+    - targets: [localhost]
+      labels:
+        job: test-2
+        __path__: %s
+	`, positionsDir, lis.Addr().String(), tmpFile.Name()))
+
+	var newCfg Config
+	dec = yaml.NewDecoder(strings.NewReader(cfgText))
+	dec.SetStrict(true)
+	require.NoError(t, dec.Decode(&newCfg))
+
+	require.NoError(t, l.ApplyConfig(newCfg))
+
+	fmt.Fprintf(tmpFile, "Hello again!\n")
+	select {
+	case <-time.After(time.Second * 30):
+		require.FailNow(t, "timed out waiting for data to be pushed")
+	case req := <-pushes:
+		require.Equal(t, "Hello again!", req.Streams[0].Entries[0].Line)
+	}
+}