Gomain.go -rw-r--r-- 7.8 KiB
1package main
2
3import (
4 "bytes"
5 "encoding/json"
6 "fmt"
7 "image/color"
8 "log"
9 "math"
10 "net/http"
11 "net/url"
12 "os"
13 "sort"
14 "strconv"
15 "strings"
16 "text/template"
17 "time"
18
19 "github.com/dustin/go-humanize"
20 "github.com/go-chi/chi"
21 "github.com/go-chi/chi/middleware"
22 "gonum.org/v1/plot"
23 "gonum.org/v1/plot/plotter"
24 "gonum.org/v1/plot/plotutil"
25 "gonum.org/v1/plot/vg"
26)
27
28var (
29 upstream string
30)
31
32type PromResponse struct {
33 Status string `json:"status"`
34 Data struct {
35 ResultType string `json:"resultType"`
36 Result []struct {
37 Metric map[string]string `json:"metric"`
38 Values [][]interface{} `json:"values"`
39 } `json:"result"`
40 } `json:"data"`
41}
42
43type Datapoint struct {
44 Time time.Time
45 Value float64
46}
47
48type PromResult struct {
49 Metric string
50 Values []Datapoint
51}
52
53func Query(q string, start time.Time, end time.Time, step int) ([]PromResult, error) {
54 body := url.Values{}
55 body.Set("query", q)
56 body.Set("start", fmt.Sprintf("%d", start.Unix()))
57 body.Set("end", fmt.Sprintf("%d", end.Unix()))
58 body.Set("step", fmt.Sprintf("%d", step))
59 resp, err := http.Post(fmt.Sprintf("%s/api/v1/query_range", upstream),
60 "application/x-www-form-urlencoded", strings.NewReader(body.Encode()))
61 if err != nil {
62 return nil, err
63 }
64 defer resp.Body.Close()
65
66 if resp.StatusCode != 200 {
67 return nil, fmt.Errorf("Received %d response from upstream", resp.StatusCode)
68 }
69
70 var data PromResponse
71 if err := json.NewDecoder(resp.Body).Decode(&data); err != nil {
72 return nil, err
73 }
74
75 if data.Data.ResultType != "matrix" {
76 return nil, fmt.Errorf("result type isn't of type matrix: %s",
77 data.Data.ResultType)
78 }
79
80 if len(data.Data.Result) == 0 {
81 return nil, fmt.Errorf("No data")
82 }
83
84 var results []PromResult
85 for _, res := range data.Data.Result {
86 r := PromResult{}
87 r.Metric = metricName(res.Metric)
88
89 var values []Datapoint
90 isValid := true
91 for _, vals := range res.Values {
92 timestamp := vals[0].(float64)
93 value := vals[1].(string)
94 fv, _ := strconv.ParseFloat(value, 64)
95 if math.IsNaN(fv) || math.IsInf(fv, 0) {
96 isValid = false
97 break
98 }
99 values = append(values, Datapoint{
100 time.Unix(int64(timestamp), 0),
101 fv,
102 })
103 }
104 r.Values = values
105
106 if isValid {
107 results = append(results, r)
108 }
109 }
110 return results, nil
111}
112
113func metricName(metric map[string]string) string {
114 if len(metric) == 0 {
115 return "{}"
116 }
117
118 out := ""
119 var inner []string
120 for key, value := range metric {
121 if key == "__name__" {
122 out = value
123 continue
124 }
125 inner = append(inner, fmt.Sprintf(`%s="%s"`, key, value))
126 }
127
128 if len(inner) == 0 {
129 return out
130 }
131
132 sort.Slice(inner, func(i, j int) bool {
133 return inner[i] < inner[j]
134 })
135
136 return out + "{" + strings.Join(inner, ",") + "}"
137}
138
139func handleLabel(p *plot.Plot, l *plotter.Line, label string, metric string) {
140 raw := metric[1 : len(metric)-1]
141 raw_tags := strings.Split(raw, ",")
142 tags := make(map[string]string)
143 for _, v := range raw_tags {
144 tag := strings.Split(v, "=")
145 if len(tag) != 2 {
146 log.Printf("Expected tag format: 'name=value'!")
147 continue
148 }
149 if len(tag[1]) > 2 && tag[1][0] == '"' && tag[1][len(tag[1])-1] == '"' {
150 tags[tag[0]] = tag[1][1 : len(tag[1])-1]
151 }
152 }
153 tmpl, err := template.New("label").Parse(label)
154 if err != nil {
155 log.Printf("Failed to parse label template: %v", err)
156 } else {
157 var label_out bytes.Buffer
158 tmpl.Execute(&label_out, tags)
159 p.Legend.Add(label_out.String(), l)
160 }
161}
162
163func registerExtension(router chi.Router, extension string, mime string) {
164 router.Get("/chart."+extension, func(w http.ResponseWriter, r *http.Request) {
165 args := r.URL.Query()
166 var query string
167 if q, ok := args["query"]; !ok {
168 w.WriteHeader(400)
169 w.Write([]byte("Expected ?query=... parameter"))
170 return
171 } else {
172 query = q[0]
173 }
174
175 start := time.Now().Add(-24 * 60 * time.Minute)
176 end := time.Now()
177 if s, ok := args["since"]; ok {
178 d, _ := time.ParseDuration(s[0])
179 start = time.Now().Add(-d)
180 }
181 if u, ok := args["until"]; ok {
182 d, _ := time.ParseDuration(u[0])
183 end = time.Now().Add(-d)
184 }
185
186 width := 12 * vg.Inch
187 height := 6 * vg.Inch
188 if ws, ok := args["width"]; ok {
189 w, _ := strconv.ParseFloat(ws[0], 32)
190 width = vg.Length(w) * vg.Inch
191 }
192 if hs, ok := args["height"]; ok {
193 h, _ := strconv.ParseFloat(hs[0], 32)
194 height = vg.Length(h) * vg.Inch
195 }
196
197 // Label template
198 var label string
199 if l, ok := args["label"]; ok {
200 label = l[0]
201 }
202
203 // Set step so that there's approximately 25 data points per inch
204 step := int(end.Sub(start).Seconds() / (25 * float64(width/vg.Inch)))
205 if s, ok := args["step"]; ok {
206 d, _ := strconv.ParseInt(s[0], 10, 32)
207 step = int(d)
208 }
209 _, stacked := args["stacked"]
210
211 data, err := Query(query, start, end, step)
212 if err != nil {
213 w.WriteHeader(400)
214 w.Write([]byte(fmt.Sprintf("%v", err)))
215 return
216 }
217
218 p, err := plot.New()
219 if err != nil {
220 panic(err)
221 }
222 if t, ok := args["title"]; ok {
223 p.Title.Text = t[0]
224 }
225 p.X.Label.Text = "Time"
226 p.X.Tick.Marker = dateTicks{start, end}
227 if ms, ok := args["max"]; ok {
228 m, _ := strconv.ParseFloat(ms[0], 64)
229 p.Y.Max = m
230 }
231
232 p.Y.Tick.Marker = humanTicks{}
233 if ms, ok := args["min"]; ok {
234 m, _ := strconv.ParseFloat(ms[0], 64)
235 p.Y.Min = m
236 }
237 p.Legend.Top = true
238
239 sums := make([]float64, len(data[0].Values))
240
241 plotters := make([]plot.Plotter, len(data))
242 var nextColor int
243 colors := plotutil.SoftColors
244 for i, res := range data {
245 var points plotter.XYs
246 for j, d := range res.Values {
247 value := d.Value
248 if stacked {
249 value += sums[j]
250 }
251 points = append(points, plotter.XY{
252 float64(d.Time.Unix()),
253 value,
254 })
255 sums[j] += d.Value
256 }
257
258 l, _, err := plotter.NewLinePoints(points)
259 if err != nil {
260 w.WriteHeader(400)
261 w.Write([]byte(fmt.Sprintf("%v", err)))
262 return
263 }
264 if stacked {
265 l.FillColor = colors[nextColor]
266 if i != len(data)-1 {
267 l.Color = color.RGBA{0, 0, 0, 0}
268 }
269 } else {
270 l.Color = colors[nextColor]
271 }
272 nextColor += 1
273 if nextColor >= len(colors) {
274 nextColor = 0
275 }
276 plotters[i] = l
277 if label != "" && len(res.Metric) > 2 && res.Metric[0] == '{' && res.Metric[len(res.Metric)-1] == '}' {
278 handleLabel(p, l, label, res.Metric)
279 } else {
280 p.Legend.Add(res.Metric, l)
281 }
282 }
283 for i := len(plotters) - 1; i >= 0; i-- {
284 p.Add(plotters[i])
285 }
286
287 writer, err := p.WriterTo(width, height, extension)
288 if err != nil {
289 w.WriteHeader(400)
290 w.Write([]byte(fmt.Sprintf("%v", err)))
291 return
292 }
293
294 w.Header().Add("Content-Type", mime)
295 writer.WriteTo(w)
296 })
297
298}
299
300func main() {
301 plotutil.DefaultDashes = [][]vg.Length{{}}
302
303 if len(os.Args) < 2 {
304 fmt.Printf("Usage: %s server\n", os.Args[0])
305 os.Exit(1)
306 }
307 upstream = os.Args[1]
308 router := chi.NewRouter()
309
310 router.Use(middleware.RealIP)
311 router.Use(middleware.Logger)
312
313 registerExtension(router, "svg", "image/svg+xml")
314 registerExtension(router, "png", "image/png")
315
316 addr := ":8142"
317 if len(os.Args) > 2 {
318 addr = os.Args[2]
319 }
320 log.Printf("Listening on %s", addr)
321 http.ListenAndServe(addr, router)
322}
323
324type dateTicks struct {
325 Start time.Time
326 End time.Time
327}
328
329// Ticks computes the default tick marks, but inserts commas
330// into the labels for the major tick marks.
331func (dt dateTicks) Ticks(min, max float64) []plot.Tick {
332 fmt := "15:04:05"
333 if dt.End.Sub(dt.Start).Hours() >= 24 {
334 fmt = "Jan 2 15:04:05"
335 }
336
337 tks := plot.DefaultTicks{}.Ticks(min, max)
338 for i, t := range tks {
339 if t.Label == "" { // Skip minor ticks, they are fine.
340 continue
341 }
342 d, _ := strconv.ParseFloat(t.Label, 64)
343 tm := time.Unix(int64(d), 0)
344 tks[i].Label = tm.Format(fmt)
345 }
346 return tks
347}
348
349type humanTicks struct{}
350
351func (ht humanTicks) Ticks(min, max float64) []plot.Tick {
352 tks := plot.DefaultTicks{}.Ticks(min, max)
353 for i, t := range tks {
354 if t.Label == "" { // Skip minor ticks, they are fine.
355 continue
356 }
357 d, _ := strconv.ParseFloat(t.Label, 64)
358 tks[i].Label = humanize.SI(d, "")
359 }
360 return tks
361}