Prism/internal/metrics/collector.go

package metrics

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"sync"
	"time"

	"github.com/prism/proxy/internal/cache"
	"github.com/prism/proxy/internal/mihomoapi"
	"github.com/prism/proxy/internal/store"
)

type Collector struct {
	store   *store.Store
	cache   *cache.Cache
	apiBase string
	secret  string
	client  *http.Client

	mu          sync.Mutex
	connCache   []map[string]any
	connCacheAt time.Time
	refreshing  bool
}

func NewCollector(s *store.Store, c *cache.Cache, apiBase, secret string) *Collector {
	return &Collector{
		store:   s,
		cache:   c,
		apiBase: trimRight(apiBase, "/"),
		secret:  secret,
		client: &http.Client{
			Timeout: 3 * time.Second,
		},
	}
}

func trimRight(s, cut string) string {
	for len(s) > 0 && len(s) >= len(cut) && s[len(s)-len(cut):] == cut {
		s = s[:len(s)-len(cut)]
	}
	return s
}

// Snapshot returns cached dashboard data and never blocks on Mihomo.
func (m *Collector) Snapshot(ctx context.Context) cache.DashboardSnapshot {
	snap := m.cache.GetDashboard()
	total, alive, _ := m.store.CountEnabledProxies(ctx)
	snap.TotalNodes = total
	snap.AliveNodes = alive
	return snap
}

func (m *Collector) RefreshIfStale(ctx context.Context, maxAge time.Duration) cache.DashboardSnapshot {
	snap := m.cache.GetDashboard()
	if !snap.UpdatedAt.IsZero() && time.Since(snap.UpdatedAt) < maxAge {
		return snap
	}
	refreshed, err := m.Refresh(ctx)
	if err != nil {
		return snap
	}
	return refreshed
}

func (m *Collector) RefreshAsync(maxAge time.Duration) {
	snap := m.cache.GetDashboard()
	if !snap.UpdatedAt.IsZero() && time.Since(snap.UpdatedAt) < maxAge {
		return
	}
	m.mu.Lock()
	if m.refreshing {
		m.mu.Unlock()
		return
	}
	m.refreshing = true
	m.mu.Unlock()
	go func() {
		defer func() {
			m.mu.Lock()
			m.refreshing = false
			m.mu.Unlock()
		}()
		ctx, cancel := context.WithTimeout(context.Background(), 6*time.Second)
		defer cancel()
		_, _ = m.Refresh(ctx)
	}()
}

func (m *Collector) Refresh(ctx context.Context) (cache.DashboardSnapshot, error) {
	var snap cache.DashboardSnapshot
	snap.UpdatedAt = time.Now()

	var (
		traffic TrafficRatesResult
		conns   map[string]any
		wg      sync.WaitGroup
	)
	wg.Add(2)
	go func() {
		defer wg.Done()
		rates, err := mihomoapi.New(m.apiBase, m.secret).TrafficRates(ctx)
		if err == nil {
			traffic = TrafficRatesResult{Up: rates.Up, Down: rates.Down, OK: true}
		}
	}()
	go func() {
		defer wg.Done()
		conns, _ = m.getJSON("/connections")
	}()
	wg.Wait()

	if traffic.OK {
		snap.UploadRate = traffic.Up
		snap.DownloadRate = traffic.Down
	}

	if conns != nil {
		if arr, ok := conns["connections"].([]any); ok {
			snap.Connections = len(arr)
			m.setConnCache(arr)
		}
	}

	total, alive, _ := m.store.CountEnabledProxies(ctx)
	snap.TotalNodes = total
	snap.AliveNodes = alive

	m.cache.SetDashboard(snap)
	_ = m.store.AddTrafficSnapshot(ctx, snap.UploadRate, snap.DownloadRate, int64(snap.Connections))
	return snap, nil
}

func (m *Collector) setConnCache(arr []any) {
	out := make([]map[string]any, 0, len(arr))
	for _, item := range arr {
		if m, ok := item.(map[string]any); ok {
			out = append(out, m)
		}
	}
	m.mu.Lock()
	m.connCache = out
	m.connCacheAt = time.Now()
	m.mu.Unlock()
}

func (m *Collector) CachedConnections() []map[string]any {
	m.mu.Lock()
	defer m.mu.Unlock()
	return copyConnSlice(m.connCache)
}

func (m *Collector) GetConnections() ([]map[string]any, error) {
	return m.GetConnectionsCached(0)
}

func (m *Collector) GetConnectionsCached(maxAge time.Duration) ([]map[string]any, error) {
	m.mu.Lock()
	if maxAge > 0 && m.connCache != nil && time.Since(m.connCacheAt) < maxAge {
		out := copyConnSlice(m.connCache)
		m.mu.Unlock()
		return out, nil
	}
	m.mu.Unlock()

	m.RefreshAsync(maxAge)
	if cached := m.CachedConnections(); len(cached) > 0 {
		return cached, nil
	}

	data, err := m.getJSON("/connections")
	if err != nil {
		m.mu.Lock()
		defer m.mu.Unlock()
		if m.connCache != nil {
			return copyConnSlice(m.connCache), nil
		}
		return nil, err
	}
	arr, ok := data["connections"].([]any)
	if !ok {
		return nil, nil
	}
	m.setConnCache(arr)
	return m.CachedConnections(), nil
}

func copyConnSlice(in []map[string]any) []map[string]any {
	if len(in) == 0 {
		return nil
	}
	out := make([]map[string]any, len(in))
	copy(out, in)
	return out
}

type TrafficRatesResult struct {
	Up   int64
	Down int64
	OK   bool
}

func (m *Collector) UpdateAPIBase(apiBase, secret string) {
	m.apiBase = trimRight(apiBase, "/")
	m.secret = secret
}

func (m *Collector) getJSON(path string) (map[string]any, error) {
	req, err := http.NewRequest(http.MethodGet, m.apiBase+path, nil)
	if err != nil {
		return nil, err
	}
	if m.secret != "" {
		req.Header.Set("Authorization", "Bearer "+m.secret)
	}
	resp, err := m.client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode >= 400 {
		return nil, fmt.Errorf("api error %d: %s", resp.StatusCode, string(body))
	}
	var out map[string]any
	if err := json.Unmarshal(body, &out); err != nil {
		return nil, err
	}
	return out, nil
}