ez-api/internal/service/stats.go

package service

import (
	"context"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/redis/go-redis/v9"
)

type StatsService struct {
	rdb *redis.Client
}

type RealtimeStats struct {
	Requests       int64
	Tokens         int64
	LastAccessedAt *time.Time
}

type MasterRealtimeSnapshot struct {
	Requests    int64
	Tokens      int64
	QPS         int64
	QPSLimit    int64
	RateLimited bool
	UpdatedAt   *time.Time
}

func NewStatsService(rdb *redis.Client) *StatsService {
	return &StatsService{rdb: rdb}
}

func (s *StatsService) GetKeyRealtimeStats(ctx context.Context, tokenHash string) (RealtimeStats, error) {
	if s == nil || s.rdb == nil {
		return RealtimeStats{}, fmt.Errorf("redis client is required")
	}
	tokenHash = strings.TrimSpace(tokenHash)
	if tokenHash == "" {
		return RealtimeStats{}, fmt.Errorf("token hash required")
	}
	if ctx == nil {
		ctx = context.Background()
	}
	reqs, err := s.rdb.Get(ctx, fmt.Sprintf("key:stats:%s:requests", tokenHash)).Int64()
	if err != nil && err != redis.Nil {
		return RealtimeStats{}, fmt.Errorf("read key requests: %w", err)
	}
	tokens, err := s.rdb.Get(ctx, fmt.Sprintf("key:stats:%s:tokens", tokenHash)).Int64()
	if err != nil && err != redis.Nil {
		return RealtimeStats{}, fmt.Errorf("read key tokens: %w", err)
	}
	lastRaw, err := s.rdb.Get(ctx, fmt.Sprintf("key:stats:%s:last_access", tokenHash)).Result()
	if err != nil && err != redis.Nil {
		return RealtimeStats{}, fmt.Errorf("read key last access: %w", err)
	}
	var lastAt *time.Time
	if lastRaw != "" {
		if sec, err := strconv.ParseInt(lastRaw, 10, 64); err == nil && sec > 0 {
			t := time.Unix(sec, 0).UTC()
			lastAt = &t
		}
	}
	return RealtimeStats{Requests: reqs, Tokens: tokens, LastAccessedAt: lastAt}, nil
}

func (s *StatsService) GetMasterRealtimeStats(ctx context.Context, masterID uint) (RealtimeStats, error) {
	if s == nil || s.rdb == nil {
		return RealtimeStats{}, fmt.Errorf("redis client is required")
	}
	if masterID == 0 {
		return RealtimeStats{}, fmt.Errorf("master id required")
	}
	if ctx == nil {
		ctx = context.Background()
	}
	reqs, err := s.rdb.Get(ctx, fmt.Sprintf("master:stats:%d:requests", masterID)).Int64()
	if err != nil && err != redis.Nil {
		return RealtimeStats{}, fmt.Errorf("read master requests: %w", err)
	}
	tokens, err := s.rdb.Get(ctx, fmt.Sprintf("master:stats:%d:tokens", masterID)).Int64()
	if err != nil && err != redis.Nil {
		return RealtimeStats{}, fmt.Errorf("read master tokens: %w", err)
	}
	return RealtimeStats{Requests: reqs, Tokens: tokens}, nil
}

func (s *StatsService) GetMasterRealtimeSnapshot(ctx context.Context, masterID uint) (MasterRealtimeSnapshot, error) {
	if s == nil || s.rdb == nil {
		return MasterRealtimeSnapshot{}, fmt.Errorf("redis client is required")
	}
	if masterID == 0 {
		return MasterRealtimeSnapshot{}, fmt.Errorf("master id required")
	}
	if ctx == nil {
		ctx = context.Background()
	}
	statsKey := fmt.Sprintf("master:stats:%d", masterID)
	rateKey := fmt.Sprintf("master:rate:%d", masterID)

	pipe := s.rdb.Pipeline()
	reqCmd := pipe.Get(ctx, statsKey+":requests")
	tokCmd := pipe.Get(ctx, statsKey+":tokens")
	qpsCmd := pipe.HGet(ctx, rateKey, "qps")
	limitCmd := pipe.HGet(ctx, rateKey, "limit")
	blockedCmd := pipe.HGet(ctx, rateKey, "blocked")
	updatedCmd := pipe.HGet(ctx, rateKey, "updated_at")
	if _, err := pipe.Exec(ctx); err != nil && err != redis.Nil {
		return MasterRealtimeSnapshot{}, fmt.Errorf("read realtime stats: %w", err)
	}

	requests := readCmdInt64(reqCmd)
	tokens := readCmdInt64(tokCmd)
	qps := readCmdInt64(qpsCmd)
	limit := readCmdInt64(limitCmd)
	blocked := readCmdInt64(blockedCmd) == 1
	updatedAt := readCmdTime(updatedCmd)
	if limit < 0 {
		limit = 0
	}

	return MasterRealtimeSnapshot{
		Requests:    requests,
		Tokens:      tokens,
		QPS:         qps,
		QPSLimit:    limit,
		RateLimited: blocked,
		UpdatedAt:   updatedAt,
	}, nil
}

func readCmdInt64(cmd *redis.StringCmd) int64 {
	if cmd == nil {
		return 0
	}
	v, err := cmd.Int64()
	if err != nil {
		return 0
	}
	return v
}

func readCmdTime(cmd *redis.StringCmd) *time.Time {
	if cmd == nil {
		return nil
	}
	v, err := cmd.Int64()
	if err != nil || v <= 0 {
		return nil
	}
	tm := time.Unix(v, 0).UTC()
	return &tm
}