reforger_crawler_main/telemetry/client.go

// Package telemetry reports metrics, events and heartbeats from a crawler
// service (indexer, scraper, ...) to the main API's POST /back/telemetry.
//
// It has no dependencies outside the standard library, so it can be imported
// by any service, or the file can simply be copied.
//
//	c := telemetry.New(telemetry.Config{
//		BaseURL:   "https://crawler.example.com",
//		SecretKey: os.Getenv("SECRET_KEY"),
//		Service:   "indexer",
//		Version:   "1.4.0",
//	})
//	c.Start(ctx)
//	defer c.Close(context.Background())
//
//	c.Counter("jobs_total", "result", "ok").Add(1)
//	c.Gauge("data_window_bytes").Set(float64(used))
//	c.Summary("download_seconds").Observe(took.Seconds())
//	c.Event(telemetry.Warn, "data limit reached", map[string]any{"limitGB": 50})
//	c.SetStatus(telemetry.Degraded, "rate limited by workshop")
//
// Recording never blocks on the network. A push that fails keeps the counter
// increments for the next attempt; a few failed batches of events and samples
// are buffered and the oldest are dropped first.
package telemetry

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"math"
	"net/http"
	"os"
	"sort"
	"strings"
	"sync"
	"time"
)

// Level of an event.
type Level string

const (
	Debug Level = "debug"
	Info  Level = "info"
	Warn  Level = "warn"
	Error Level = "error"
)

// Status is the self-reported health shown on the dashboard.
type Status string

const (
	OK       Status = "ok"
	Degraded Status = "degraded"
	Failing  Status = "error"
)

// Config configures a Client. BaseURL, SecretKey and Service are required.
type Config struct {
	BaseURL   string        // e.g. https://crawler.example.com (no trailing /back)
	SecretKey string        // the API's SECRET_KEY
	Service   string        // [a-z0-9_], e.g. "indexer"
	Instance  string        // defaults to the hostname
	Version   string        // shown on the dashboard
	Interval  time.Duration // push interval, default 15s (also the heartbeat)
	Meta      map[string]any
	// HTTPClient defaults to a client with a 10s timeout.
	HTTPClient *http.Client
	// OnError is called when a push fails (optional).
	OnError func(error)
}

// Client buffers metrics and events and pushes them periodically.
type Client struct {
	cfg     Config
	started time.Time

	mu        sync.Mutex
	counters  map[string]*Counter
	gauges    map[string]*Gauge
	summaries map[string]*Summary
	events    []event
	status    Status
	statusMsg string
	pending   [][]byte // failed batches waiting for a retry

	stop chan struct{}
	done chan struct{}
}

type metric struct {
	Name   string            `json:"name"`
	Type   string            `json:"type"`
	Value  float64           `json:"value,omitempty"`
	Count  int64             `json:"count,omitempty"`
	Sum    float64           `json:"sum,omitempty"`
	Min    float64           `json:"min,omitempty"`
	Max    float64           `json:"max,omitempty"`
	Labels map[string]string `json:"labels,omitempty"`
	Ts     int64             `json:"ts"`
}

type event struct {
	Level   Level          `json:"level"`
	Message string         `json:"message"`
	Ts      int64          `json:"ts"`
	Fields  map[string]any `json:"fields,omitempty"`
}

type batch struct {
	Service              string         `json:"service"`
	Instance             string         `json:"instance"`
	Version              string         `json:"version"`
	StartedAt            int64          `json:"startedAt"`
	HeartbeatIntervalSec int            `json:"heartbeatIntervalSec"`
	Status               Status         `json:"status"`
	StatusMessage        string         `json:"statusMessage"`
	Meta                 map[string]any `json:"meta,omitempty"`
	Metrics              []metric       `json:"metrics"`
	Events               []event        `json:"events"`
}

const (
	maxPending       = 20
	maxEvents        = 200
	maxMetricsPerReq = 1000
)

// New creates a client. Call Start to begin pushing.
func New(cfg Config) *Client {
	if cfg.Instance == "" {
		cfg.Instance, _ = os.Hostname()
	}
	if cfg.Interval <= 0 {
		cfg.Interval = 15 * time.Second
	}
	if cfg.HTTPClient == nil {
		cfg.HTTPClient = &http.Client{Timeout: 10 * time.Second}
	}
	cfg.BaseURL = strings.TrimRight(cfg.BaseURL, "/")
	return &Client{
		cfg:       cfg,
		started:   time.Now(),
		counters:  map[string]*Counter{},
		gauges:    map[string]*Gauge{},
		summaries: map[string]*Summary{},
		status:    OK,
	}
}

func key(name string, labels []string) string {
	return name + "\x00" + strings.Join(labels, "\x00")
}

func labelMap(kv []string) map[string]string {
	if len(kv) < 2 {
		return nil
	}
	m := make(map[string]string, len(kv)/2)
	for i := 0; i+1 < len(kv); i += 2 {
		m[kv[i]] = kv[i+1]
	}
	return m
}

// Counter returns a counter; labels are key/value pairs.
func (c *Client) Counter(name string, labels ...string) *Counter {
	k := key(name, labels)
	c.mu.Lock()
	defer c.mu.Unlock()
	ctr := c.counters[k]
	if ctr == nil {
		ctr = &Counter{name: name, labels: labelMap(labels)}
		c.counters[k] = ctr
	}
	return ctr
}

// Gauge returns a gauge; labels are key/value pairs.
func (c *Client) Gauge(name string, labels ...string) *Gauge {
	k := key(name, labels)
	c.mu.Lock()
	defer c.mu.Unlock()
	g := c.gauges[k]
	if g == nil {
		g = &Gauge{name: name, labels: labelMap(labels)}
		c.gauges[k] = g
	}
	return g
}

// Summary returns a summary (count/sum/min/max per push); labels are key/value pairs.
func (c *Client) Summary(name string, labels ...string) *Summary {
	k := key(name, labels)
	c.mu.Lock()
	defer c.mu.Unlock()
	s := c.summaries[k]
	if s == nil {
		s = &Summary{name: name, labels: labelMap(labels)}
		c.summaries[k] = s
	}
	return s
}

// Event queues a log event for the dashboard.
func (c *Client) Event(level Level, message string, fields map[string]any) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if len(c.events) >= maxEvents {
		c.events = c.events[1:]
	}
	c.events = append(c.events, event{Level: level, Message: message, Ts: time.Now().UnixMilli(), Fields: fields})
}

// SetStatus sets the self-reported health shown on the dashboard.
func (c *Client) SetStatus(s Status, message string) {
	c.mu.Lock()
	c.status, c.statusMsg = s, message
	c.mu.Unlock()
}

// Counter accumulates increments between pushes.
type Counter struct {
	mu     sync.Mutex
	name   string
	labels map[string]string
	delta  float64
}

func (m *Counter) Add(v float64) {
	if v <= 0 || math.IsNaN(v) || math.IsInf(v, 0) {
		return
	}
	m.mu.Lock()
	m.delta += v
	m.mu.Unlock()
}

func (m *Counter) Inc() { m.Add(1) }

// Gauge holds the latest value.
type Gauge struct {
	mu     sync.Mutex
	name   string
	labels map[string]string
	value  float64
	set    bool
}

func (m *Gauge) Set(v float64) {
	if math.IsNaN(v) || math.IsInf(v, 0) {
		return
	}
	m.mu.Lock()
	m.value, m.set = v, true
	m.mu.Unlock()
}

// Summary aggregates observations between pushes.
type Summary struct {
	mu       sync.Mutex
	name     string
	labels   map[string]string
	count    int64
	sum      float64
	min, max float64
}

func (m *Summary) Observe(v float64) {
	if math.IsNaN(v) || math.IsInf(v, 0) {
		return
	}
	m.mu.Lock()
	if m.count == 0 || v < m.min {
		m.min = v
	}
	if m.count == 0 || v > m.max {
		m.max = v
	}
	m.count++
	m.sum += v
	m.mu.Unlock()
}

// collect drains the current values into a batch payload.
func (c *Client) collect() []byte {
	now := time.Now()
	c.mu.Lock()
	b := batch{
		Service:              c.cfg.Service,
		Instance:             c.cfg.Instance,
		Version:              c.cfg.Version,
		StartedAt:            c.started.Unix(),
		HeartbeatIntervalSec: int(c.cfg.Interval.Seconds()),
		Status:               c.status,
		StatusMessage:        c.statusMsg,
		Meta:                 c.cfg.Meta,
		Events:               c.events,
	}
	c.events = nil
	counters := make([]*Counter, 0, len(c.counters))
	for _, m := range c.counters {
		counters = append(counters, m)
	}
	gauges := make([]*Gauge, 0, len(c.gauges))
	for _, m := range c.gauges {
		gauges = append(gauges, m)
	}
	summaries := make([]*Summary, 0, len(c.summaries))
	for _, m := range c.summaries {
		summaries = append(summaries, m)
	}
	c.mu.Unlock()

	ts := now.UnixMilli()
	for _, m := range counters {
		m.mu.Lock()
		if m.delta > 0 {
			b.Metrics = append(b.Metrics, metric{Name: m.name, Type: "counter", Value: m.delta, Labels: m.labels, Ts: ts})
			m.delta = 0
		}
		m.mu.Unlock()
	}
	for _, m := range gauges {
		m.mu.Lock()
		if m.set {
			b.Metrics = append(b.Metrics, metric{Name: m.name, Type: "gauge", Value: m.value, Labels: m.labels, Ts: ts})
		}
		m.mu.Unlock()
	}
	for _, m := range summaries {
		m.mu.Lock()
		if m.count > 0 {
			b.Metrics = append(b.Metrics, metric{Name: m.name, Type: "summary", Count: m.count, Sum: m.sum, Min: m.min, Max: m.max, Labels: m.labels, Ts: ts})
			m.count, m.sum, m.min, m.max = 0, 0, 0, 0
		}
		m.mu.Unlock()
	}
	sort.Slice(b.Metrics, func(i, j int) bool { return b.Metrics[i].Name < b.Metrics[j].Name })
	if len(b.Metrics) > maxMetricsPerReq {
		b.Metrics = b.Metrics[:maxMetricsPerReq]
	}
	if b.Metrics == nil {
		b.Metrics = []metric{}
	}
	if b.Events == nil {
		b.Events = []event{}
	}
	payload, _ := json.Marshal(b)
	return payload
}

func (c *Client) post(ctx context.Context, payload []byte) error {
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.cfg.BaseURL+"/back/telemetry", bytes.NewReader(payload))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("X-SECRET-KEY", c.cfg.SecretKey)
	resp, err := c.cfg.HTTPClient.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	body, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
	if resp.StatusCode/100 != 2 {
		return fmt.Errorf("telemetry push: %s: %s", resp.Status, strings.TrimSpace(string(body)))
	}
	return nil
}

// Flush pushes buffered batches and the current values now.
func (c *Client) Flush(ctx context.Context) error {
	payload := c.collect()
	c.mu.Lock()
	queue := append(c.pending, payload)
	c.pending = nil
	c.mu.Unlock()

	for i, p := range queue {
		if err := c.post(ctx, p); err != nil {
			c.mu.Lock()
			c.pending = append(queue[i:], c.pending...)
			if len(c.pending) > maxPending {
				c.pending = c.pending[len(c.pending)-maxPending:]
			}
			c.mu.Unlock()
			if c.cfg.OnError != nil {
				c.cfg.OnError(err)
			}
			return err
		}
	}
	return nil
}

// Start pushes in the background every Interval until ctx ends or Close is called.
func (c *Client) Start(ctx context.Context) {
	c.stop = make(chan struct{})
	c.done = make(chan struct{})
	go func() {
		defer close(c.done)
		t := time.NewTicker(c.cfg.Interval)
		defer t.Stop()
		_ = c.Flush(ctx)
		for {
			select {
			case <-ctx.Done():
				return
			case <-c.stop:
				return
			case <-t.C:
				_ = c.Flush(ctx)
			}
		}
	}()
}

// Close stops the background loop and makes a final push.
func (c *Client) Close(ctx context.Context) error {
	if c.stop != nil {
		close(c.stop)
		<-c.done
		c.stop = nil
	}
	return c.Flush(ctx)
}