Luminary/internal/monitor/service.go

package monitor

import (
	"context"
	"log"
	"time"

	"github.com/rose_cat707/luminary/internal/gateway"
	"github.com/rose_cat707/luminary/internal/model"
	"github.com/rose_cat707/luminary/internal/settings"
	"github.com/rose_cat707/luminary/internal/store/cache"
	sqlitestore "github.com/rose_cat707/luminary/internal/store/sqlite"
	"gorm.io/gorm"
)

type Service struct {
	gateway       *gateway.Service
	cache         *cache.Store
	db            *gorm.DB
	settings      *settings.Runtime
	stop          chan struct{}
	healthRestart chan struct{}
	flushRestart  chan struct{}
	flushFn       func([]cache.UsageDelta)
}

func New(gw *gateway.Service, c *cache.Store, store *sqlitestore.Store, rt *settings.Runtime) *Service {
	return &Service{
		gateway:       gw,
		cache:         c,
		db:            store.DB(),
		settings:      rt,
		stop:          make(chan struct{}),
		healthRestart: make(chan struct{}, 1),
		flushRestart:  make(chan struct{}, 1),
	}
}

func (s *Service) Start(flushFn func([]cache.UsageDelta)) {
	s.flushFn = flushFn
	go s.healthLoop()
	go s.usageFlushLoop()
	go s.dailyResetLoop()
}

func (s *Service) Stop() {
	close(s.stop)
}

func (s *Service) NotifySettingsChanged() {
	select {
	case s.healthRestart <- struct{}{}:
	default:
	}
	select {
	case s.flushRestart <- struct{}{}:
	default:
	}
}

func (s *Service) healthInterval() time.Duration {
	interval := s.settings.HealthCheckInterval()
	if interval < time.Second {
		return 30 * time.Second
	}
	return interval
}

func (s *Service) flushInterval() time.Duration {
	interval := s.settings.UsageFlushInterval()
	if interval < time.Second {
		return 5 * time.Second
	}
	return interval
}

func (s *Service) healthLoop() {
	var ticker *time.Ticker
	resetTicker := func() {
		if ticker != nil {
			ticker.Stop()
		}
		ticker = time.NewTicker(s.healthInterval())
	}
	resetTicker()
	defer ticker.Stop()
	s.runChecks()
	for {
		select {
		case <-ticker.C:
			s.runChecks()
		case <-s.healthRestart:
			resetTicker()
		case <-s.stop:
			return
		}
	}
}

func (s *Service) usageFlushLoop() {
	var ticker *time.Ticker
	resetTicker := func() {
		if ticker != nil {
			ticker.Stop()
		}
		ticker = time.NewTicker(s.flushInterval())
	}
	resetTicker()
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
			s.flushUsage()
		case <-s.flushRestart:
			resetTicker()
		case <-s.stop:
			s.flushUsage()
			return
		}
	}
}

func (s *Service) flushUsage() {
	if s.flushFn == nil {
		return
	}
	deltas := s.cache.DrainUsageDeltas()
	if len(deltas) > 0 {
		s.flushFn(deltas)
	}
}

func (s *Service) runChecks() {
	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
	defer cancel()
	for _, p := range s.cache.ListProviders() {
		status, latency, errMsg := s.gateway.HealthCheckProvider(ctx, p.ID)
		rec := model.HealthCheckRecord{
			ProviderID: p.ID,
			Status:     status,
			LatencyMs:  latency,
			ErrorMsg:   errMsg,
			CheckedAt:  time.Now(),
		}
		s.cache.SetHealthStatus(p.ID, rec)
		if err := s.db.Create(&rec).Error; err != nil {
			log.Printf("save health check: %v", err)
		}
	}
}

func (s *Service) dailyResetLoop() {
	for {
		now := time.Now()
		next := time.Date(now.Year(), now.Month(), now.Day()+1, 0, 0, 0, 0, now.Location())
		timer := time.NewTimer(time.Until(next))
		select {
		case <-timer.C:
			s.resetDailyCounters()
		case <-s.stop:
			timer.Stop()
			return
		}
	}
}

func (s *Service) resetDailyCounters() {
	s.cache.ResetProviderKeyDailyCounts()
	if err := s.db.Model(&model.ProviderKey{}).Updates(map[string]any{
		"requests_today": 0,
		"tokens_today":   0,
	}).Error; err != nil {
		log.Printf("reset daily provider key counters: %v", err)
	}
}