opencatd-open/backend/internal/usage/recorder.go

package usage

import (
	"context"
	"fmt"
	"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
	CacheCreationTokens int
	Cost                float64
	IsError             bool
	IsCanceled          bool
	RequestID           string
	KeyID               uint64
	Protocol            string
	ErrorCode           string
	LatencyMS           int
	InputPrice          float64
	OutputPrice         float64
	CacheReadPrice      float64
	TraceID             string // TraceID for distributed tracing
	ModelID             uint64 // Model ID from channel-model binding
}

// 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
		}

		var errCode *string
		if e.ErrorCode != "" {
			errCode = &e.ErrorCode
		}

		log := &store.UsageLog{
			UserID:              e.UserID,
			KeyID:               e.KeyID,
			ChannelID:           e.ChannelID,
			ModelID:             e.ModelID,
			ModelName:           e.ModelName,
			Protocol:            e.Protocol,
			InputTokens:         int64(e.PromptTokens),
			OutputTokens:        int64(e.CompletionTokens),
			CacheReadTokens:     int64(e.CacheReadTokens),
			CacheCreationTokens: int64(e.CacheCreationTokens),
			InputPrice:          e.InputPrice,
			OutputPrice:         e.OutputPrice,
			CacheReadPrice:      e.CacheReadPrice,
			Cost:                e.Cost,
			LatencyMS:           e.LatencyMS,
			Status:              status,
			ErrorCode:           errCode,
			RequestID:           e.RequestID,
			TraceID:             e.TraceID,
		}
		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)
	}

	// Daily rollup for success and canceled requests
	dailyMap := make(map[string]*store.UsageDaily)
	for _, e := range events {
		if e.IsError {
			continue
		}
		date := time.Now().Format("2006-01-02")
		key := fmt.Sprintf("%d:%d:%s", e.UserID, e.ModelID, date)
		d := dailyMap[key]
		if d == nil {
			d = &store.UsageDaily{
				UserID:          e.UserID,
				ModelID:         e.ModelID,
				Date:            date,
				Requests:        0,
				InputTokens:     0,
				OutputTokens:    0,
				CacheReadTokens: 0,
				Cost:            0,
			}
			dailyMap[key] = d
		}
		d.Requests++
		d.InputTokens += int64(e.PromptTokens)
		d.OutputTokens += int64(e.CompletionTokens)
		d.CacheReadTokens += int64(e.CacheReadTokens)
		d.Cost += e.Cost
	}
	for _, d := range dailyMap {
		if err := r.dailyDAO.UpsertDailyUsage(context.Background(), d); err != nil {
			log.Printf("Failed to upsert daily usage: %v", err)
		}
	}

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