cloudflare/pint

Public

mirrored from https://github.com/cloudflare/pintAvailable

CodeCommitsIssuesPull requestsActionsInsightsSecurity
v0.25.0

Branches

Tags

  • No tags available.
0Branches0Tags
Go to file
Add file
Code

Clone

HTTPS

Download ZIP

internal/promapi/prometheus.go

157lines · modecode

1package promapi
2
3import (
4 "sync"
5 "time"
6
7 lru "github.com/hashicorp/golang-lru"
8 "github.com/prometheus/client_golang/api"
9 v1 "github.com/prometheus/client_golang/api/prometheus/v1"
10 "github.com/rs/zerolog/log"
11 "go.uber.org/ratelimit"
12
13 "github.com/cloudflare/pint/internal/output"
14)
15
16var cacheExpiry = time.Minute * 5
17
18type QueryError struct {
19 err error
20 msg string
21}
22
23func (qe QueryError) Error() string {
24 return qe.msg
25}
26
27func (qe QueryError) Unwrap() error {
28 return qe.err
29}
30
31type querier interface {
32 Endpoint() string
33 String() string
34 CacheKey() string
35 Run() (any, error)
36}
37
38type queryRequest struct {
39 query querier
40 result chan queryResult
41}
42
43type queryResult struct {
44 value any
45 err error
46}
47
48type Prometheus struct {
49 name string
50 uri string
51 api v1.API
52 timeout time.Duration
53 concurrency int
54 cache *lru.ARCCache
55 locker *partitionLocker
56 rateLimiter ratelimit.Limiter
57 wg sync.WaitGroup
58 queries chan queryRequest
59}
60
61func NewPrometheus(name, uri string, timeout time.Duration, concurrency, cacheSize, rl int) *Prometheus {
62 client, err := api.NewClient(api.Config{Address: uri})
63 if err != nil {
64 // config validation should prevent this from ever happening
65 // panic so we don't need to return an error and it's easier to
66 // use this code in tests
67 panic(err)
68 }
69 cache, _ := lru.NewARC(cacheSize)
70
71 prom := Prometheus{
72 name: name,
73 uri: uri,
74 api: v1.NewAPI(client),
75 timeout: timeout,
76 cache: cache,
77 locker: newPartitionLocker((&sync.Mutex{})),
78 rateLimiter: ratelimit.New(rl),
79 concurrency: concurrency,
80 }
81 return &prom
82}
83
84func (prom *Prometheus) Close() {
85 log.Debug().Str("name", prom.name).Str("uri", prom.uri).Msg("Stopping query workers")
86 close(prom.queries)
87 prom.wg.Wait()
88}
89
90func (prom *Prometheus) StartWorkers() {
91 log.Debug().
92 Str("name", prom.name).
93 Str("uri", prom.uri).
94 Int("workers", prom.concurrency).
95 Msg("Starting query workers")
96
97 prom.queries = make(chan queryRequest, prom.concurrency*10)
98
99 for w := 1; w <= prom.concurrency; w++ {
100 prom.wg.Add(1)
101 go func() {
102 defer prom.wg.Done()
103 queryWorker(prom, prom.queries)
104 }()
105 }
106}
107
108func queryWorker(prom *Prometheus, queries chan queryRequest) {
109 for job := range queries {
110 job := job
111
112 cacheKey := job.query.CacheKey()
113 if cacheKey != "" {
114 if cached, ok := prom.cache.Get(cacheKey); ok {
115 job.result <- queryResult{value: cached}
116 prometheusCacheHitsTotal.WithLabelValues(prom.name, job.query.Endpoint()).Inc()
117 log.Debug().
118 Str("uri", prom.uri).
119 Str("query", job.query.String()).
120 Str("key", cacheKey).
121 Msg("Cache hit")
122 continue
123 }
124 }
125
126 prometheusQueriesTotal.WithLabelValues(prom.name, job.query.Endpoint()).Inc()
127 prometheusQueriesRunning.WithLabelValues(prom.name, job.query.Endpoint()).Inc()
128 prom.rateLimiter.Take()
129 start := time.Now()
130 result, err := job.query.Run()
131 dur := time.Since(start)
132 log.Debug().
133 Str("uri", prom.uri).
134 Str("query", job.query.String()).
135 Str("endpoint", job.query.Endpoint()).
136 Str("duration", output.HumanizeDuration(dur)).
137 Msg("Query completed")
138 prometheusQueriesRunning.WithLabelValues(prom.name, job.query.Endpoint()).Dec()
139 if err != nil {
140 prometheusQueryErrorsTotal.WithLabelValues(prom.name, job.query.Endpoint(), errReason(err)).Inc()
141 log.Error().
142 Err(err).
143 Str("uri", prom.uri).
144 Str("query", job.query.String()).
145 Msg("Query returned an error")
146 job.result <- queryResult{err: err}
147 continue
148 }
149
150 if cacheKey != "" {
151 prom.cache.Add(cacheKey, result)
152 }
153 prometheusCacheSize.WithLabelValues(prom.name).Set(float64(prom.cache.Len()))
154
155 job.result <- queryResult{value: result}
156 }
157}
158