Cyrene_v2/internal/channel/adapter/flowgate.go

package adapter

import (
	"context"
	"sync"
	"time"
)

// Reviewer decides which buffered messages are worth keeping. It receives the
// current pending set (in arrival order) and returns the indices it wants to keep.
// Messages at a returned index survive; all others are pruned and audited.
//
// It is a pure noise/importance filter (design/12 §6 策略四). Implementations must
// never decide identity: who the sender is was already resolved structurally by the
// channel layer. A nil Reviewer makes the gate fall back to a plain FIFO cap.
type Reviewer func(ctx context.Context, msgs []UnifiedMessage) (keep []int, err error)

// GateVerdict reports how the gate classified one inbound message.
type GateVerdict struct {
	// Proceed reports whether the message should be delivered downstream.
	Proceed bool
	// Reason is a short token describing a drop: "rate_limited", "review_pruned"
	// or "capped". It is empty when Proceed is true.
	Reason string
	// DroppedMsg is the message that failed admission (nil when Proceed is true).
	DroppedMsg *UnifiedMessage
	// Pruned lists every message the gate dropped during this admission. On the
	// review path it is the noise the reviewer rejected; on the cap path it is the
	// oldest messages evicted to honor the total cap. Each is recorded to the
	// OnPruned audit hook before being dropped, so nothing is lost silently.
	Pruned []UnifiedMessage
}

// FlowGateOptions configures a per-adapter ingress gate. Every limit is
// independent and a zero value disables that limit, so an account that does not
// declare the new fields keeps its previous unlocked behavior (backward
// compatible, design/12 §6 实现要点).
type FlowGateOptions struct {
	// GroupMsgPerMin caps group-chat inbound messages per minute (design/12 §6
	// 策略三). Zero disables the group limit.
	GroupMsgPerMin int
	// MsgPerMin caps total inbound non-owner messages per minute. Zero disables it.
	MsgPerMin int
	// MaxMessages is the total pending-message cap. When the pending set nears it
	// the Reviewer is invoked to prune noise; with no Reviewer the oldest messages
	// are evicted FIFO. Zero disables the cap (design/12 §6 策略四).
	MaxMessages int
	// Window is the sliding window over which the per-minute limits are counted.
	// Zero uses defaultRateWindow.
	Window time.Duration
	// Review is the optional review model. When nil the gate enforces the cap with
	// FIFO eviction instead of model judgment.
	Review Reviewer
	// OnPruned is called once per pruned/dropped message with a short reason token
	// ("review_pruned" or "capped") so the caller can record an audit trail.
	// It may be nil (then the caller only sees Pruned on the verdict).
	OnPruned func(UnifiedMessage, string)
	// Now overrides the clock for tests. Nil uses time.Now.
	Now func() time.Time
}

const (
	// defaultRateWindow is the sliding window over which per-minute limits count.
	defaultRateWindow = time.Minute
	// defaultReviewPct is how full the pending set must be before the review model
	// runs. Design/12 §6 策略四 says "maxMessages 的 80%": this is the cost guard
	// that keeps review from firing on every message.
	defaultReviewPct = 80
	// FlowLimitDefault is the protective default of 6 group messages per minute
	// (design/12 §6 策略三). It is applied when an account pins no group rate limit.
	FlowLimitDefault = 6
	// FlowCapDefault is the protective default total-cap of 50 pending messages
	// (design/12 §6 策略四).
	FlowCapDefault = 50
)

// FlowGate is a concurrency-safe, per-adapter ingress gate implementing the two
// measures from design/12 §6 策略三/四:
//
//   - 策略三: a sliding-window rate limit on group chat (and, optionally, all
//     non-owner messages). Over-limit messages are dropped (downgraded) rather
//     than blocking a flood of messages from saturating the loop.
//   - 策略四: a bounded pending-message set. When it nears its cap the Reviewer is
//     invoked to prune low-value messages; with no Reviewer the oldest messages are
//     evicted FIFO. Either way every dropped message is audited.
//
// Owner-private messages are not routed through the gate (the caller bypasses it),
// so the owner is never throttled or pruned. Identity is never examined here.
type FlowGate struct {
	groupPerMin int
	msgPerMin   int
	maxMessages int
	reviewAt    int
	window      time.Duration
	review      Reviewer
	onPruned    func(UnifiedMessage, string)
	now         func() time.Time

	mu         sync.Mutex
	groupHits  []time.Time
	totalHits  []time.Time
	buffer     []queuedMsg
	lastReview time.Time
	seq        uint64
}

// queuedMsg wraps a buffer entry with a stable sequence so pruning decisions can
// be audited and the fate of the just-admitted message be tracked across index
// shifts (value equality on UnifiedMessage is unsafe because it holds slices).
type queuedMsg struct {
	seq uint64
	msg UnifiedMessage
}

// NewFlowGate builds a gate from opts. Options that are zero disable their limit,
// except Window which defaults to one minute when left zero.
func NewFlowGate(opts FlowGateOptions) *FlowGate {
	g := &FlowGate{
		groupPerMin: opts.GroupMsgPerMin,
		msgPerMin:   opts.MsgPerMin,
		maxMessages: opts.MaxMessages,
		window:      opts.Window,
		review:      opts.Review,
		onPruned:    opts.OnPruned,
		now:         opts.Now,
	}
	if g.window <= 0 {
		g.window = defaultRateWindow
	}
	if g.maxMessages > 0 {
		g.reviewAt = g.maxMessages * defaultReviewPct / 100
		if g.reviewAt < 1 {
			g.reviewAt = 1
		}
	}
	if g.now == nil {
		g.now = time.Now
	}
	return g
}

// Admit applies the rate limit then the total-cap/review gate to one inbound
// message. It returns a verdict telling the caller whether to proceed; messages
// that the gate drops are returned (and audited) so nothing is lost silently.
func (g *FlowGate) Admit(ctx context.Context, msg UnifiedMessage) GateVerdict {
	now := g.now()
	g.mu.Lock()
	defer g.mu.Unlock()

	// 策略三: sliding-window rate limits. Group chat has its own cap; MsgPerMin is
	// an optional total cap across all non-owner messages. Over-limit messages are
	// dropped at the door (downgraded) instead of blocking.
	if g.groupPerMin > 0 && msg.ChannelType == ChannelGroup {
		g.groupHits = g.trimHits(g.groupHits, now)
		if len(g.groupHits) >= g.groupPerMin {
			return g.drop("rate_limited", msg)
		}
		g.groupHits = append(g.groupHits, now)
	}
	if g.msgPerMin > 0 {
		g.totalHits = g.trimHits(g.totalHits, now)
		if len(g.totalHits) >= g.msgPerMin {
			return g.drop("rate_limited", msg)
		}
		g.totalHits = append(g.totalHits, now)
	}

	if g.maxMessages <= 0 {
		return GateVerdict{Proceed: true}
	}

	// 策略四: add to the pending set, then prune when it nears the cap.
	g.seq++
	seq := g.seq
	g.buffer = append(g.buffer, queuedMsg{seq: seq, msg: msg})

	var dropped []queuedMsg
	if g.review != nil && len(g.buffer) >= g.reviewAt && !g.reviewRecently(now) {
		keep, err := g.review(ctx, queueMessages(g.buffer))
		if err == nil {
			g.buffer, dropped = g.applyReview(g.buffer, keep)
			g.lastReview = now
			g.auditDropped(dropped, "review_pruned")
		}
	}
	if len(g.buffer) > g.maxMessages {
		n := len(g.buffer) - g.maxMessages
		var evicted []queuedMsg
		g.buffer, evicted = g.dropFront(g.buffer, n)
		dropped = append(dropped, evicted...)
		g.auditDropped(evicted, "capped")
	}

	if !g.containsSeq(seq) {
		reason := "capped"
		if containsSeq(dropped, seq) {
			reason = "review_pruned"
		}
		v := GateVerdict{Proceed: false, Reason: reason, Pruned: queueMessages(dropped)}
		c := msg
		v.DroppedMsg = &c
		return v
	}
	return GateVerdict{Proceed: true, Pruned: queueMessages(dropped)}
}

// drop builds a refused verdict from a single message that never entered the
// buffer (rate-limited at the door), so the audit hook still sees it.
func (g *FlowGate) drop(reason string, msg UnifiedMessage) GateVerdict {
	c := msg
	v := GateVerdict{Proceed: false, Reason: reason, DroppedMsg: &c, Pruned: []UnifiedMessage{msg}}
	g.auditMsg(msg, reason)
	return v
}

// reviewRecently throttles review so it does not run on every message: it fires at
// most once per window, which is the cost guard for the review model.
func (g *FlowGate) reviewRecently(now time.Time) bool {
	if g.lastReview.IsZero() {
		return false
	}
	return now.Sub(g.lastReview) < g.window
}

// trimHits drops event timestamps older than the sliding window.
func (g *FlowGate) trimHits(hits []time.Time, now time.Time) []time.Time {
	cutoff := now.Add(-g.window)
	keep := 0
	for _, t := range hits {
		if t.After(cutoff) {
			hits[keep] = t
			keep++
		}
	}
	return hits[:keep]
}

// applyReview keeps only the entries at the returned indices (validated, de-duped,
// ascending) and returns the kept slice plus the dropped entries in original order.
func (g *FlowGate) applyReview(buffer []queuedMsg, keep []int) (kept, dropped []queuedMsg) {
	seen := make(map[int]bool, len(keep))
	for _, i := range keep {
		if i >= 0 && i < len(buffer) {
			seen[i] = true
		}
	}
	out := make([]queuedMsg, 0, len(buffer))
	drop := make([]queuedMsg, 0)
	for i, entry := range buffer {
		if seen[i] {
			out = append(out, entry)
		} else {
			drop = append(drop, entry)
		}
	}
	return out, drop
}

// dropFront removes n entries from the front (oldest) and returns the rest plus the
// dropped ones.
func (g *FlowGate) dropFront(buffer []queuedMsg, n int) (kept, dropped []queuedMsg) {
	if n >= len(buffer) {
		return nil, append([]queuedMsg(nil), buffer...)
	}
	return buffer[n:], append([]queuedMsg(nil), buffer[:n]...)
}

// containsSeq reports whether an entry with the given sequence is still buffered.
func (g *FlowGate) containsSeq(seq uint64) bool {
	for _, e := range g.buffer {
		if e.seq == seq {
			return true
		}
	}
	return false
}

// containsSeq reports whether seq appears among the dropped entries.
func containsSeq(dropped []queuedMsg, seq uint64) bool {
	for _, e := range dropped {
		if e.seq == seq {
			return true
		}
	}
	return false
}

// auditDropped records every dropped entry to the OnPruned audit hook.
func (g *FlowGate) auditDropped(dropped []queuedMsg, reason string) {
	for _, e := range dropped {
		g.auditMsg(e.msg, reason)
	}
}

// auditMsg records one dropped message to the OnPruned audit hook.
func (g *FlowGate) auditMsg(msg UnifiedMessage, reason string) {
	if g.onPruned != nil {
		g.onPruned(msg, reason)
	}
}

// queueMessages extracts the payload messages of queued entries in order.
func queueMessages(queue []queuedMsg) []UnifiedMessage {
	out := make([]UnifiedMessage, 0, len(queue))
	for _, e := range queue {
		out = append(out, e.msg)
	}
	return out
}