llmgw/internal/accounting/recorder.go

package accounting

import (
	"context"
	"log/slog"
	"sync"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

// RecorderConfig tunes the batch writer.
type RecorderConfig struct {
	// FlushInterval and BatchSize decide when a batch is written. Whichever
	// comes first.
	FlushInterval time.Duration
	BatchSize     int
	// Capacity bounds the buffer. Beyond it, records are dropped rather than
	// blocking a request — see Record.
	Capacity int
}

// DefaultRecorderConfig returns sensible values.
func DefaultRecorderConfig() RecorderConfig {
	return RecorderConfig{FlushInterval: 200 * time.Millisecond, BatchSize: 500, Capacity: 8192}
}

func (c *RecorderConfig) normalize() {
	d := DefaultRecorderConfig()
	if c.FlushInterval <= 0 {
		c.FlushInterval = d.FlushInterval
	}
	if c.BatchSize <= 0 {
		c.BatchSize = d.BatchSize
	}
	if c.Capacity <= 0 {
		c.Capacity = d.Capacity
	}
}

// Toucher is the part of the authenticator the recorder flushes alongside
// usage, so last-used timestamps ride the same transaction schedule.
type Toucher interface {
	FlushTouches(ctx context.Context) error
}

// Recorder batches usage records into the store.
//
// The buffered channel is the whole design. A request must never wait on a
// database write, so Record is a non-blocking send and the writer drains it on
// its own schedule.
type Recorder struct {
	cfg   RecorderConfig
	store store.Store
	touch Toucher
	log   *slog.Logger

	ch   chan store.Usage
	done chan struct{}
	wg   sync.WaitGroup

	mu      sync.Mutex
	dropped int64
	written int64
}

// NewRecorder builds a recorder. Start begins writing.
func NewRecorder(s store.Store, touch Toucher, cfg RecorderConfig, log *slog.Logger) *Recorder {
	cfg.normalize()
	if log == nil {
		log = slog.Default()
	}
	return &Recorder{
		cfg: cfg, store: s, touch: touch, log: log,
		ch:   make(chan store.Usage, cfg.Capacity),
		done: make(chan struct{}),
	}
}

// Record queues one request's accounting.
//
// Never blocks. If the buffer is full — a database stalled long enough to fill
// eight thousand records — the record is dropped and counted. Losing some
// accounting is bad; making every request wait on a stalled database is worse,
// because it converts a reporting problem into an outage.
func (r *Recorder) Record(u store.Usage) {
	select {
	case r.ch <- u:
	default:
		r.mu.Lock()
		r.dropped++
		n := r.dropped
		r.mu.Unlock()
		if n == 1 || n%1000 == 0 {
			r.log.Warn("dropping usage records: the store is not keeping up",
				"dropped_total", n)
		}
	}
}

// Start runs the writer until ctx ends.
func (r *Recorder) Start(ctx context.Context) {
	r.wg.Add(1)
	go func() {
		defer r.wg.Done()
		defer close(r.done)
		r.run(ctx)
	}()
}

// Stop waits for the writer to drain and finish.
func (r *Recorder) Stop() { r.wg.Wait() }

func (r *Recorder) run(ctx context.Context) {
	t := time.NewTicker(r.cfg.FlushInterval)
	defer t.Stop()

	batch := make([]store.Usage, 0, r.cfg.BatchSize)
	for {
		select {
		case <-ctx.Done():
			// Drain what is already queued before leaving. These are completed
			// requests that have already been served; throwing them away would
			// mean a clean shutdown silently loses the last moments of billing.
			r.drain(&batch)
			r.flush(context.WithoutCancel(ctx), batch)
			return

		case u := <-r.ch:
			batch = append(batch, u)
			if len(batch) >= r.cfg.BatchSize {
				r.flush(ctx, batch)
				batch = batch[:0]
			}

		case <-t.C:
			r.drain(&batch)
			if len(batch) > 0 {
				r.flush(ctx, batch)
				batch = batch[:0]
			}
			if err := r.touchFlush(ctx); err != nil {
				r.log.Warn("failed to write last-used timestamps", "err", err)
			}
		}
	}
}

// drain moves everything queued into the batch without blocking.
func (r *Recorder) drain(batch *[]store.Usage) {
	for len(*batch) < r.cfg.BatchSize {
		select {
		case u := <-r.ch:
			*batch = append(*batch, u)
		default:
			return
		}
	}
}

func (r *Recorder) flush(ctx context.Context, batch []store.Usage) {
	if len(batch) == 0 {
		return
	}
	if err := r.store.WriteUsage(ctx, batch); err != nil {
		// The records are gone either way; what matters is that this is loud.
		// Silently losing billing data is the kind of thing discovered at the
		// end of a month.
		r.log.Error("failed to write usage records", "count", len(batch), "err", err)
		return
	}
	r.mu.Lock()
	r.written += int64(len(batch))
	r.mu.Unlock()
}

func (r *Recorder) touchFlush(ctx context.Context) error {
	if r.touch == nil {
		return nil
	}
	return r.touch.FlushTouches(ctx)
}

// Stats reports what the recorder has done, for metrics and tests.
func (r *Recorder) Stats() (written, dropped int64) {
	r.mu.Lock()
	defer r.mu.Unlock()
	return r.written, r.dropped
}