Add support for leader election

In HA setup two or more Prometheus instances scrape the same data and send it to adapters.
We want to prevent multiple adapters to write the same/similar data (no consistency guarantees in HA Prometheus setup) in parallel. The idea is to use leader election and allow only one adapter (leader) to write to the database. Provided leader election implementation relies on PostgreSQL advisory locks and does not provide strong data guarantess (some duplicates or data loss possible during failover - btw Prometheus HA doesn't provide any guarantees that multiple Prometheus nodes would see the same data). REST interface for leader election should enable plugging in any external leader election system you might use already (eg. Zookeeper)

Author: niksajakovljevic

Gopkg.lock

@@ -35,6 +35,7 @@ import (
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/common/model"
 
+	"database/sql"
 	"fmt"
 	"github.com/prometheus/client_model/go"
 	"github.com/prometheus/prometheus/prompb"
@@ -47,6 +48,8 @@ type config struct {
 	pgPrometheusConfig pgprometheus.Config
 	logLevel           string
 	readOnly           bool
+	haGroupLockId      int
+	restElection       bool
 }
 
 const (
@@ -90,7 +93,8 @@ var (
 		},
 		[]string{"path"},
 	)
-	writeThroughtput = util.NewThroughputCalc(tickInterval)
+	writeThroughput = util.NewThroughputCalc(tickInterval)
+	elector         *util.Elector
 )
 
 func init() {
@@ -99,7 +103,7 @@ func init() {
 	prometheus.MustRegister(failedSamples)
 	prometheus.MustRegister(sentBatchDuration)
 	prometheus.MustRegister(httpRequestDuration)
-	writeThroughtput.Start()
+	writeThroughput.Start()
 }
 
 func main() {
@@ -110,6 +114,12 @@ func main() {
 	http.Handle(cfg.telemetryPath, prometheus.Handler())
 
 	writer, reader := buildClients(cfg)
+	pgClient, ok := writer.(*pgprometheus.Client)
+	if ok {
+		elector = initElector(cfg, pgClient.DB)
+	} else {
+		log.Info("msg", "Running in read-only mode. This instance can't participate in leader election")
+	}
 
 	http.Handle("/write", timeHandler("write", write(writer)))
 	http.Handle("/read", timeHandler("read", read(reader)))
@@ -137,7 +147,8 @@ func parseFlags() *config {
 	flag.StringVar(&cfg.telemetryPath, "web.telemetry-path", "/metrics", "Address to listen on for web endpoints.")
 	flag.StringVar(&cfg.logLevel, "log.level", "debug", "The log level to use [ \"error\", \"warn\", \"info\", \"debug\" ].")
 	flag.BoolVar(&cfg.readOnly, "read.only", false, "Read-only mode. Don't write to database.")
-
+	flag.IntVar(&cfg.haGroupLockId, "leader-election.pg-advisory-lock-id", 0, "Unique advisory lock id per adapter high-availability group. Set it if you want to use leader election implementation based on PostgreSQL advisory lock")
+	flag.BoolVar(&cfg.restElection, "leader-election.rest", false, "Enable REST interface for the leader election")
 	flag.Parse()
 
 	return cfg
@@ -173,6 +184,28 @@ func buildClients(cfg *config) (writer, reader) {
 	return pgClient, pgClient
 }
 
+func initElector(cfg *config, db *sql.DB) *util.Elector {
+	if cfg.restElection && cfg.haGroupLockId != 0 {
+		log.Error("msg", "Use either REST or PgAdvisoryLock for the leader election")
+		os.Exit(1)
+	}
+	if cfg.restElection {
+		return util.NewElector(util.NewRestElection())
+	}
+	if cfg.haGroupLockId != 0 {
+		lock, err := util.NewPgAdvisoryLock(cfg.haGroupLockId, db)
+		if err != nil {
+			log.Error("msg", "Error creating advisory lock", "haGroupLockId", cfg.haGroupLockId, "err", err)
+			os.Exit(1)
+		}
+		log.Info("msg", "Initialized leader election based on PostgreSQL advisory lock")
+		return util.NewElector(lock)
+	} else {
+		log.Warn("msg", "No adapter leader election. Group lock id is not set. Possible duplicate write load if running adapter in high-availability mode")
+		return nil
+	}
+}
+
 func write(writer writer) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		compressed, err := ioutil.ReadAll(r.Body)
@@ -208,10 +241,10 @@ func write(writer writer) http.Handler {
 		if err != nil {
 			log.Warn("msg", "Couldn't get a counter", "labelValue", writer.Name(), "err", err)
 		}
-		writeThroughtput.SetCurrent(getCounterValue(counter))
+		writeThroughput.SetCurrent(getCounterValue(counter))
 
 		select {
-		case d := <-writeThroughtput.Values:
+		case d := <-writeThroughput.Values:
 			log.Info("msg", "Samples write throughput", "samples/sec", d)
 		default:
 		}
@@ -306,7 +339,21 @@ func protoToSamples(req *prompb.WriteRequest) model.Samples {
 
 func sendSamples(w writer, samples model.Samples) error {
 	begin := time.Now()
-	err := w.Write(samples)
+	shouldWrite := true
+	var err error
+	if elector != nil {
+		shouldWrite, err = elector.IsLeader()
+		if err != nil {
+			log.Error("msg", "IsLeader check failed", "err", err)
+			return err
+		}
+	}
+	if shouldWrite {
+		err = w.Write(samples)
+	} else {
+		log.Debug("msg", fmt.Sprintf("Election id %v: Instance is not a leader. Can't write data", elector.Id()))
+		return nil
+	}
 	duration := time.Since(begin).Seconds()
 	if err != nil {
 		failedSamples.WithLabelValues(w.Name()).Add(float64(len(samples)))

main.go

@@ -64,7 +64,7 @@ func ParseFlags(cfg *Config) *Config {
 
 // Client sends Prometheus samples to PostgreSQL
 type Client struct {
-	db  *sql.DB
+	DB  *sql.DB
 	cfg *Config
 }
 
@@ -101,7 +101,7 @@ func NewClient(cfg *Config) *Client {
 	db.SetMaxIdleConns(cfg.maxIdleConns)
 
 	client := &Client{
-		db:  db,
+		DB:  db,
 		cfg: cfg,
 	}
 
@@ -121,7 +121,7 @@ func NewClient(cfg *Config) *Client {
 }
 
 func (c *Client) setupPgPrometheus() error {
-	tx, err := c.db.Begin()
+	tx, err := c.DB.Begin()
 
 	if err != nil {
 		return err
@@ -193,7 +193,7 @@ func metricString(m model.Metric) string {
 // Write implements the Writer interface and writes metric samples to the database
 func (c *Client) Write(samples model.Samples) error {
 	begin := time.Now()
-	tx, err := c.db.Begin()
+	tx, err := c.DB.Begin()
 
 	if err != nil {
 		log.Error("msg", "Error on Begin when writing samples", "err", err)
@@ -360,7 +360,7 @@ func (c *Client) Read(req *prompb.ReadRequest) (*prompb.ReadResponse, error) {
 
 		log.Debug("msg", "Executed query", "query", command)
 
-		rows, err := c.db.Query(command)
+		rows, err := c.DB.Query(command)
 
 		if err != nil {
 			return nil, err
@@ -439,7 +439,7 @@ func (c *Client) Read(req *prompb.ReadRequest) (*prompb.ReadResponse, error) {
 
 // HealthCheck implements the healtcheck interface
 func (c *Client) HealthCheck() error {
-	rows, err := c.db.Query("SELECT 1")
+	rows, err := c.DB.Query("SELECT 1")
 
 	if err != nil {
 		log.Debug("msg", "Health check error", "err", err)

postgresql/client.go

@@ -7,6 +7,7 @@ import (
 	"github.com/prometheus/common/model"
 	"github.com/prometheus/prometheus/prompb"
 	"github.com/timescale/prometheus-postgresql-adapter/log"
+	"github.com/timescale/prometheus-postgresql-adapter/util"
 	"testing"
 )
 
@@ -30,17 +31,17 @@ func TestBuildCommand(t *testing.T) {
 		StartTimestampMs: 0,
 		EndTimestampMs:   20000,
 		Matchers: []*prompb.LabelMatcher{
-			&prompb.LabelMatcher{
+			{
 				Type:  prompb.LabelMatcher_EQ,
 				Name:  "__name__",
 				Value: "cpu_usage",
 			},
-			&prompb.LabelMatcher{
+			{
 				Type:  prompb.LabelMatcher_EQ,
 				Name:  "job",
 				Value: "nginx",
 			},
-			&prompb.LabelMatcher{
+			{
 				Type:  prompb.LabelMatcher_RE,
 				Name:  "host",
 				Value: "local.*",
@@ -58,30 +59,7 @@ func TestBuildCommand(t *testing.T) {
 }
 
 func TestWriteCommand(t *testing.T) {
-	flag.Parse()
-	if len(*database) == 0 {
-		t.Skip()
-	}
-
-	db, err := sql.Open("postgres", "host=localhost user=postgres sslmode=disable")
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	_, err = db.Exec(fmt.Sprintf("DROP DATABASE IF EXISTS %s", *database))
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	_, err = db.Exec(fmt.Sprintf("CREATE DATABASE %s", *database))
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	err = db.Close()
-	if err != nil {
-		t.Fatal(err)
-	}
+	dbSetup(t)
 
 	cfg := &Config{}
 	ParseFlags(cfg)
@@ -90,30 +68,30 @@ func TestWriteCommand(t *testing.T) {
 	c := NewClient(cfg)
 
 	sample := []*model.Sample{
-		&model.Sample{
+		{
 			Metric: model.Metric{
 				model.MetricNameLabel: "test_metric",
 				"label1":              "1",
 			},
 			Value:     123.1,
 			Timestamp: 1234567,
 		},
-		&model.Sample{
+		{
 			Metric: model.Metric{
 				model.MetricNameLabel: "test_metric",
 				"label1":              "1",
 			},
 			Value:     123.2,
 			Timestamp: 1234568,
 		},
-		&model.Sample{
+		{
 			Metric: model.Metric{
 				model.MetricNameLabel: "test_metric",
 			},
 			Value:     123.2,
 			Timestamp: 1234569,
 		},
-		&model.Sample{
+		{
 			Metric: model.Metric{
 				model.MetricNameLabel: "test_metric_2",
 				"label1":              "1",
@@ -125,7 +103,7 @@ func TestWriteCommand(t *testing.T) {
 
 	c.Write(sample)
 
-	db, err = sql.Open("postgres", fmt.Sprintf("host=localhost dbname=%s user=postgres sslmode=disable", *database))
+	db, err := sql.Open("postgres", fmt.Sprintf("host=localhost dbname=%s user=postgres sslmode=disable", *database))
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -140,3 +118,59 @@ func TestWriteCommand(t *testing.T) {
 		t.Fatal("Wrong cnt: ", cnt)
 	}
 }
+
+func TestPgAdvisoryLock(t *testing.T) {
+	db := dbSetup(t)
+	lock, err := util.NewPgAdvisoryLock(1, db)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !lock.Locked() {
+		t.Error("Couldn't obtain the lock")
+	}
+
+	newLock, err := util.NewPgAdvisoryLock(1, db)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if newLock.Locked() {
+		t.Error("Lock should have already been taken")
+	}
+
+	if err = lock.Release(); err != nil {
+		t.Errorf("Failed to release a lock. Error: %v", err)
+	}
+
+	if lock.Locked() {
+		t.Error("Should be unlocked after release")
+	}
+
+	newLock.TryLock()
+
+	if !newLock.Locked() {
+		t.Error("New lock should take over")
+	}
+}
+
+func dbSetup(t *testing.T) *sql.DB {
+	flag.Parse()
+	if len(*database) == 0 {
+		t.Skip()
+	}
+
+	db, err := sql.Open("postgres", "host=localhost user=postgres sslmode=disable")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	_, err = db.Exec(fmt.Sprintf("DROP DATABASE IF EXISTS %s", *database))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	_, err = db.Exec(fmt.Sprintf("CREATE DATABASE %s", *database))
+	if err != nil {
+		t.Fatal(err)
+	}
+	return db
+}