package main import ( "encoding/json" "fmt" "io" "net/http" "os" "sync" "github.com/go-co-op/gocron/v2" "github.com/prometheus/client_golang/prometheus" "github.com/prometheus/client_golang/prometheus/promhttp" "github.com/urfave/cli/v2" log "github.com/sirupsen/logrus" ) type DomainState struct { // Hits per page (file) Pages map[string]int PagesMetrics map[string]prometheus.Gauge // Referrals Referrals map[string]int ReferralsMetrics map[string]prometheus.Gauge } type State struct { // Maps domain names to the sliding window Domains map[string]DomainState // Lock guarding the access to Domains Lock sync.Mutex } func increment(key string, dict *map[string]int) { v, found := (*dict)[key] if !found { (*dict)[key] = 1 } else { (*dict)[key] = v + 1 } } type Track struct { Domain string `json:"domain"` Referer string `json:"referer"` Path string `json:"path"` } func ingestHandler(registry *prometheus.Registry, state *State) http.HandlerFunc { return func(w http.ResponseWriter, req *http.Request) { if req.Method != http.MethodPost { log.Debugf("Unexpected %s request", req.Method) w.WriteHeader(http.StatusBadRequest) return } body, err := io.ReadAll(req.Body) if err != nil { log.Debugf("Failed to read body: %v", err) w.WriteHeader(http.StatusBadRequest) return } var track Track err = json.Unmarshal(body, &track) if err != nil { log.Debugf("Failed to unmarshal request body: %v", err) w.WriteHeader(http.StatusBadRequest) return } // Acquire lock state.Lock.Lock() defer state.Lock.Unlock() // Update everything log.Debugf("%v", state.Domains) log.Debugf("Got Domain: %s", track.Domain) domainState, found := state.Domains[track.Domain] if !found { log.Debugf("Domain %s not found. Creating state", track.Domain) domainState = DomainState{ Pages: make(map[string]int), PagesMetrics: make(map[string]prometheus.Gauge), Referrals: make(map[string]int), ReferralsMetrics: make(map[string]prometheus.Gauge), } state.Domains[track.Domain] = domainState } // Record the referrer referrer := track.Referer if referrer != "" { v, found := domainState.Referrals[referrer] if found { value := v + 1 domainState.Referrals[referrer] = value metric, _ := domainState.ReferralsMetrics[referrer] metric.Set(float64(value)) } else { metric := prometheus.NewGauge( prometheus.GaugeOpts{ Name: "referals", ConstLabels: prometheus.Labels{"domain": track.Domain, "referer": referrer}, }, ) metric.Set(1) domainState.Referrals[referrer] = 1 domainState.ReferralsMetrics[referrer] = metric log.Debugf("Registering gauge \"referals\" with referer=%s", referrer) registry.MustRegister(metric) } } // Record the file path := track.Path v, found := domainState.Pages[path] log.Debugf("v=%d, found=%v", v, found) if found { value := v + 1 domainState.Pages[path] = value metric, _ := domainState.PagesMetrics[path] metric.Set(float64(value)) } else { metric := prometheus.NewGauge( prometheus.GaugeOpts{ Name: "pages", ConstLabels: prometheus.Labels{"domain": track.Domain, "path": path}, }, ) metric.Set(1) domainState.Pages[path] = 1 domainState.PagesMetrics[path] = metric log.Debugf("Registering gauge \"path\" with path=%s", path) registry.MustRegister(metric) } state.Domains[track.Domain] = domainState w.WriteHeader(http.StatusAccepted) } } func run(ctx *cli.Context) error { host := ctx.String("host") port := ctx.String("port") if ctx.Bool("debug") { log.SetLevel(log.DebugLevel) } // Setup metrics registry := prometheus.NewRegistry() state := State{ Domains: make(map[string]DomainState), } // Setup the scheduler scheduler, err := gocron.NewScheduler() if err != nil { log.Error("Failed to create scheduler") return err } _, err = scheduler.NewJob( gocron.DailyJob(1, gocron.NewAtTimes(gocron.NewAtTime(0, 0, 0))), gocron.NewTask( func() { // Acquire the lock log.Info("Starting reset procedure...") state.Lock.Lock() defer state.Lock.Unlock() for domainName, domain := range state.Domains { log.Infof("Resetting counters for %s", domainName) // Reset pages for page := range domain.Pages { domain.Pages[page] = 0 domain.PagesMetrics[page].Set(0) } // Reset referals for referer := range domain.Referrals { domain.Referrals[referer] = 0 domain.ReferralsMetrics[referer].Set(0) } } }, ), ) if err != nil { log.Fatalf("Failed to create task") return err } scheduler.Start() // Set up the HTTP server http.Handle( "/metrics", promhttp.HandlerFor( registry, promhttp.HandlerOpts{ EnableOpenMetrics: true, }, ), ) http.Handle( "/track", ingestHandler(registry, &state), ) addr := fmt.Sprintf("%s:%s", host, port) log.Infof("Starting server at %s", addr) return http.ListenAndServe(addr, nil) } func main() { app := &cli.App{ Action: run, Flags: []cli.Flag{ &cli.StringFlag{ Name: "host", Usage: "Host to expose metrics on", Value: "127.0.0.1", EnvVars: []string{"HTTP_HOST"}, }, &cli.StringFlag{ Name: "port", Usage: "Port to expose metrics on", Value: "9999", EnvVars: []string{"HTTP_PORT"}, }, &cli.BoolFlag{ Name: "debug", Usage: "Enable verbose logging", Value: false, EnvVars: []string{"DEBUG"}, }, }, } if err := app.Run(os.Args); err != nil { log.Fatal(err) } }