badger/main.go

package main

import (
	"context"
	"errors"
	"fmt"
	"log"
	"net/http"
	"os"
	"strconv"
	"time"

	svg "github.com/ajstarks/svgo"
	"github.com/gin-gonic/gin"
	"github.com/prometheus/client_golang/api"
	v1 "github.com/prometheus/client_golang/api/prometheus/v1"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/prometheus/common/model"
	"github.com/sirupsen/logrus"
)

const (
	radius             = 0
	fontSize           = 12
	textMarginX        = 6.0
	textMarginY        = 14
	badgeHeight        = 20
	textStyle          = "text-anchor:start;font-size:12px;fill:white;font-family:-apple-system,BlinkMacSystemFont,Segoe UI,Helvetica,Arial,sans-serif,Apple Color Emoji,Segoe UI Emoji;"
	badgeExpireSeconds = 60
	labelWidthPortion  = 0.65
	ColorSchemeGreen   = "fill-opacity:1.00; fill:rgb(49, 197, 83);"
	ColorSchemeOrange  = "fill-opacity:1.00; fill:rgb(255, 87, 34);"
	ColorSchemeBlue    = "fill-opacity:1.00; fill:rgb(32, 99, 155);"
)

type ColorScheme string

var (
	ErrNoSuchBadgeFound = errors.New("no such badge can be found, please make sure you have provided the correct information")
	queriesMap          = map[string]queryParam{
		"mel":                           {Color: ColorSchemeGreen, Query: "air_index_iaqi_pm25_v{city_name=\"Melbourne CBD\"}", ValueField: "city_name"},
		"chengdu":                       {Color: ColorSchemeGreen, Query: "air_index_iaqi_pm25_v{city_name=\"Chengdu (成都)\"}", ValueField: "city_name"},
		"dockerhub":                     {Color: ColorSchemeGreen, Query: "docker_hub_pull_count{name=\"aria2-ui\"}", ValueField: "name"},
		"github":                        {Color: ColorSchemeGreen, Query: "github_repository_stars{name=\"aria2-ariang-docker\"}", ValueField: "name"},
		"badge-requests-count":          {Color: ColorSchemeBlue, Query: "ceil(sum(increase(badge_requested_total[500d])))", Label: "Badge Req Count"},
		"junv-github-profile":           {Color: ColorSchemeBlue, Query: "ceil(sum(increase(badge_requested_total{badge=\"junv-github-profile\"}[500d])))", Label: "Page Views"},
		"aria2-ariang-docker":           {Color: ColorSchemeBlue, Query: "ceil(sum(increase(badge_requested_total{badge=\"aria2-ariang-docker\"}[500d])))", Label: "Page Views"},
		"aria2-ariang-x-docker-compose": {Color: ColorSchemeBlue, Query: "ceil(sum(increase(badge_requested_total{badge=\"aria2-ariang-x-docker-compose\"}[500d])))", Label: "Page Views"},
		"work-in-australia":             {Color: ColorSchemeBlue, Query: "ceil(sum(increase(badge_requested_total{badge=\"work-in-australia\"}[500d])))", Label: "Page Views"},
		"aria-ui-dockerhub":             {Color: ColorSchemeBlue, Query: "ceil(sum(increase(badge_requested_total{badge=\"aria-ui-dockerhub\"}[500d])))", Label: "Page Views"},
	}

	badgeRequested = promauto.NewCounterVec(prometheus.CounterOpts{
		Name: "badge_requested_total",
		Help: "The total request number of badges",
	}, []string{"badge"})

	badgeGenerationTime = promauto.NewHistogram(prometheus.HistogramOpts{
		Name: "badge_generation_time",
		Help: "The time consumed for badge generation",
	})
)

type Metric struct {
	Label string
	Value float64
	Time  time.Time
}

type queryParam struct {
	Color      ColorScheme
	Query      string
	ValueField string // nullable
	Label      string // when ValueField is null, then the program will use this field as the label
}

func main() {
	client, err := api.NewClient(api.Config{
		Address: "https://prometheus-api.home.toozhao.com",
	})
	if err != nil {
		fmt.Printf("Error creating client: %v\n", err)
		os.Exit(1)
	}

	r := gin.Default()

	r.GET("/", func(c *gin.Context) {
		c.JSON(200, gin.H{"msg": "Wanna create some badges for your website or Github repo from various data sources? It's coming soon."})
	})

	r.GET("/ping", func(c *gin.Context) {
		c.JSON(200, gin.H{"msg": "Hello!"})
	})

	r.GET("/svg/:name", func(c *gin.Context) {
		if err := generateSVG(client, c); err != nil {
			logrus.Error(err.Error())
			if err == ErrNoSuchBadgeFound {
				c.JSON(http.StatusNotFound, gin.H{"error": err.Error()})
				return
			}

			c.JSON(http.StatusInternalServerError, gin.H{"error": "please try it again"})
		}
	})

	r.StaticFile("/favicon.ico", "static/favicon.ico")
	r.StaticFile("/robots.txt", "static/robots.txt")

	go func() {
		http.Handle("/metrics", promhttp.Handler())
		if err := http.ListenAndServe(":9090", nil); err != nil {
			logrus.Error(err)
			os.Exit(1)
		}

	}()

	log.Fatal(r.Run())
}

func generateSVG(client api.Client, c *gin.Context) error {
	startTime := time.Now()

	queryName := c.Param("name")
	query, ok := queriesMap[queryName]
	if !ok {
		return ErrNoSuchBadgeFound
	}
	metric := queryMetric(client, query)

	metricValue := int(metric.Value)
	label := metric.Label

	valueRectWidth := 40
	labelContainerWidth := 120

	if len(label) <= 10 {
		labelContainerWidth = 74
	}

	if metricValue > 10000 && metricValue < 1000000 {
		valueRectWidth = 50
	} else if metricValue >= 1000000 {
		valueRectWidth = 66
	}
	badgeWidth := labelContainerWidth + valueRectWidth

	w := c.Writer
	w.Header().Set("Content-Type", "image/svg+xml")
	w.Header().Set("Cache-Control", fmt.Sprintf("max-age=%d, s-maxage=%d", badgeExpireSeconds, badgeExpireSeconds))
	canvas := svg.New(w)
	canvas.Start(badgeWidth, badgeHeight)
	canvas.Rect(0, 0, badgeWidth, badgeHeight, "fill-opacity:1.00; fill:rgb(90,90,90); padding: 2px 5px;")

	//labelContainerWidth := len(metric.Label)*fontSize + textMarginX

	canvas.Rect(0, 0, int(labelContainerWidth), badgeHeight, string(query.Color))

	//when there is metric label then use default label
	if len(query.ValueField) == 0 {
		label = query.Label
	}
	canvas.Text(textMarginX, textMarginY, label, textStyle)

	valueText := strconv.Itoa(metricValue)
	canvas.Text(labelContainerWidth+textMarginX, textMarginY, valueText, textStyle)
	canvas.End()

	badgeGenerationTime.Observe(float64(time.Since(startTime).Microseconds()))
	badgeRequested.With(map[string]string{"badge": queryName}).Inc()

	return nil
}

func queryMetric(client api.Client, query queryParam) *Metric {
	v1api := v1.NewAPI(client)
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	result, warnings, err := v1api.Query(ctx, query.Query, time.Now())
	if err != nil {
		fmt.Printf("Error querying Prometheus: %v\n", err)
		os.Exit(1)
	}
	if len(warnings) > 0 {
		fmt.Printf("Warnings: %v\n", warnings)
	}
	var metric Metric
	for _, v := range result.(model.Vector) {
		metric = Metric{
			Label: string(v.Metric[model.LabelName(query.ValueField)]),
			Value: float64(v.Value),
			Time:  v.Timestamp.Time(),
		}
	}
	return &metric
}