Add label cache, and only fetch missing labels at query time

In simple queries, approximately 1/3rd our queries runtime is spent in
key_value_array. To reduce the time spent, this commit adds a cache of
ids-to-labels in the connector, so we can only fetch those labels we
have never seen before.

Author: JLockerman

pkg/clockcache/cache.go

@@ -36,9 +36,6 @@ type element struct {
 }
 
 func WithMax(max uint64) *Cache {
-	if max < 1 {
-		panic("must have max greater than 0")
-	}
 	return &Cache{
 		elements: make(map[interface{}]*element, max),
 		storage:  make([]element, 0, max),

pkg/pgclient/client.go

@@ -4,9 +4,10 @@ import (
 	"context"
 	"flag"
 	"fmt"
-	"github.com/prometheus/client_golang/prometheus"
 	"runtime"
 
+	"github.com/prometheus/client_golang/prometheus"
+
 	"github.com/allegro/bigcache"
 	"github.com/jackc/pgx/v4/pgxpool"
 
@@ -29,6 +30,7 @@ type Config struct {
 	dbConnectRetries int
 	AsyncAcks        bool
 	ReportInterval   int
+	LabelsCacheSize  uint64
 }
 
 // ParseFlags parses the configuration flags specific to PostgreSQL and TimescaleDB
@@ -42,6 +44,7 @@ func ParseFlags(cfg *Config) *Config {
 	flag.IntVar(&cfg.dbConnectRetries, "db-connect-retries", 0, "How many times to retry connecting to the database")
 	flag.BoolVar(&cfg.AsyncAcks, "async-acks", false, "Ack before data is written to DB")
 	flag.IntVar(&cfg.ReportInterval, "tput-report", 0, "interval in seconds at which throughput should be reported")
+	flag.Uint64Var(&cfg.LabelsCacheSize, "labels-cache-size", 10000, "maximum number of labels to cache")
 	return cfg
 }
 
@@ -84,7 +87,7 @@ func NewClient(cfg *Config, readHist prometheus.ObserverVec) (*Client, error) {
 		log.Error("err starting ingestor", err)
 		return nil, err
 	}
-	reader := pgmodel.NewPgxReaderWithMetricCache(connectionPool, cache)
+	reader := pgmodel.NewPgxReaderWithMetricCache(connectionPool, cache, cfg.LabelsCacheSize)
 
 	queryable := query.NewQueryable(reader.GetQuerier())
 

pkg/pgmodel/end_to_end_tests/nan_test.go

@@ -130,7 +130,7 @@ func TestSQLStaleNaN(t *testing.T) {
 		}
 
 		for _, c := range query {
-			r := NewPgxReader(db, nil)
+			r := NewPgxReader(db, nil, 100)
 			resp, err := r.Read(&c.rrq)
 			startMs := c.rrq.Queries[0].StartTimestampMs
 			endMs := c.rrq.Queries[0].EndTimestampMs

pkg/pgmodel/end_to_end_tests/promql_endpoint_integration_test.go

@@ -3,7 +3,6 @@ package end_to_end_tests
 import (
 	"encoding/json"
 	"fmt"
-	"github.com/prometheus/common/route"
 	"io/ioutil"
 	"net/http"
 	"net/http/httptest"
@@ -13,6 +12,8 @@ import (
 	"testing"
 	"time"
 
+	"github.com/prometheus/common/route"
+
 	"github.com/jackc/pgx/v4/pgxpool"
 	"github.com/prometheus/common/model"
 	"github.com/timescale/timescale-prometheus/pkg/api"
@@ -347,7 +348,7 @@ func TestPromQLQueryEndpoint(t *testing.T) {
 			return
 		}
 
-		r := pgmodel.NewPgxReader(readOnly, nil)
+		r := pgmodel.NewPgxReader(readOnly, nil, 100)
 		queryable := query.NewQueryable(r.GetQuerier())
 		queryEngine := query.NewEngine(log.GetLogger(), time.Minute)
 
@@ -401,7 +402,7 @@ func TestPromQLLabelEndpoints(t *testing.T) {
 			return
 		}
 
-		r := pgmodel.NewPgxReader(readOnly, nil)
+		r := pgmodel.NewPgxReader(readOnly, nil, 100)
 		queryable := query.NewQueryable(r.GetQuerier())
 
 		labelNamesHandler := api.Labels(queryable)

pkg/pgmodel/end_to_end_tests/query_integration_test.go

@@ -534,7 +534,7 @@ func TestSQLQuery(t *testing.T) {
 			t.Fatalf("Cannot run test, not an instance of testing.T")
 		}
 
-		r := NewPgxReader(readOnly, nil)
+		r := NewPgxReader(readOnly, nil, 100)
 		for _, c := range testCases {
 			tester.Run(c.name, func(t *testing.T) {
 				resp, err := r.Read(&c.readRequest)
@@ -915,7 +915,7 @@ func TestPromQL(t *testing.T) {
 			return
 		}
 
-		r := NewPgxReader(readOnly, nil)
+		r := NewPgxReader(readOnly, nil, 100)
 		for _, c := range testCases {
 			tester.Run(c.name, func(t *testing.T) {
 				connResp, connErr := r.Read(c.readRequest)
@@ -1184,7 +1184,7 @@ func TestPushdown(t *testing.T) {
 			return
 		}
 
-		r := NewPgxReader(readOnly, nil)
+		r := NewPgxReader(readOnly, nil, 100)
 		queryable := query.NewQueryable(r.GetQuerier())
 		queryEngine := query.NewEngine(log.GetLogger(), time.Minute)
 

pkg/pgmodel/labels.go

@@ -161,13 +161,8 @@ func (l *Labels) Less(i, j int) bool {
 }
 
 func (l *Labels) Swap(i, j int) {
-	tmp := l.names[j]
-	l.names[j] = l.names[i]
-	l.names[i] = tmp
-
-	tmp = l.values[j]
-	l.values[j] = l.values[i]
-	l.values[i] = tmp
+	l.names[j], l.names[i] = l.names[i], l.names[j]
+	l.values[j], l.values[i] = l.values[i], l.values[j]
 }
 
 // FromLabelMatchers parses protobuf label matchers to Prometheus label matchers.

pkg/pgmodel/pgx.go

@@ -25,6 +25,7 @@ import (
 	"github.com/prometheus/prometheus/pkg/labels"
 	"github.com/prometheus/prometheus/promql/parser"
 	"github.com/prometheus/prometheus/storage"
+	"github.com/timescale/timescale-prometheus/pkg/clockcache"
 	"github.com/timescale/timescale-prometheus/pkg/log"
 	"github.com/timescale/timescale-prometheus/pkg/prompb"
 )
@@ -883,12 +884,13 @@ func (p *pendingBuffer) addReq(req insertDataRequest) bool {
 
 // NewPgxReaderWithMetricCache returns a new DBReader that reads from PostgreSQL using PGX
 // and caches metric table names using the supplied cacher.
-func NewPgxReaderWithMetricCache(c *pgxpool.Pool, cache MetricCache) *DBReader {
+func NewPgxReaderWithMetricCache(c *pgxpool.Pool, cache MetricCache, labelsCacheSize uint64) *DBReader {
 	pi := &pgxQuerier{
 		conn: &pgxConnImpl{
 			conn: c,
 		},
 		metricTableNames: cache,
+		labels:           clockcache.WithMax(labelsCacheSize),
 	}
 
 	return &DBReader{
@@ -897,10 +899,10 @@ func NewPgxReaderWithMetricCache(c *pgxpool.Pool, cache MetricCache) *DBReader {
 }
 
 // NewPgxReader returns a new DBReader that reads that from PostgreSQL using PGX.
-func NewPgxReader(c *pgxpool.Pool, readHist prometheus.ObserverVec) *DBReader {
+func NewPgxReader(c *pgxpool.Pool, readHist prometheus.ObserverVec, labelsCacheSize uint64) *DBReader {
 	metrics, _ := bigcache.NewBigCache(DefaultCacheConfig())
 	cache := &MetricNameCache{metrics}
-	return NewPgxReaderWithMetricCache(c, cache)
+	return NewPgxReaderWithMetricCache(c, cache, labelsCacheSize)
 }
 
 type metricTimeRangeFilter struct {
@@ -912,6 +914,8 @@ type metricTimeRangeFilter struct {
 type pgxQuerier struct {
 	conn             pgxConn
 	metricTableNames MetricCache
+	// contains [int64]labels.Label
+	labels *clockcache.Cache
 }
 
 // HealthCheck implements the healtchecker interface
@@ -933,7 +937,7 @@ func (q *pgxQuerier) Select(mint int64, maxt int64, sortSeries bool, hints *stor
 		return nil, nil, nil, err
 	}
 
-	ss, warn, err := buildSeriesSet(rows, sortSeries)
+	ss, warn, err := buildSeriesSet(rows, sortSeries, q)
 	return ss, topNode, warn, err
 }
 
@@ -963,7 +967,7 @@ func (q *pgxQuerier) Query(query *prompb.Query) ([]*prompb.TimeSeries, error) {
 	results := make([]*prompb.TimeSeries, 0, len(rows))
 
 	for _, r := range rows {
-		ts, err := buildTimeSeries(r)
+		ts, err := buildTimeSeries(r, q)
 
 		if err != nil {
 			return nil, err
@@ -1021,6 +1025,82 @@ func (q *pgxQuerier) LabelValues(labelName string) ([]string, error) {
 	return labelValues, nil
 }
 
+const GetLabelsSQL = "SELECT (key_value_array($1::int[])).*"
+
+type labelQuerier interface {
+	getLabelsForIds(ids []int64) (lls labels.Labels, err error)
+}
+
+func (q *pgxQuerier) getPrompbLabelsForIds(ids []int64) (lls []prompb.Label, err error) {
+	ll, err := q.getLabelsForIds(ids)
+	if err != nil {
+		return
+	}
+	lls = make([]prompb.Label, len(ll))
+	for i := range ll {
+		lls[i] = prompb.Label{Name: ll[i].Name, Value: ll[i].Value}
+	}
+	return
+}
+
+func (q *pgxQuerier) getLabelsForIds(ids []int64) (lls labels.Labels, err error) {
+	keys := make([]interface{}, len(ids))
+	values := make([]interface{}, len(ids))
+	for i := range ids {
+		keys[i] = ids[i]
+	}
+	numHits := q.labels.GetValues(keys, values)
+
+	if numHits < len(ids) {
+		err = q.fetchMissingLabels(keys[numHits:], ids[numHits:], values[numHits:])
+		if err != nil {
+			return
+		}
+	}
+
+	lls = make([]labels.Label, 0, len(values))
+	for i := range values {
+		lls = append(lls, values[i].(labels.Label))
+	}
+
+	return
+}
+
+func (q *pgxQuerier) fetchMissingLabels(misses []interface{}, missedIds []int64, newLabels []interface{}) error {
+	for i := range misses {
+		missedIds[i] = misses[i].(int64)
+	}
+	rows, err := q.conn.Query(context.Background(), GetLabelsSQL, missedIds)
+	if err != nil {
+		return err
+	}
+	defer rows.Close()
+
+	for rows.Next() {
+		var keys []string
+		var vals []string
+		err = rows.Scan(&keys, &vals)
+		if err != nil {
+			return err
+		}
+		if len(keys) != len(vals) {
+			return fmt.Errorf("query returned a mismatch in timestamps and values: %d, %d", len(keys), len(vals))
+		}
+		if len(keys) != len(misses) {
+			return fmt.Errorf("query returned wrong number of labels: %d, %d", len(misses), len(keys))
+		}
+
+		for i := range misses {
+			newLabels[i] = labels.Label{Name: keys[i], Value: vals[i]}
+		}
+		numInserted := q.labels.InsertBatch(misses, newLabels)
+		if numInserted < len(misses) {
+			log.Warn("msg", "labels cache starving, may need to increase size")
+		}
+	}
+	return nil
+}
+
 func (q *pgxQuerier) getResultRows(startTimestamp int64, endTimestamp int64, hints *storage.SelectHints, path []parser.Node, matchers []*labels.Matcher) ([]pgx.Rows, parser.Node, error) {
 
 	metric, cases, values, err := buildSubQueries(matchers)