feat(natsstore): keep cooldown state in a bucket with per-key TTL

internal/store/natsstore/cooldown_store.go

@@ -0,0 +1,218 @@
+package natsstore
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"strings"
+	"sync"
+	"time"
+
+	cliproxyauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth"
+)
+
+const (
+	// cooldownGrace keeps a record a little past the retry time it describes,
+	// so a replica with a slow clock still sees the cooldown end.
+	cooldownGrace = time.Minute
+	// cooldownIdleTTL bounds a record that names no retry time.
+	cooldownIdleTTL = time.Hour
+	// cooldownTombstoneTTL is how long a cleared record keeps refusing older
+	// writes from a replica that has not seen it cleared yet.
+	cooldownTombstoneTTL = 10 * time.Minute
+	// maxCooldownWriteAttempts bounds the retries of one record when replicas race.
+	maxCooldownWriteAttempts = 4
+)
+
+var (
+	_ cliproxyauth.CooldownStateStore         = (*CooldownStore)(nil)
+	_ cliproxyauth.CooldownStateStoreProvider = (*Store)(nil)
+)
+
+// CooldownStore keeps runtime cooldown state in a bucket with per-key TTL, so a
+// cooldown disappears on its own once its retry time has passed. It merges the
+// way the PostgreSQL store does: a write never replaces a record that carries a
+// newer UpdatedAt, and a cleared record leaves a tombstone that refuses older writes.
+type CooldownStore struct {
+	kv       *KV
+	mu       sync.Mutex
+	previous map[string]time.Time
+}
+
+// cooldownEntry is the stored form of one auth/model pair.
+type cooldownEntry struct {
+	Deleted   bool                              `json:"deleted,omitempty"`
+	UpdatedAt time.Time                         `json:"updated_at"`
+	Record    *cliproxyauth.CooldownStateRecord `json:"record,omitempty"`
+}
+
+// NewCooldownStore wraps the cooldown bucket.
+func NewCooldownStore(kv *KV) *CooldownStore {
+	return &CooldownStore{kv: kv, previous: make(map[string]time.Time)}
+}
+
+// CooldownStateStore returns the cooldown store the manager persists through.
+func (s *Store) CooldownStateStore() cliproxyauth.CooldownStateStore {
+	if s == nil || s.cooldown == nil {
+		return nil
+	}
+	return s.cooldown
+}
+
+func cooldownID(record cliproxyauth.CooldownStateRecord) string {
+	authID := strings.TrimSpace(record.AuthID)
+	if authID == "" {
+		return ""
+	}
+	return authID + "\x1f" + strings.TrimSpace(record.Model)
+}
+
+func normalizeCooldownTime(value, fallback time.Time) time.Time {
+	if value.IsZero() {
+		value = fallback
+	}
+	return value.UTC().Truncate(time.Microsecond)
+}
+
+// Load returns every live record and remembers what it saw, so the next Save
+// can tell which records were cleared since.
+func (s *CooldownStore) Load(ctx context.Context) ([]cliproxyauth.CooldownStateRecord, error) {
+	if s == nil || s.kv == nil {
+		return nil, errors.New("nats cooldown store: not initialized")
+	}
+	if ctx == nil {
+		ctx = context.Background()
+	}
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	entries, err := s.kv.List(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("nats cooldown store: load state: %w", err)
+	}
+	records := make([]cliproxyauth.CooldownStateRecord, 0, len(entries))
+	previous := make(map[string]time.Time, len(entries))
+	for _, entry := range entries {
+		var stored cooldownEntry
+		if err := json.Unmarshal(entry.Value, &stored); err != nil {
+			return nil, fmt.Errorf("nats cooldown store: decode state: %w", err)
+		}
+		if stored.Deleted || stored.Record == nil {
+			continue
+		}
+		if cooldownID(*stored.Record) != entry.ID {
+			return nil, errors.New("nats cooldown store: decoded state does not match its key")
+		}
+		records = append(records, *stored.Record)
+		previous[entry.ID] = stored.UpdatedAt
+	}
+	s.previous = previous
+	return records, nil
+}
+
+// Save writes the records and clears the ones a previous Load or Save knew that are
+// missing now. A record newer than the one being written wins.
+func (s *CooldownStore) Save(ctx context.Context, records []cliproxyauth.CooldownStateRecord) error {
+	if s == nil || s.kv == nil {
+		return errors.New("nats cooldown store: not initialized")
+	}
+	if ctx == nil {
+		ctx = context.Background()
+	}
+	now := normalizeCooldownTime(time.Now(), time.Time{})
+	current := make(map[string]time.Time, len(records))
+	type write struct {
+		id    string
+		entry cooldownEntry
+		ttl   time.Duration
+	}
+	writes := make([]write, 0, len(records))
+	for i := range records {
+		record := records[i]
+		id := cooldownID(record)
+		if id == "" {
+			return errors.New("nats cooldown store: state has empty auth ID")
+		}
+		record.UpdatedAt = normalizeCooldownTime(record.UpdatedAt, now)
+		current[id] = record.UpdatedAt
+		writes = append(writes, write{id: id, entry: cooldownEntry{UpdatedAt: record.UpdatedAt, Record: &record}, ttl: cooldownTTL(record, now)})
+	}
+
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	for _, w := range writes {
+		if err := s.put(ctx, w.id, w.entry, w.ttl, func(stored cooldownEntry) bool {
+			return !stored.UpdatedAt.After(w.entry.UpdatedAt)
+		}); err != nil {
+			return fmt.Errorf("nats cooldown store: save state: %w", err)
+		}
+	}
+	for id, seen := range s.previous {
+		if _, ok := current[id]; ok {
+			continue
+		}
+		clearedAt := now
+		if !clearedAt.After(seen) {
+			clearedAt = seen.Add(time.Microsecond)
+		}
+		// Only the record this replica saw is cleared. One written since by
+		// another replica, or already cleared, stays as it is.
+		err := s.put(ctx, id, cooldownEntry{Deleted: true, UpdatedAt: clearedAt}, cooldownTombstoneTTL, func(stored cooldownEntry) bool {
+			return !stored.Deleted && !stored.UpdatedAt.After(seen)
+		})
+		if err != nil {
+			return fmt.Errorf("nats cooldown store: clear state: %w", err)
+		}
+	}
+	s.previous = current
+	return nil
+}
+
+// cooldownTTL keeps a record until its retry time has passed.
+func cooldownTTL(record cliproxyauth.CooldownStateRecord, now time.Time) time.Duration {
+	if record.NextRetryAfter.IsZero() {
+		return cooldownIdleTTL
+	}
+	ttl := record.NextRetryAfter.Sub(now) + cooldownGrace
+	if ttl < cooldownGrace {
+		return cooldownGrace
+	}
+	return ttl
+}
+
+// put writes entry when the key is empty or replace allows it. A key's TTL is set
+// only when it is created, so replacing a record deletes it at the revision that
+// was read and creates it again; losing either step to another replica rereads.
+func (s *CooldownStore) put(ctx context.Context, id string, entry cooldownEntry, ttl time.Duration, replace func(cooldownEntry) bool) error {
+	value, err := json.Marshal(entry)
+	if err != nil {
+		return fmt.Errorf("encode %s: %w", id, err)
+	}
+	for range maxCooldownWriteAttempts {
+		stored, err := s.kv.Get(ctx, id)
+		switch {
+		case errors.Is(err, ErrNotFound):
+		case err != nil:
+			return err
+		default:
+			var existing cooldownEntry
+			if json.Unmarshal(stored.Value, &existing) == nil && !replace(existing) {
+				return nil
+			}
+			if err := s.kv.Delete(ctx, id, stored.Revision); err != nil {
+				if errors.Is(err, ErrConflict) {
+					continue
+				}
+				return err
+			}
+		}
+		if _, err := s.kv.CreateWithTTL(ctx, id, value, ttl); err != nil {
+			if errors.Is(err, ErrConflict) {
+				continue
+			}
+			return err
+		}
+		return nil
+	}
+	return &ConflictError{Bucket: s.kv.bucket, ID: id}
+}

internal/store/natsstore/cooldown_store_test.go

@@ -0,0 +1,179 @@
+package natsstore
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	cliproxyauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth"
+)
+
+func cooldownRecord(authID, model string, updatedAt time.Time) cliproxyauth.CooldownStateRecord {
+	return cliproxyauth.CooldownStateRecord{
+		Provider:       "example",
+		AuthID:         authID,
+		Model:          model,
+		Status:         "error",
+		Reason:         "quota",
+		NextRetryAfter: time.Now().Add(10 * time.Minute),
+		UpdatedAt:      updatedAt,
+	}
+}
+
+func loadCooldowns(t *testing.T, s *Store) map[string]cliproxyauth.CooldownStateRecord {
+	t.Helper()
+	records, err := s.CooldownStateStore().Load(context.Background())
+	if err != nil {
+		t.Fatal(err)
+	}
+	byKey := make(map[string]cliproxyauth.CooldownStateRecord, len(records))
+	for _, record := range records {
+		byKey[record.AuthID+"/"+record.Model] = record
+	}
+	return byKey
+}
+
+func TestCooldownStateIsSharedBetweenReplicas(t *testing.T) {
+	f := newStoreFixture(t)
+	a, b := f.open(), f.open()
+	at := time.Now().Add(-time.Minute).UTC().Truncate(time.Microsecond)
+
+	if err := a.CooldownStateStore().Save(context.Background(), []cliproxyauth.CooldownStateRecord{
+		cooldownRecord("one", "model-a", at),
+		cooldownRecord("one", "", at),
+	}); err != nil {
+		t.Fatal(err)
+	}
+	got := loadCooldowns(t, b)
+	if len(got) != 2 {
+		t.Fatalf("loaded %d records, want 2: %+v", len(got), got)
+	}
+	record := got["one/model-a"]
+	if record.Reason != "quota" || !record.UpdatedAt.Equal(at) {
+		t.Fatalf("record = %+v, want reason quota updated at %v", record, at)
+	}
+}
+
+func TestCooldownSaveNeverReplacesANewerRecord(t *testing.T) {
+	f := newStoreFixture(t)
+	a, b := f.open(), f.open()
+	ctx := context.Background()
+	now := time.Now()
+
+	newer := cooldownRecord("one", "m", now)
+	newer.Reason = "newer"
+	if err := a.CooldownStateStore().Save(ctx, []cliproxyauth.CooldownStateRecord{newer}); err != nil {
+		t.Fatal(err)
+	}
+	older := cooldownRecord("one", "m", now.Add(-time.Hour))
+	older.Reason = "older"
+	if err := b.CooldownStateStore().Save(ctx, []cliproxyauth.CooldownStateRecord{older}); err != nil {
+		t.Fatal(err)
+	}
+	if got := loadCooldowns(t, a)["one/m"]; got.Reason != "newer" {
+		t.Fatalf("reason = %q, want the newer record to survive", got.Reason)
+	}
+
+	newest := cooldownRecord("one", "m", now.Add(time.Hour))
+	newest.Reason = "newest"
+	if err := b.CooldownStateStore().Save(ctx, []cliproxyauth.CooldownStateRecord{newest}); err != nil {
+		t.Fatal(err)
+	}
+	if got := loadCooldowns(t, a)["one/m"]; got.Reason != "newest" {
+		t.Fatalf("reason = %q, want a newer write to replace the record", got.Reason)
+	}
+}
+
+func TestCooldownSaveClearsRecordsItNoLongerHolds(t *testing.T) {
+	f := newStoreFixture(t)
+	a, b := f.open(), f.open()
+	ctx := context.Background()
+	at := time.Now().Add(-time.Minute)
+
+	if err := a.CooldownStateStore().Save(ctx, []cliproxyauth.CooldownStateRecord{cooldownRecord("one", "m", at), cooldownRecord("two", "m", at)}); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := a.CooldownStateStore().Load(ctx); err != nil {
+		t.Fatal(err)
+	}
+	if err := a.CooldownStateStore().Save(ctx, []cliproxyauth.CooldownStateRecord{cooldownRecord("two", "m", at.Add(time.Second))}); err != nil {
+		t.Fatal(err)
+	}
+	got := loadCooldowns(t, b)
+	if _, ok := got["one/m"]; ok || len(got) != 1 {
+		t.Fatalf("records = %+v, want only two/m", got)
+	}
+
+	// A replica that missed the clearing cannot bring the record back.
+	if err := b.CooldownStateStore().Save(ctx, []cliproxyauth.CooldownStateRecord{cooldownRecord("one", "m", at), cooldownRecord("two", "m", at.Add(time.Second))}); err != nil {
+		t.Fatal(err)
+	}
+	if _, ok := loadCooldowns(t, a)["one/m"]; ok {
+		t.Fatal("an older write resurrected a cleared record")
+	}
+}
+
+func TestCooldownClearLeavesARecordAnotherReplicaWrote(t *testing.T) {
+	f := newStoreFixture(t)
+	a, b := f.open(), f.open()
+	ctx := context.Background()
+	now := time.Now()
+
+	if err := a.CooldownStateStore().Save(ctx, []cliproxyauth.CooldownStateRecord{cooldownRecord("one", "m", now.Add(-time.Minute))}); err != nil {
+		t.Fatal(err)
+	}
+	if err := b.CooldownStateStore().Save(ctx, []cliproxyauth.CooldownStateRecord{cooldownRecord("one", "m", now)}); err != nil {
+		t.Fatal(err)
+	}
+	// a still remembers its own, older version and now drops it.
+	if err := a.CooldownStateStore().Save(ctx, nil); err != nil {
+		t.Fatal(err)
+	}
+	if _, ok := loadCooldowns(t, b)["one/m"]; !ok {
+		t.Fatal("clearing an old version removed a newer record")
+	}
+}
+
+func TestCooldownSaveRejectsAnEmptyAuthID(t *testing.T) {
+	f := newStoreFixture(t)
+	if err := f.open().CooldownStateStore().Save(context.Background(), []cliproxyauth.CooldownStateRecord{{AuthID: " "}}); err == nil {
+		t.Fatal("a record without an auth ID was saved")
+	}
+}
+
+func TestCooldownTTLFollowsTheRetryTime(t *testing.T) {
+	now := time.Now()
+	cases := []struct {
+		name string
+		next time.Time
+		want time.Duration
+	}{
+		{"future", now.Add(time.Hour), time.Hour + cooldownGrace},
+		{"passed", now.Add(-time.Hour), cooldownGrace},
+		{"unset", time.Time{}, cooldownIdleTTL},
+	}
+	for _, tc := range cases {
+		got := cooldownTTL(cliproxyauth.CooldownStateRecord{NextRetryAfter: tc.next}, now)
+		if got != tc.want {
+			t.Errorf("%s: ttl = %v, want %v", tc.name, got, tc.want)
+		}
+	}
+}
+
+func TestCooldownRecordExpiresOnTheServer(t *testing.T) {
+	f := newStoreFixture(t)
+	s := f.open()
+	ctx := context.Background()
+	record := cooldownRecord("one", "m", time.Now())
+	entry := cooldownEntry{UpdatedAt: record.UpdatedAt.UTC().Truncate(time.Microsecond), Record: &record}
+	if err := s.cooldown.put(ctx, cooldownID(record), entry, time.Second, func(cooldownEntry) bool { return true }); err != nil {
+		t.Fatal(err)
+	}
+	if _, ok := loadCooldowns(t, s)["one/m"]; !ok {
+		t.Fatal("record is not there before it expires")
+	}
+	waitFor(t, "record to expire", func() bool {
+		_, ok := loadCooldowns(t, s)["one/m"]
+		return !ok
+	})
+}