1package refusalcap
  2
  3import (
  4	"context"
  5	"encoding/json"
  6	"log/slog"
  7	"sync"
  8	"time"
  9
 10	"tug/internal/platform/origin"
 11	"tug/internal/platform/ratelimit"
 12	"tug/internal/storage"
 13)
 14
 15const (
 16	ActionRowsDropped = "audit.rows_dropped"
 17	DroppedSubject    = "audit"
 18	FlushEvery        = time.Minute
 19	stopTimeout       = 5 * time.Second
 20)
 21
 22var (
 23	RefusalRate = ratelimit.Rate{Burst: 20, Every: 3 * time.Second}
 24	SuccessRate = ratelimit.Rate{Burst: 60, Every: time.Second}
 25)
 26
 27type Appender interface {
 28	AppendAudit(context.Context, storage.AppendAuditInput) (storage.AuditEvent, error)
 29}
 30
 31type bucket uint8
 32
 33const (
 34	refusals bucket = iota
 35	successes
 36)
 37
 38type Cap struct {
 39	store   Appender
 40	logger  *slog.Logger
 41	buckets *ratelimit.Table[bucket]
 42	mu      sync.Mutex
 43	dropped map[string]int
 44}
 45
 46func New(store Appender, now func() time.Time, logger *slog.Logger) *Cap {
 47	if now == nil {
 48		now = time.Now
 49	}
 50	if logger == nil {
 51		logger = slog.Default()
 52	}
 53	return &Cap{store: store, logger: logger, buckets: ratelimit.New[bucket](2, now), dropped: map[string]int{}}
 54}
 55
 56func (c *Cap) Refused(ctx context.Context, input storage.AppendAuditInput) {
 57	c.write(ctx, refusals, RefusalRate, input)
 58}
 59
 60func (c *Cap) Succeeded(ctx context.Context, input storage.AppendAuditInput) {
 61	c.write(ctx, successes, SuccessRate, input)
 62}
 63
 64func (c *Cap) write(ctx context.Context, b bucket, rate ratelimit.Rate, input storage.AppendAuditInput) {
 65	if _, ok := c.buckets.Take(b, rate); !ok {
 66		c.mu.Lock()
 67		c.dropped[input.Action]++
 68		c.mu.Unlock()
 69		return
 70	}
 71	if _, err := c.store.AppendAudit(ctx, input); err != nil {
 72		c.logger.Error("audit row not written", "action", input.Action, "err", err)
 73	}
 74}
 75
 76func (c *Cap) Flush(ctx context.Context) error {
 77	c.mu.Lock()
 78	counts := c.dropped
 79	if len(counts) == 0 {
 80		c.mu.Unlock()
 81		return nil
 82	}
 83	c.dropped = map[string]int{}
 84	c.mu.Unlock()
 85	detail, err := json.Marshal(counts)
 86	if err == nil {
 87		_, err = c.store.AppendAudit(origin.With(ctx, origin.Client{}), storage.AppendAuditInput{Actor: "system", Action: ActionRowsDropped, Subject: DroppedSubject, Detail: string(detail)})
 88	}
 89	if err != nil {
 90		c.mu.Lock()
 91		for action, n := range counts {
 92			c.dropped[action] += n
 93		}
 94		c.mu.Unlock()
 95	}
 96	return err
 97}
 98
 99func (c *Cap) Run(ctx context.Context) {
100	tick := time.NewTicker(FlushEvery)
101	defer tick.Stop()
102	for {
103		select {
104		case <-ctx.Done():
105			stop, cancel := context.WithTimeout(context.WithoutCancel(ctx), stopTimeout)
106			defer cancel()
107			if err := c.Flush(stop); err != nil {
108				c.logger.Error("audit overflow row not written", "err", err)
109			}
110			return
111		case <-tick.C:
112			if err := c.Flush(ctx); err != nil {
113				c.logger.Error("audit overflow row not written", "err", err)
114			}
115		}
116	}
117}