broad-announce/internal/quarantine/quarantine.go

// Package quarantine provides per-source error-rate enforcement
// for the ingest pipeline (SPEC §22 layer 7). When a source
// generates too many errors within a sliding window, it is
// temporarily banned (quarantined) so it cannot submit alerts.
//
// Unlike the circuit breaker (layer 6, in-process), quarantine
// uses Redis so the ban is shared across all ingestd instances.
// A source banned by instance A is also banned by instance B.
//
// How it works:
//   - Every inbound alert that fails at any layer (rate limit,
//     circuit open, marshal error, publish error, etc.) records
//     one "hit" in Redis.
//   - A hit is a ZADD of (timestamp_ms as score, timestamp_ms as
//     value) into a sorted set keyed
//     "ingestd:quarantine_hits:<source_id>".
//   - A separate String key "ingestd:quarantine_banned:<source_id>"
//     with a TTL of QuarantineDuration holds the ban.
//   - On every ProcessAlert entry, we check the ban key first.
//     If present, reject immediately with result="quarantined".
//   - After every rejection (before returning from ProcessAlert),
//     we ZADD the hit. If the ban key is absent but the hit count
//     in the window >= QuarantineHitsThreshold, we SET the ban key
//     with TTL QuarantineDuration.
//
// This means a source accumulates hits continuously while it is
// quarantined — as soon as the ban expires it can immediately be
// re-quarantined if it keeps generating errors. This is intentional:
// a source that is broken does not get a "free pass" after the
// ban expires.
package quarantine

import (
	"context"
	"fmt"
	"strconv"
	"time"

	"github.com/redis/go-redis/v9"
)

// Config holds the static knobs. All fields must be set.
type Config struct {
	// HitsThreshold is the number of error hits within
	// HitsWindow that triggers a ban.
	HitsThreshold int
	// HitsWindow is the rolling window for counting error hits.
	HitsWindow time.Duration
	// BanDuration is how long a ban lasts once triggered.
	BanDuration time.Duration
}

// DefaultConfig is a reasonable starting point: 100 errors in
// 5 minutes triggers a 10-minute ban.
func DefaultConfig() Config {
	return Config{
		HitsThreshold: 100,
		HitsWindow:    5 * time.Minute,
		BanDuration:   10 * time.Minute,
	}
}

// Manager owns the Redis connection and provides the Check/Record
// API to the ingest pipeline. It is safe for concurrent use.
type Manager struct {
	config Config
	rdb    *redis.Client
}

// New creates a new quarantine manager.
func New(rdb *redis.Client, cfg Config) *Manager {
	return &Manager{config: cfg, rdb: rdb}
}

// IsBanned returns true and the remaining TTL if sourceID is
// currently banned. The remaining TTL is 0 if not banned.
func (q *Manager) IsBanned(ctx context.Context, sourceID string) (bool, time.Duration, error) {
	key := bannedKey(sourceID)
	ttl, err := q.rdb.TTL(ctx, key).Result()
	if err != nil && err != redis.Nil {
		return false, 0, fmt.Errorf("quarantine ttl check: %w", err)
	}
	if ttl > 0 {
		return true, ttl, nil
	}
	return false, 0, nil
}

// RecordHit records one error for sourceID. If the hit count
// in the window now exceeds HitsThreshold, the source is
// immediately banned for BanDuration.
//
// RecordHit is idempotent: calling it twice for the same
// timestamp counts as two hits (the hit timestamp is precise
// to the millisecond, so two calls in the same millisecond
// would need a unique member value to avoid being collapsed
// by the sorted set's score uniqueness — we append a random
// suffix to make each ZADD member unique).
func (q *Manager) RecordHit(ctx context.Context, sourceID string) error {
	hitKey := hitsKey(sourceID)
	banKey := bannedKey(sourceID)
	now := time.Now()
	nowMs := now.UnixMilli()
	member := strconv.FormatInt(nowMs, 10) + ":" + randMember()

	pipe := q.rdb.Pipeline()

	// Add the hit to the sorted set with score = nowMs.
	pipe.ZAdd(ctx, hitKey, redis.Z{Score: float64(nowMs), Member: member})
	// Expire the hits key after HitsWindow + BanDuration so the
	// set auto-cleans and doesn't grow forever.
	pipe.Expire(ctx, hitKey, q.config.HitsWindow+q.config.BanDuration+time.Minute)

	// Prune hits older than HitsWindow.
	cutoff := now.Add(-q.config.HitsWindow).UnixMilli()
	pipe.ZRemRangeByScore(ctx, hitKey, "-inf", strconv.FormatInt(cutoff, 10))

	// Count hits in window.
	count, err := pipe.ZCard(ctx, hitKey).Result()
	if err != nil {
		return fmt.Errorf("quarantine zcard: %w", err)
	}

	if int(count) >= q.config.HitsThreshold {
		// Trip the ban.
		pipe.Set(ctx, banKey, "1", q.config.BanDuration)
	}

	_, err = pipe.Exec(ctx)
	if err != nil {
		return fmt.Errorf("quarantine exec: %w", err)
	}

	return nil
}

// HitsInWindow returns the current number of error hits for
// sourceID within the configured hits window. Exposed for the
// metrics emission site.
func (q *Manager) HitsInWindow(ctx context.Context, sourceID string) (int, error) {
	hitKey := hitsKey(sourceID)
	now := time.Now()
	cutoff := now.Add(-q.config.HitsWindow).UnixMilli()
	count, err := q.rdb.ZCount(ctx, hitKey, strconv.FormatInt(cutoff, 10), "+inf").Result()
	if err != nil {
		return 0, fmt.Errorf("quarantine hits count: %w", err)
	}
	return int(count), nil
}

func hitsKey(sourceID string) string { return "ingestd:quarantine_hits:" + sourceID }
func bannedKey(sourceID string) string {
	return "ingestd:quarantine_banned:" + sourceID
}

// randMember returns a short random string to ensure ZADD
// members are unique when two calls happen in the same ms.
func randMember() string {
	// 6 hex chars is plenty for uniqueness within a process.
	b := make([]byte, 3)
	for i := range b {
		b[i] = byte(time.Now().UnixNano() >> (i * 8) & 0xff)
	}
	return fmt.Sprintf("%x", b)
}