sift/internal/metrics/metrics.go

// Package metrics is sift's Prometheus registry and the loopback-scrape predicate.
package metrics

import (
	"net"
	"net/http"
	"net/netip"
	"strconv"
	"strings"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

var durationBuckets = []float64{0.1, 0.5, 1, 2, 5, 10, 15, 20, 30, 45}

var fallbackOK = map[string]struct{}{
	"queue_full": {}, "queue_wait": {}, "canceled": {},
	"embed": {}, "dixie": {}, "timeout": {}, "empty": {}, "unset": {},
}

// Snapshot is process state gathered on each Prometheus scrape.
type Snapshot struct {
	Blocks, Scrapes, SelfIPFailures uint64
	DigestCalls                     uint64
	DigestFallbacks                 map[string]uint64
	DixieInflight, DixieWaiters     int
	DigestHit, DigestMiss           uint64
	EmbedHit, EmbedMiss             uint64
	SearchHit, SearchMiss           uint64
	SearchQueries                   map[string]map[string]uint64 // caller → source
	SearchCredits                   map[string]map[string]uint64
	SearchDaily                     map[string]uint64
	SearchCapped                    map[string]uint64
	Restricted                      map[string]struct{}
}

// Source is implemented by api.Server.
type Source interface {
	Snapshot() Snapshot
}

// pdfOutcomes are the outcomes ObservePDF counts directly; anything else is
// folded into "other" so a raw outcome string never becomes a label.
var pdfOutcomes = map[string]struct{}{
	"ok": {}, "no_text": {}, "encrypted": {}, "unreadable": {},
	"timeout": {}, "too_large": {}, "busy": {},
}

// Set is a private registry of sift_* series.
type Set struct {
	Reg                  *prometheus.Registry
	Requests             *prometheus.CounterVec
	Duration             *prometheus.HistogramVec
	PDFTotal             *prometheus.CounterVec
	PDFSeconds           prometheus.Histogram
	PDFPages             prometheus.Counter
	PDFWait              prometheus.Histogram
	PDFCacheHits         prometheus.Counter
	BatchJobsTotal       *prometheus.CounterVec
	BatchTargetsTotal    *prometheus.CounterVec
	BatchQueueDepth      prometheus.Gauge
	BatchTargetsInflight prometheus.Gauge
	BatchJobSeconds      prometheus.Histogram
	BatchTargetSeconds   prometheus.Histogram
	BatchResultBytes     prometheus.Histogram
	handler              http.Handler
}

func New(src Source) *Set {
	reg := prometheus.NewRegistry()
	requests := prometheus.NewCounterVec(prometheus.CounterOpts{
		Name: "sift_requests_total",
		Help: "HTTP API requests by endpoint and status.",
	}, []string{"endpoint", "code"})
	duration := prometheus.NewHistogramVec(prometheus.HistogramOpts{
		Name:    "sift_request_duration_seconds",
		Help:    "API wall time (fetch/search) or digest time (digest).",
		Buckets: durationBuckets,
	}, []string{"op"})
	pdfTotal := prometheus.NewCounterVec(prometheus.CounterOpts{
		Name: "sift_pdf_extract_total",
		Help: "PDF extraction attempts by outcome.",
	}, []string{"outcome"})
	pdfSeconds := prometheus.NewHistogram(prometheus.HistogramOpts{
		Name:    "sift_pdf_extract_seconds",
		Help:    "PDF extraction wall time.",
		Buckets: durationBuckets,
	})
	pdfPages := prometheus.NewCounter(prometheus.CounterOpts{
		Name: "sift_pdf_pages_total",
		Help: "PDF pages extracted.",
	})
	pdfWait := prometheus.NewHistogram(prometheus.HistogramOpts{
		Name:    "sift_pdf_pool_wait_seconds",
		Help:    "Time a PDF job waited in the pool before its extraction started.",
		Buckets: durationBuckets,
	})
	pdfCacheHits := prometheus.NewCounter(prometheus.CounterOpts{
		Name: "sift_pdf_cache_hits_total",
		Help: "PDF answers served from the document cache.",
	})
	batchJobs := prometheus.NewCounterVec(prometheus.CounterOpts{
		Name: "sift_batch_jobs_total",
		Help: "Batch jobs by outcome.",
	}, []string{"outcome"})
	batchTargets := prometheus.NewCounterVec(prometheus.CounterOpts{
		Name: "sift_batch_targets_total",
		Help: "Batch targets by outcome.",
	}, []string{"outcome"})
	batchQueueDepth := prometheus.NewGauge(prometheus.GaugeOpts{
		Name: "sift_batch_queue_depth",
		Help: "Live batch jobs (queued+running).",
	})
	batchTargetsInflight := prometheus.NewGauge(prometheus.GaugeOpts{
		Name: "sift_batch_targets_inflight",
		Help: "Batch targets currently being fetched.",
	})
	batchJobSeconds := prometheus.NewHistogram(prometheus.HistogramOpts{
		Name:    "sift_batch_job_seconds",
		Help:    "Batch wall time from submit to terminal status.",
		Buckets: durationBuckets,
	})
	batchTargetSeconds := prometheus.NewHistogram(prometheus.HistogramOpts{
		Name:    "sift_batch_target_seconds",
		Help:    "Batch target fetch wall time.",
		Buckets: durationBuckets,
	})
	batchResultBytes := prometheus.NewHistogram(prometheus.HistogramOpts{
		Name:    "sift_batch_result_bytes",
		Help:    "Stored bytes per batch target result.",
		Buckets: batchResultBytesBuckets,
	})
	reg.MustRegister(requests, duration, pdfTotal, pdfSeconds, pdfPages, pdfWait, pdfCacheHits,
		batchJobs, batchTargets, batchQueueDepth, batchTargetsInflight,
		batchJobSeconds, batchTargetSeconds, batchResultBytes, &collector{src: src})
	return &Set{
		Reg:                  reg,
		Requests:             requests,
		Duration:             duration,
		PDFTotal:             pdfTotal,
		PDFSeconds:           pdfSeconds,
		PDFPages:             pdfPages,
		PDFWait:              pdfWait,
		PDFCacheHits:         pdfCacheHits,
		BatchJobsTotal:       batchJobs,
		BatchTargetsTotal:    batchTargets,
		BatchQueueDepth:      batchQueueDepth,
		BatchTargetsInflight: batchTargetsInflight,
		BatchJobSeconds:      batchJobSeconds,
		BatchTargetSeconds:   batchTargetSeconds,
		BatchResultBytes:     batchResultBytes,
		handler:              promhttp.HandlerFor(reg, promhttp.HandlerOpts{EnableOpenMetrics: false}),
	}
}

func (s *Set) Observe(endpoint string, dur time.Duration, code int) {
	if s == nil {
		return
	}
	s.Requests.WithLabelValues(endpoint, strconv.Itoa(code)).Inc()
	switch endpoint {
	case "/v1/fetch", "/v2/scrape":
		s.Duration.WithLabelValues("fetch").Observe(dur.Seconds())
	case "/v1/search", "/v2/search":
		s.Duration.WithLabelValues("search").Observe(dur.Seconds())
	}
}

func (s *Set) ObserveDigest(d time.Duration) {
	if s == nil {
		return
	}
	s.Duration.WithLabelValues("digest").Observe(d.Seconds())
}

// ObservePDF records one PDF extraction attempt: the outcome (folded to "other"
// when it is not a known value, so a raw string never becomes a label), the wall
// time, and the pages extracted.
func (s *Set) ObservePDF(outcome string, d time.Duration, pages int) {
	if s == nil {
		return
	}
	if _, ok := pdfOutcomes[outcome]; !ok {
		outcome = "other"
	}
	s.PDFTotal.WithLabelValues(outcome).Inc()
	s.PDFSeconds.Observe(d.Seconds())
	s.PDFPages.Add(float64(pages))
}

// ObservePDFWait records the time a PDF job spent in the pool before its
// extraction started.
func (s *Set) ObservePDFWait(d time.Duration) {
	if s == nil {
		return
	}
	s.PDFWait.Observe(d.Seconds())
}

// PDFCacheHit records one PDF answer served from the document cache.
func (s *Set) PDFCacheHit() {
	if s == nil {
		return
	}
	s.PDFCacheHits.Inc()
}

func (s *Set) Handler() http.Handler {
	if s == nil {
		return http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
			http.Error(w, "metrics unavailable", http.StatusServiceUnavailable)
		})
	}
	return s.handler
}

// LoopbackOpen is curl http://127.0.0.1:8890/metrics. tailscale serve on :8443
// is not open: it sets forwarded/Tailscale headers and a non-loopback Host.
func LoopbackOpen(r *http.Request) bool {
	if r == nil {
		return false
	}
	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		host = r.RemoteAddr
	}
	addr, err := netip.ParseAddr(host)
	if err != nil || !addr.IsLoopback() {
		return false
	}
	for _, h := range []string{
		"X-Forwarded-For", "X-Forwarded-Proto", "X-Forwarded-Host",
		"Tailscale-User-Login", "Tailscale-Headers-Info",
	} {
		if r.Header.Get(h) != "" {
			return false
		}
	}
	h := r.Host
	if hh, _, err := net.SplitHostPort(h); err == nil {
		h = hh
	}
	if strings.EqualFold(h, "localhost") {
		return true
	}
	ip, err := netip.ParseAddr(h)
	return err == nil && ip.IsLoopback()
}

func CoerceFallback(reason string) string {
	if _, ok := fallbackOK[reason]; ok {
		return reason
	}
	return "other"
}

func skipRestrictedSource(source string) bool {
	return strings.HasSuffix(source, ":restricted")
}

type collector struct {
	src Source
}

func (c *collector) Describe(ch chan<- *prometheus.Desc) {
	ch <- dBlocks
	ch <- dSelfIP
	ch <- dDigestCalls
	ch <- dDigestFB
	ch <- dDixieIn
	ch <- dDixieWait
	ch <- dCache
	ch <- dQTotal
	ch <- dQCaller
	ch <- dCTotal
	ch <- dCCaller
	ch <- dDaily
	ch <- dCapped
}

func (c *collector) Collect(ch chan<- prometheus.Metric) {
	if c.src == nil {
		return
	}
	s := c.src.Snapshot()
	ch <- prometheus.MustNewConstMetric(dBlocks, prometheus.CounterValue, float64(s.Blocks))
	ch <- prometheus.MustNewConstMetric(dSelfIP, prometheus.CounterValue, float64(s.SelfIPFailures))
	ch <- prometheus.MustNewConstMetric(dDigestCalls, prometheus.CounterValue, float64(s.DigestCalls))
	fb := map[string]uint64{}
	for reason, n := range s.DigestFallbacks {
		fb[CoerceFallback(reason)] += n
	}
	for reason, n := range fb {
		ch <- prometheus.MustNewConstMetric(dDigestFB, prometheus.CounterValue, float64(n), reason)
	}
	ch <- prometheus.MustNewConstMetric(dDixieIn, prometheus.GaugeValue, float64(s.DixieInflight))
	ch <- prometheus.MustNewConstMetric(dDixieWait, prometheus.GaugeValue, float64(s.DixieWaiters))
	emitCache := func(cache string, hit, miss uint64) {
		ch <- prometheus.MustNewConstMetric(dCache, prometheus.CounterValue, float64(hit), cache, "hit")
		ch <- prometheus.MustNewConstMetric(dCache, prometheus.CounterValue, float64(miss), cache, "miss")
	}
	emitCache("digest", s.DigestHit, s.DigestMiss)
	emitCache("embed", s.EmbedHit, s.EmbedMiss)
	emitCache("search", s.SearchHit, s.SearchMiss)

	sumQ := map[string]uint64{}
	sumC := map[string]uint64{}
	for caller, bySrc := range s.SearchQueries {
		restricted := false
		if s.Restricted != nil {
			_, restricted = s.Restricted[caller]
		}
		for src, n := range bySrc {
			if skipRestrictedSource(src) {
				continue
			}
			sumQ[src] += n
			if !restricted {
				ch <- prometheus.MustNewConstMetric(dQCaller, prometheus.CounterValue, float64(n), caller, src)
			}
		}
	}
	for caller, bySrc := range s.SearchCredits {
		restricted := false
		if s.Restricted != nil {
			_, restricted = s.Restricted[caller]
		}
		for src, n := range bySrc {
			if skipRestrictedSource(src) {
				continue
			}
			sumC[src] += n
			if !restricted {
				ch <- prometheus.MustNewConstMetric(dCCaller, prometheus.CounterValue, float64(n), caller, src)
			}
		}
	}
	for src, n := range sumQ {
		ch <- prometheus.MustNewConstMetric(dQTotal, prometheus.CounterValue, float64(n), src)
	}
	for src, n := range sumC {
		ch <- prometheus.MustNewConstMetric(dCTotal, prometheus.CounterValue, float64(n), src)
	}
	for src, n := range s.SearchDaily {
		if skipRestrictedSource(src) {
			continue
		}
		ch <- prometheus.MustNewConstMetric(dDaily, prometheus.GaugeValue, float64(n), src)
	}
	for src, n := range s.SearchCapped {
		if skipRestrictedSource(src) {
			continue
		}
		ch <- prometheus.MustNewConstMetric(dCapped, prometheus.CounterValue, float64(n), src)
	}
}

var (
	dBlocks      = prometheus.NewDesc("sift_guard_blocks_total", "SSRF/guard rejections.", nil, nil)
	dSelfIP      = prometheus.NewDesc("sift_self_ip_failures_total", "Self-IP discovery failures.", nil, nil)
	dDigestCalls = prometheus.NewDesc("sift_digest_calls_total", "Dixie digest calls started (excludes cache hits).", nil, nil)
	dDigestFB    = prometheus.NewDesc("sift_digest_fallbacks_total", "Digest fallbacks by reason.", []string{"reason"}, nil)
	dDixieIn     = prometheus.NewDesc("sift_dixie_inflight", "Sift-side Dixie permits in use.", nil, nil)
	dDixieWait   = prometheus.NewDesc("sift_dixie_queue_waiters", "Waiters in the sift-side Dixie FIFO.", nil, nil)
	dCache       = prometheus.NewDesc("sift_cache_requests_total", "Cache lookups.", []string{"cache", "result"}, nil)
	dQTotal      = prometheus.NewDesc("sift_search_queries_total", "Search queries by source (all callers).", []string{"source"}, nil)
	dQCaller     = prometheus.NewDesc("sift_search_queries_caller_total", "Search queries by fleet caller and source.", []string{"caller", "source"}, nil)
	dCTotal      = prometheus.NewDesc("sift_search_credits_total", "Search credits by source (all callers).", []string{"source"}, nil)
	dCCaller     = prometheus.NewDesc("sift_search_credits_caller_total", "Search credits by fleet caller and source.", []string{"caller", "source"}, nil)
	dDaily       = prometheus.NewDesc("sift_search_daily_credits", "UTC-day paid-source credits (fleet partition).", []string{"source"}, nil)
	dCapped      = prometheus.NewDesc("sift_search_capped_total", "Times a paid source was skipped for the daily cap.", []string{"source"}, nil)
)