opencatd-open/internal/usage/recorder.go

package usage

import (
	"context"
	"log"
	"opencatd-open/internal/dao"
	"opencatd-open/internal/store"
	"sync"
	"time"
)

// Event represents a usage event to be recorded
type Event struct {
	UserID           uint64
	ModelName        string
	ChannelID        uint64
	PromptTokens     int
	CompletionTokens int
	CacheReadTokens  int
	Cost             float64
	IsError          bool
	IsCanceled       bool
	RequestID        string
}

// Recorder handles async usage recording
type Recorder struct {
	usageDAO      *dao.UsageDAO
	dailyDAO      *dao.DailyUsageDAO
	ch            chan Event
	batchSize     int
	flushInterval time.Duration
	wg            sync.WaitGroup
}

// NewRecorder creates a new usage recorder
func NewRecorder(usageDAO *dao.UsageDAO, dailyDAO *dao.DailyUsageDAO) *Recorder {
	return &Recorder{
		usageDAO:      usageDAO,
		dailyDAO:      dailyDAO,
		ch:            make(chan Event, 10000),
		batchSize:     100,
		flushInterval: 5 * time.Second,
	}
}

// Start starts the recorder's background workers
func (r *Recorder) Start(ctx context.Context) {
	r.wg.Add(1)
	go r.processLoop(ctx)
}

// Stop gracefully stops the recorder
func (r *Recorder) Stop() {
	close(r.ch)
	r.wg.Wait()
}

// Record queues a usage event for async recording
func (r *Recorder) Record(event Event) {
	select {
	case r.ch <- event:
	default:
		log.Printf("Usage channel full, dropping event for user %d model %s", event.UserID, event.ModelName)
	}
}

func (r *Recorder) processLoop(ctx context.Context) {
	defer r.wg.Done()

	batch := make([]Event, 0, r.batchSize)
	ticker := time.NewTicker(r.flushInterval)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			if len(batch) > 0 {
				r.flush(batch)
			}
			return
		case event, ok := <-r.ch:
			if !ok {
				if len(batch) > 0 {
					r.flush(batch)
				}
				return
			}
			batch = append(batch, event)
			if len(batch) >= r.batchSize {
				r.flush(batch)
				batch = make([]Event, 0, r.batchSize)
			}
		case <-ticker.C:
			if len(batch) > 0 {
				r.flush(batch)
				batch = make([]Event, 0, r.batchSize)
			}
		}
	}
}

func (r *Recorder) flush(events []Event) {
	if len(events) == 0 {
		return
	}

	// Batch create usage logs
	logs := make([]*store.UsageLog, 0, len(events))

	for _, e := range events {
		status := store.UsageStatusSuccess
		if e.IsError {
			status = store.UsageStatusError
		}
		if e.IsCanceled {
			status = store.UsageStatusCanceled
		}

		log := &store.UsageLog{
			UserID:           e.UserID,
			ModelName:        e.ModelName,
			ChannelID:        e.ChannelID,
			InputTokens:      int64(e.PromptTokens),
			OutputTokens:     int64(e.CompletionTokens),
			CacheReadTokens:  int64(e.CacheReadTokens),
			Cost:             e.Cost,
			Status:           status,
			RequestID:        e.RequestID,
		}
		logs = append(logs, log)
	}

	// Write to database
	if err := r.usageDAO.BatchCreate(context.Background(), logs); err != nil {
		log.Printf("Failed to batch create usage logs: %v", err)
	}

	log.Printf("Flushed %d usage logs", len(logs))
}