reforger_crawler_main/metrics/traffic.go

package metrics

import (
	"crypto/rand"
	"crypto/sha256"
	"encoding/binary"
	"log"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"gitea.tbdevent.eu/TBD/reforger_crawler_main/initializers"
	"gitea.tbdevent.eu/TBD/reforger_crawler_main/models"
	"github.com/gin-gonic/gin"
	"gorm.io/gorm/clause"
)

// Middleware records latency and traffic for every request except static admin
// UI files and health checks.
func Middleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		path := c.Request.URL.Path
		if c.Request.Method == "OPTIONS" || path == "/health" || strings.HasPrefix(path, "/admin/ui") {
			c.Next()
			return
		}
		start := time.Now()
		c.Next()
		ms := float64(time.Since(start).Microseconds()) / 1000

		route := c.FullPath()
		if route == "" {
			route = "(unmatched)"
		}
		status := c.Writer.Status()
		group := routeGroup(path)
		client := "direct"
		if c.GetHeader("Origin") != "" || c.GetHeader("Referer") != "" {
			client = "web"
		}
		HTTPDuration.Observe(ms, c.Request.Method, route, strconv.Itoa(status/100)+"xx", group, client)
		if size := c.Writer.Size(); size > 0 {
			HTTPResponseBytes.Add(float64(size), group)
		}

		if group == "v1" && status < 500 {
			visitors.add(c.ClientIP(), c.Request.UserAgent())
		}
		if status == 200 {
			switch route {
			case "/v1/addon/:id":
				Top.Add("addon", c.Param("id"))
			case "/v1/getPossible":
				Top.Add("search", normalizeSearch(c.Query("q")))
			}
		}
	}
}

func routeGroup(path string) string {
	for _, g := range []string{"v1", "back", "admin", "metrics"} {
		if path == "/"+g || strings.HasPrefix(path, "/"+g+"/") {
			return g
		}
	}
	return "other"
}

func normalizeSearch(q string) string {
	q = strings.Join(strings.Fields(strings.ToLower(q)), " ")
	if len(q) > 64 {
		q = q[:64]
	}
	return q
}

// visitorSet counts distinct (IP, user agent) pairs per UTC day. Only salted
// 64-bit hashes are kept and the salt never leaves memory.
type visitorSet struct {
	mu     sync.Mutex
	day    string
	salt   []byte
	humans map[uint64]struct{}
	bots   map[uint64]struct{}
}

var visitors = &visitorSet{}

const maxVisitors = 2_000_000

func (v *visitorSet) rotate(day string) {
	v.day = day
	v.salt = make([]byte, 32)
	_, _ = rand.Read(v.salt)
	v.humans = make(map[uint64]struct{})
	v.bots = make(map[uint64]struct{})
}

func isBot(ua string) bool {
	ua = strings.ToLower(ua)
	if ua == "" {
		return true
	}
	for _, s := range []string{"bot", "spider", "crawl", "curl", "wget", "python", "go-http", "headless", "preview", "monitor"} {
		if strings.Contains(ua, s) {
			return true
		}
	}
	return false
}

func (v *visitorSet) add(ip, ua string) {
	day := time.Now().UTC().Format("2006-01-02")
	v.mu.Lock()
	defer v.mu.Unlock()
	if v.day != day {
		v.rotate(day)
	}
	h := sha256.New()
	h.Write(v.salt)
	h.Write([]byte(ip))
	h.Write([]byte{0})
	h.Write([]byte(ua))
	key := binary.BigEndian.Uint64(h.Sum(nil)[:8])
	set := v.humans
	if isBot(ua) {
		set = v.bots
	}
	if len(set) < maxVisitors {
		set[key] = struct{}{}
	}
}

// Visitors returns today's approximate unique human and bot visitors. The set
// lives in memory, so after a restart the day's highest recorded value is used
// until the new set catches up.
func Visitors() (humans, bots int) {
	now := time.Now().UTC()
	day := now.Format("2006-01-02")
	visitors.mu.Lock()
	if visitors.day == day {
		humans, bots = len(visitors.humans), len(visitors.bots)
	}
	visitors.mu.Unlock()

	dayStart := now.Truncate(24 * time.Hour).Unix()
	if dayStart < startedAt.Unix() {
		if t, err := Totals(Selector{Service: MainService, Name: "visitors_unique", GroupBy: []string{"kind"}}, dayStart, now.Unix()); err == nil {
			for _, g := range t {
				if g.Max == nil {
					continue
				}
				switch g.Labels["kind"] {
				case "human":
					humans = max(humans, int(*g.Max))
				case "bot":
					bots = max(bots, int(*g.Max))
				}
			}
		}
	}
	return humans, bots
}

// liveVisitors is the in-memory count only, used when recording the gauge.
func liveVisitors() (humans, bots int) {
	day := time.Now().UTC().Format("2006-01-02")
	visitors.mu.Lock()
	defer visitors.mu.Unlock()
	if visitors.day != day {
		return 0, 0
	}
	return len(visitors.humans), len(visitors.bots)
}

// topCounter counts requested items per UTC day and persists the most popular.
type topCounter struct {
	mu      sync.Mutex
	day     string
	counts  map[string]map[string]int64
	pending []topSnapshot // previous days not yet written
}

type topSnapshot struct {
	day    string
	counts map[string]map[string]int64
}

// Top tracks the most looked-up addons and search queries.
var Top = &topCounter{counts: map[string]map[string]int64{}}

const (
	maxTopKeys    = 50_000
	topPersistMax = 500
)

func (t *topCounter) Add(kind, item string) {
	if item == "" {
		return
	}
	if len(item) > 128 {
		item = item[:128]
	}
	day := time.Now().UTC().Format("2006-01-02")
	t.mu.Lock()
	defer t.mu.Unlock()
	if t.day != day {
		if t.day != "" {
			t.pending = append(t.pending, topSnapshot{t.day, t.counts})
		}
		t.day = day
		t.counts = map[string]map[string]int64{}
	}
	m := t.counts[kind]
	if m == nil {
		m = map[string]int64{}
		t.counts[kind] = m
	}
	if _, ok := m[item]; !ok && len(m) >= maxTopKeys {
		for k, v := range m {
			if v <= 1 {
				delete(m, k)
			}
		}
		if len(m) >= maxTopKeys {
			return
		}
	}
	m[item]++
}

// Today returns today's in-memory counts for a kind.
func (t *topCounter) Today(kind string) (string, map[string]int64) {
	t.mu.Lock()
	defer t.mu.Unlock()
	out := make(map[string]int64, len(t.counts[kind]))
	for k, v := range t.counts[kind] {
		out[k] = v
	}
	return t.day, out
}

// seed loads today's persisted counts so a restart does not reset them.
func (t *topCounter) seed() {
	day := time.Now().UTC().Format("2006-01-02")
	var rows []models.DailyTopItem
	if err := initializers.DB.Where("day = ?", day).Find(&rows).Error; err != nil {
		log.Printf("metrics: failed to load top items: %v", err)
		return
	}
	t.mu.Lock()
	defer t.mu.Unlock()
	if t.day != day {
		t.day = day
		t.counts = map[string]map[string]int64{}
	}
	for _, r := range rows {
		if t.counts[r.Kind] == nil {
			t.counts[r.Kind] = map[string]int64{}
		}
		t.counts[r.Kind][r.Item] += r.Count
	}
}

func topRows(snap topSnapshot) []models.DailyTopItem {
	var rows []models.DailyTopItem
	for kind, m := range snap.counts {
		type kv struct {
			k string
			v int64
		}
		list := make([]kv, 0, len(m))
		for k, v := range m {
			list = append(list, kv{k, v})
		}
		sort.Slice(list, func(i, j int) bool { return list[i].v > list[j].v })
		if len(list) > topPersistMax {
			list = list[:topPersistMax]
		}
		for _, e := range list {
			rows = append(rows, models.DailyTopItem{Day: snap.day, Kind: kind, Item: e.k, Count: e.v})
		}
	}
	return rows
}

func (t *topCounter) flush() error {
	t.mu.Lock()
	snaps := append([]topSnapshot{}, t.pending...)
	t.pending = nil
	cur := topSnapshot{day: t.day, counts: map[string]map[string]int64{}}
	for kind, m := range t.counts {
		cp := make(map[string]int64, len(m))
		for k, v := range m {
			cp[k] = v
		}
		cur.counts[kind] = cp
	}
	t.mu.Unlock()
	if cur.day != "" {
		snaps = append(snaps, cur)
	}

	for i, snap := range snaps {
		rows := topRows(snap)
		if len(rows) == 0 {
			continue
		}
		err := initializers.DB.Clauses(clause.OnConflict{
			Columns:   []clause.Column{{Name: "day"}, {Name: "kind"}, {Name: "item"}},
			DoUpdates: clause.AssignmentColumns([]string{"count"}),
		}).CreateInBatches(&rows, 200).Error
		if err != nil {
			// keep finished days for the next attempt; today is re-read from memory
			t.mu.Lock()
			for _, old := range snaps[i:] {
				if old.day != cur.day {
					t.pending = append(t.pending, old)
				}
			}
			t.mu.Unlock()
			return err
		}
	}
	return nil
}