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 provides at least once semantics for metrics data (some duplicates possible during failover and re-election).

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,7 @@ type config struct {
 	pgPrometheusConfig pgprometheus.Config
 	logLevel           string
 	readOnly           bool
+	haGroupLockId      int
 }
 
 const (
@@ -91,6 +93,7 @@ var (
 		[]string{"path"},
 	)
 	writeThroughtput = util.NewThroughputCalc(tickInterval)
+	election         util.Election
 )
 
 func init() {
@@ -110,6 +113,7 @@ func main() {
 	http.Handle(cfg.telemetryPath, prometheus.Handler())
 
 	writer, reader := buildClients(cfg)
+	election = initElection(cfg, writer.(*pgprometheus.Client).DB)
 
 	http.Handle("/write", timeHandler("write", write(writer)))
 	http.Handle("/read", timeHandler("read", read(reader)))
@@ -137,6 +141,7 @@ 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, "ha.group-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.Parse()
 
@@ -173,6 +178,21 @@ func buildClients(cfg *config) (writer, reader) {
 	return pgClient, pgClient
 }
 
+func initElection(cfg *config, db *sql.DB) util.Election {
+	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 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)
@@ -306,7 +326,31 @@ func protoToSamples(req *prompb.WriteRequest) model.Samples {
 
 func sendSamples(w writer, samples model.Samples) error {
 	begin := time.Now()
-	err := w.Write(samples)
+	shouldWrite := true
+	if election != nil {
+		isLeader, err := election.IsLeader()
+		if err != nil {
+			log.Error("msg", "IsLeader check failed", "err", err)
+			return err
+		}
+		if isLeader {
+			shouldWrite = isLeader
+		} else {
+			isLeader, err := election.BecomeLeader()
+			if err != nil {
+				log.Error("msg", "Error occurred while trying to become a leader", "err", err)
+				return err
+			}
+			shouldWrite = isLeader
+		}
+	}
+	var err error
+	if shouldWrite {
+		err = w.Write(samples)
+	} else {
+		log.Debug("msg", fmt.Sprintf("Election id %v: Instance is not a leader. Can't write data", election.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
+}

postgresql/client_test.go

@@ -0,0 +1,13 @@
+package util
+
+// Election defines an interface for adapter leader election.
+// If you are running Prometheus in HA mode where each Prometheus instance sends data to corresponding adapter you probably
+// want to allow writes into the database from only one adapter at the time. We need to elect a leader who can write to
+// the database. If leader goes down, another leader is elected. Look at `lock.go` for an implementation based on PostgreSQL
+// advisory locks. Should be easy to plug in different leader election implementations.
+type Election interface {
+	Id() string
+	BecomeLeader() (bool, error)
+	IsLeader() (bool, error)
+	Resign() error
+}