llmgw/internal/accounting/recorder_test.go

package accounting

import (
	"context"
	"errors"
	"log/slog"
	"sync"
	"sync/atomic"
	"testing"
	"time"

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

// fakeStore records what was written and can be made slow or broken.
type fakeStore struct {
	store.Store

	mu      sync.Mutex
	batches [][]store.Usage

	block   chan struct{}
	fail    atomic.Bool
	touched atomic.Int64
}

func (f *fakeStore) WriteUsage(_ context.Context, rows []store.Usage) error {
	if f.block != nil {
		<-f.block
	}
	if f.fail.Load() {
		return errors.New("store is down")
	}
	f.mu.Lock()
	defer f.mu.Unlock()
	f.batches = append(f.batches, append([]store.Usage(nil), rows...))
	return nil
}

func (f *fakeStore) rows() []store.Usage {
	f.mu.Lock()
	defer f.mu.Unlock()
	var out []store.Usage
	for _, b := range f.batches {
		out = append(out, b...)
	}
	return out
}

func (f *fakeStore) batchCount() int {
	f.mu.Lock()
	defer f.mu.Unlock()
	return len(f.batches)
}

type fakeToucher struct{ n *atomic.Int64 }

func (f fakeToucher) FlushTouches(context.Context) error {
	f.n.Add(1)
	return nil
}

func TestRecordsAreBatchedNotWrittenOneByOne(t *testing.T) {
	// The whole point of the buffer: a hundred requests must not be a hundred
	// transactions.
	f := &fakeStore{}
	r := NewRecorder(f, nil, RecorderConfig{FlushInterval: 20 * time.Millisecond, BatchSize: 500},
		slog.New(slog.DiscardHandler))

	ctx, cancel := context.WithCancel(context.Background())
	r.Start(ctx)

	for i := range 100 {
		r.Record(store.Usage{KeyID: "k", PromptTokens: int64(i)})
	}
	waitFor(t, time.Second, func() bool { return len(f.rows()) == 100 })

	if n := f.batchCount(); n > 5 {
		t.Errorf("100 records took %d transactions", n)
	}
	cancel()
	r.Stop()
}

func TestRecordNeverBlocksARequest(t *testing.T) {
	// A stalled database must not become an outage. Losing accounting is bad;
	// making every request wait on the store is worse.
	f := &fakeStore{block: make(chan struct{})}
	r := NewRecorder(f, nil, RecorderConfig{
		FlushInterval: time.Millisecond, BatchSize: 1, Capacity: 8,
	}, slog.New(slog.DiscardHandler))

	ctx, cancel := context.WithCancel(context.Background())
	r.Start(ctx)
	defer func() { close(f.block); cancel(); r.Stop() }()

	done := make(chan struct{})
	go func() {
		defer close(done)
		for range 1000 {
			r.Record(store.Usage{KeyID: "k"})
		}
	}()

	select {
	case <-done:
	case <-time.After(2 * time.Second):
		t.Fatal("Record blocked while the store was stalled")
	}

	_, dropped := r.Stats()
	if dropped == 0 {
		t.Error("the overflow was not counted, so nobody would ever know")
	}
}

func TestAShutdownDoesNotLoseWhatIsAlreadyQueued(t *testing.T) {
	// Those are completed requests that were already served. Throwing them away
	// means a clean shutdown silently loses the last moments of billing.
	f := &fakeStore{}
	r := NewRecorder(f, nil, RecorderConfig{FlushInterval: time.Hour, BatchSize: 10_000},
		slog.New(slog.DiscardHandler))

	ctx, cancel := context.WithCancel(context.Background())
	r.Start(ctx)

	for range 50 {
		r.Record(store.Usage{KeyID: "k"})
	}
	cancel()
	r.Stop()

	if got := len(f.rows()); got != 50 {
		t.Errorf("a clean shutdown wrote %d of 50 queued records", got)
	}
}

func TestAFailedWriteIsLoudRatherThanSilent(t *testing.T) {
	// Silently losing billing data is the kind of thing discovered at the end
	// of a month.
	f := &fakeStore{}
	f.fail.Store(true)
	logged := &countingHandler{}
	r := NewRecorder(f, nil, RecorderConfig{FlushInterval: 10 * time.Millisecond, BatchSize: 1},
		slog.New(logged))

	ctx, cancel := context.WithCancel(context.Background())
	r.Start(ctx)
	r.Record(store.Usage{KeyID: "k"})

	waitFor(t, time.Second, func() bool { return logged.count(slog.LevelError) > 0 })
	cancel()
	r.Stop()

	if written, _ := r.Stats(); written != 0 {
		t.Errorf("a failed write was counted as %d written", written)
	}
}

func TestLastUsedRidesTheSameSchedule(t *testing.T) {
	var n atomic.Int64
	f := &fakeStore{}
	r := NewRecorder(f, fakeToucher{n: &n}, RecorderConfig{FlushInterval: 10 * time.Millisecond},
		slog.New(slog.DiscardHandler))

	ctx, cancel := context.WithCancel(context.Background())
	r.Start(ctx)
	waitFor(t, time.Second, func() bool { return n.Load() > 0 })
	cancel()
	r.Stop()
}

type countingHandler struct {
	mu     sync.Mutex
	levels []slog.Level
}

func (c *countingHandler) Enabled(context.Context, slog.Level) bool { return true }
func (c *countingHandler) WithAttrs([]slog.Attr) slog.Handler       { return c }
func (c *countingHandler) WithGroup(string) slog.Handler            { return c }
func (c *countingHandler) Handle(_ context.Context, r slog.Record) error {
	c.mu.Lock()
	c.levels = append(c.levels, r.Level)
	c.mu.Unlock()
	return nil
}

func (c *countingHandler) count(l slog.Level) int {
	c.mu.Lock()
	defer c.mu.Unlock()
	n := 0
	for _, got := range c.levels {
		if got == l {
			n++
		}
	}
	return n
}

func waitFor(t *testing.T, limit time.Duration, pred func() bool) {
	t.Helper()
	deadline := time.Now().Add(limit)
	for time.Now().Before(deadline) {
		if pred() {
			return
		}
		time.Sleep(5 * time.Millisecond)
	}
	t.Fatal("condition never became true")
}