geo/server/internal/shared/cache/metrics.go

package cache

import (
	"context"
	"errors"
	"sort"
	"strings"
	"sync"
	"sync/atomic"
	"time"

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

type metricsKey struct {
	Operation string
	Result    string
}

type cacheMetricsRecorder struct {
	operations       sync.Map
	l1HitTotal       atomic.Int64
	l1MissTotal      atomic.Int64
	asyncDropTotal   atomic.Int64
	asyncPanicTotal  atomic.Int64
	asyncQueueLength atomic.Int64
}

type cacheOperationMetrics struct {
	Count          atomic.Int64
	DurationNanos atomic.Int64
}

var cacheMetrics = &cacheMetricsRecorder{}

type observedCache struct {
	inner  Cache
	driver string
}

func newObservedCache(inner Cache, driver string) Cache {
	driver = strings.TrimSpace(driver)
	if driver == "" {
		driver = "unknown"
	}
	return &observedCache{inner: inner, driver: driver}
}

func (c *observedCache) Get(ctx context.Context, key string) ([]byte, error) {
	start := time.Now()
	value, err := c.inner.Get(ctx, key)
	recordCacheOperation("get", cacheResult(err), time.Since(start))
	return value, err
}

func (c *observedCache) GetMulti(ctx context.Context, keys []string) (map[string][]byte, error) {
	start := time.Now()
	values, err := c.inner.GetMulti(ctx, keys)
	result := cacheResult(err)
	if err == nil {
		recordCacheOperation("get_multi_hit", "ok", time.Duration(0))
		missCount := len(keys) - len(values)
		if missCount > 0 {
			for i := 0; i < missCount; i++ {
				recordCacheOperation("get_multi_miss", "not_found", time.Duration(0))
			}
		}
	}
	recordCacheOperation("get_multi", result, time.Since(start))
	return values, err
}

func (c *observedCache) Set(ctx context.Context, key string, value []byte, ttl time.Duration) error {
	start := time.Now()
	err := c.inner.Set(ctx, key, value, ttl)
	recordCacheOperation("set", cacheResult(err), time.Since(start))
	return err
}

func (c *observedCache) SetMulti(ctx context.Context, items []Item) error {
	start := time.Now()
	err := c.inner.SetMulti(ctx, items)
	recordCacheOperation("set_multi", cacheResult(err), time.Since(start))
	return err
}

func (c *observedCache) Delete(ctx context.Context, key string) error {
	start := time.Now()
	err := c.inner.Delete(ctx, key)
	recordCacheOperation("delete", cacheResult(err), time.Since(start))
	return err
}

func (c *observedCache) DeleteMany(ctx context.Context, keys []string) error {
	start := time.Now()
	err := c.inner.DeleteMany(ctx, keys)
	recordCacheOperation("delete_many", cacheResult(err), time.Since(start))
	return err
}

func (c *observedCache) DeletePrefix(ctx context.Context, prefix string) error {
	start := time.Now()
	err := c.inner.DeletePrefix(ctx, prefix)
	recordCacheOperation("delete_prefix", cacheResult(err), time.Since(start))
	return err
}

func cacheResult(err error) string {
	switch {
	case err == nil:
		return "ok"
	case errors.Is(err, ErrNotFound):
		return "not_found"
	case errors.Is(err, ErrAsyncQueueFull):
		return "async_queue_full"
	default:
		return "error"
	}
}

func recordCacheOperation(operation, result string, duration time.Duration) {
	if operation == "" {
		operation = "unknown"
	}
	if result == "" {
		result = "unknown"
	}
	valueAny, _ := cacheMetrics.operations.LoadOrStore(metricsKey{
		Operation: operation,
		Result:    result,
	}, &cacheOperationMetrics{})
	value := valueAny.(*cacheOperationMetrics)
	value.Count.Add(1)
	if duration > 0 {
		value.DurationNanos.Add(duration.Nanoseconds())
	}
}

func recordCacheL1(result string) {
	if result == "hit" {
		cacheMetrics.l1HitTotal.Add(1)
		return
	}
	cacheMetrics.l1MissTotal.Add(1)
}

func recordCacheAsyncDrop(operation string) {
	cacheMetrics.asyncDropTotal.Add(1)
	recordCacheOperation("async_drop_"+operation, "dropped", time.Duration(0))
}

func recordCacheAsyncPanic() {
	cacheMetrics.asyncPanicTotal.Add(1)
}

func recordCacheAsyncQueue(length int) {
	cacheMetrics.asyncQueueLength.Store(int64(length))
}

type MetricsSnapshot struct {
	Operations       []OperationMetric `json:"operations"`
	L1HitTotal       int64             `json:"l1_hit_total"`
	L1MissTotal      int64             `json:"l1_miss_total"`
	AsyncDropTotal   int64             `json:"async_drop_total"`
	AsyncPanicTotal  int64             `json:"async_panic_total"`
	AsyncQueueLength int64             `json:"async_queue_length"`
}

type OperationMetric struct {
	Operation       string  `json:"operation"`
	Result          string  `json:"result"`
	Count           int64   `json:"count"`
	DurationSeconds float64 `json:"duration_seconds"`
}

func MetricsSnapshotValue() MetricsSnapshot {
	snapshot := MetricsSnapshot{
		L1HitTotal:       cacheMetrics.l1HitTotal.Load(),
		L1MissTotal:      cacheMetrics.l1MissTotal.Load(),
		AsyncDropTotal:   cacheMetrics.asyncDropTotal.Load(),
		AsyncPanicTotal:  cacheMetrics.asyncPanicTotal.Load(),
		AsyncQueueLength: cacheMetrics.asyncQueueLength.Load(),
	}
	cacheMetrics.operations.Range(func(keyAny, valueAny any) bool {
		key, ok := keyAny.(metricsKey)
		if !ok {
			return true
		}
		value, ok := valueAny.(*cacheOperationMetrics)
		if !ok {
			return true
		}
		snapshot.Operations = append(snapshot.Operations, OperationMetric{
			Operation:       key.Operation,
			Result:          key.Result,
			Count:           value.Count.Load(),
			DurationSeconds: float64(value.DurationNanos.Load()) / float64(time.Second),
		})
		return true
	})
	sort.Slice(snapshot.Operations, func(i, j int) bool {
		if snapshot.Operations[i].Operation == snapshot.Operations[j].Operation {
			return snapshot.Operations[i].Result < snapshot.Operations[j].Result
		}
		return snapshot.Operations[i].Operation < snapshot.Operations[j].Operation
	})
	return snapshot
}

func ResetMetricsForTest() {
	cacheMetrics = &cacheMetricsRecorder{}
}

type MetricsCollector struct{}

func (MetricsCollector) Describe(chan<- *prometheus.Desc) {}

func (MetricsCollector) Collect(ch chan<- prometheus.Metric) {
	snapshot := MetricsSnapshotValue()
	for _, item := range snapshot.Operations {
		ch <- prometheus.MustNewConstMetric(
			prometheus.NewDesc("cache_operations_total", "Total cache operations by operation and result.", []string{"operation", "result"}, nil),
			prometheus.CounterValue,
			float64(item.Count),
			item.Operation,
			item.Result,
		)
		ch <- prometheus.MustNewConstMetric(
			prometheus.NewDesc("cache_operation_duration_seconds_total", "Total cache operation duration by operation and result.", []string{"operation", "result"}, nil),
			prometheus.CounterValue,
			item.DurationSeconds,
			item.Operation,
			item.Result,
		)
	}
	ch <- prometheus.MustNewConstMetric(
		prometheus.NewDesc("cache_l1_access_total", "Total L1 cache access by result.", []string{"result"}, nil),
		prometheus.CounterValue,
		float64(snapshot.L1HitTotal),
		"hit",
	)
	ch <- prometheus.MustNewConstMetric(
		prometheus.NewDesc("cache_l1_access_total", "Total L1 cache access by result.", []string{"result"}, nil),
		prometheus.CounterValue,
		float64(snapshot.L1MissTotal),
		"miss",
	)
	ch <- prometheus.MustNewConstMetric(
		prometheus.NewDesc("cache_async_drops_total", "Total async cache fill tasks dropped because the queue was full.", nil, nil),
		prometheus.CounterValue,
		float64(snapshot.AsyncDropTotal),
	)
	ch <- prometheus.MustNewConstMetric(
		prometheus.NewDesc("cache_async_panics_total", "Total panics recovered in async cache fill workers.", nil, nil),
		prometheus.CounterValue,
		float64(snapshot.AsyncPanicTotal),
	)
	ch <- prometheus.MustNewConstMetric(
		prometheus.NewDesc("cache_async_queue_length", "Current async cache fill queue length.", nil, nil),
		prometheus.GaugeValue,
		float64(snapshot.AsyncQueueLength),
	)
}