starlog/rate_limit.go

package starlog

import (
	"math"
	"strconv"
	"strings"
	"sync"
	"time"
)

type RateLimitScope int

const (
	RateLimitScopeGlobal RateLimitScope = iota
	RateLimitScopeByKey
)

type RateLimitDropPolicy int

const (
	RateLimitDrop RateLimitDropPolicy = iota
	RateLimitPassThrough
)

type RateLimitDropData struct {
	Time               time.Time
	Key                string
	Reason             string
	Level              int
	LevelName          string
	LoggerName         string
	Message            string
	PassThrough        bool
	DroppedCount       uint64
	PassedThroughCount uint64
	Suppressed         uint64
	Summary            bool
	SummarySuppressed  uint64
}

type RateLimitStats struct {
	Enabled       bool
	Rate          float64
	Burst         int
	Scope         RateLimitScope
	Allowed       uint64
	Dropped       uint64
	PassedThrough uint64
	Suppressed    uint64
	Summaries     uint64
	LastDropTime  time.Time
	LastDropKey   string
	LastReason    string
	CurrentKeys   int
}

type RateLimitConfig struct {
	Enable bool

	Levels       []int
	Rate         float64
	Burst        int
	Scope        RateLimitScope
	KeyFunc      func(*Entry) string
	MaxKeys      int
	KeyTTL       time.Duration
	DropPolicy   RateLimitDropPolicy
	OnDrop       func(RateLimitDropData)
	ExemptLevels []int
	ExemptMatcher func(*Entry) bool

	SummaryInterval time.Duration
}

type rateLimitBucket struct {
	tokens     float64
	lastRefill time.Time
	lastSeen   time.Time
}

type rateLimiter struct {
	mu sync.Mutex

	cfg          RateLimitConfig
	limitedLevel map[int]struct{}
	exemptLevel  map[int]struct{}
	buckets      map[string]*rateLimitBucket

	nowFunc            func() time.Time
	lastCleanupTime    time.Time
	lastSummaryTime    time.Time
	allowedCount       uint64
	droppedCount       uint64
	passedThroughCount uint64
	suppressedCount    uint64
	summaryCount       uint64
	lastDropTime       time.Time
	lastDropKey        string
	lastDropReason     string
}

func DefaultRateLimitConfig() RateLimitConfig {
	return RateLimitConfig{
		Enable:          false,
		Levels:          nil,
		Rate:            0,
		Burst:           1,
		Scope:           RateLimitScopeGlobal,
		KeyFunc:         nil,
		MaxKeys:         4096,
		KeyTTL:          10 * time.Minute,
		DropPolicy:      RateLimitDrop,
		OnDrop:          nil,
		ExemptLevels:    []int{LvError, LvCritical, LvPanic, LvFatal},
		ExemptMatcher:   nil,
		SummaryInterval: 0,
	}
}

func cloneIntSlice(values []int) []int {
	if len(values) == 0 {
		return nil
	}
	result := make([]int, len(values))
	copy(result, values)
	return result
}

func cloneRateLimitConfig(cfg RateLimitConfig) RateLimitConfig {
	cloned := cfg
	cloned.Levels = cloneIntSlice(cfg.Levels)
	cloned.ExemptLevels = cloneIntSlice(cfg.ExemptLevels)
	return cloned
}

func normalizeRateLimitConfig(cfg RateLimitConfig) RateLimitConfig {
	if cfg.Rate < 0 {
		cfg.Rate = 0
	}
	if cfg.Burst <= 0 {
		if cfg.Rate > 0 {
			cfg.Burst = int(math.Ceil(cfg.Rate))
		}
		if cfg.Burst <= 0 {
			cfg.Burst = 1
		}
	}
	switch cfg.Scope {
	case RateLimitScopeGlobal, RateLimitScopeByKey:
	default:
		cfg.Scope = RateLimitScopeGlobal
	}
	switch cfg.DropPolicy {
	case RateLimitDrop, RateLimitPassThrough:
	default:
		cfg.DropPolicy = RateLimitDrop
	}
	if cfg.MaxKeys <= 0 {
		cfg.MaxKeys = 4096
	}
	if cfg.KeyTTL <= 0 {
		cfg.KeyTTL = 10 * time.Minute
	}
	if cfg.SummaryInterval < 0 {
		cfg.SummaryInterval = 0
	}
	return cloneRateLimitConfig(cfg)
}

func buildLevelSet(levels []int) map[int]struct{} {
	result := make(map[int]struct{}, len(levels))
	for _, level := range levels {
		result[level] = struct{}{}
	}
	return result
}

func newRateLimiter() *rateLimiter {
	cfg := normalizeRateLimitConfig(DefaultRateLimitConfig())
	return &rateLimiter{
		cfg:          cfg,
		limitedLevel: buildLevelSet(cfg.Levels),
		exemptLevel:  buildLevelSet(cfg.ExemptLevels),
		buckets:      make(map[string]*rateLimitBucket),
		nowFunc:      time.Now,
	}
}

func (limiter *rateLimiter) setNowFuncForTest(nowFunc func() time.Time) {
	if limiter == nil {
		return
	}
	limiter.mu.Lock()
	if nowFunc == nil {
		limiter.nowFunc = time.Now
	} else {
		limiter.nowFunc = nowFunc
	}
	limiter.mu.Unlock()
}

func (limiter *rateLimiter) now() time.Time {
	if limiter == nil || limiter.nowFunc == nil {
		return time.Now()
	}
	return limiter.nowFunc()
}

func (limiter *rateLimiter) SetConfig(cfg RateLimitConfig) {
	if limiter == nil {
		return
	}
	normalized := normalizeRateLimitConfig(cfg)
	limiter.mu.Lock()
	limiter.cfg = normalized
	limiter.limitedLevel = buildLevelSet(normalized.Levels)
	limiter.exemptLevel = buildLevelSet(normalized.ExemptLevels)
	limiter.buckets = make(map[string]*rateLimitBucket)
	limiter.lastCleanupTime = time.Time{}
	limiter.lastSummaryTime = time.Time{}
	limiter.mu.Unlock()
}

func (limiter *rateLimiter) Config() RateLimitConfig {
	if limiter == nil {
		return normalizeRateLimitConfig(DefaultRateLimitConfig())
	}
	limiter.mu.Lock()
	defer limiter.mu.Unlock()
	return cloneRateLimitConfig(limiter.cfg)
}

func (limiter *rateLimiter) Stats() RateLimitStats {
	if limiter == nil {
		return RateLimitStats{}
	}
	limiter.mu.Lock()
	defer limiter.mu.Unlock()
	return RateLimitStats{
		Enabled:       limiter.cfg.Enable,
		Rate:          limiter.cfg.Rate,
		Burst:         limiter.cfg.Burst,
		Scope:         limiter.cfg.Scope,
		Allowed:       limiter.allowedCount,
		Dropped:       limiter.droppedCount,
		PassedThrough: limiter.passedThroughCount,
		Suppressed:    limiter.suppressedCount,
		Summaries:     limiter.summaryCount,
		LastDropTime:  limiter.lastDropTime,
		LastDropKey:   limiter.lastDropKey,
		LastReason:    limiter.lastDropReason,
		CurrentKeys:   len(limiter.buckets),
	}
}

func (limiter *rateLimiter) ResetStats() {
	if limiter == nil {
		return
	}
	limiter.mu.Lock()
	limiter.allowedCount = 0
	limiter.droppedCount = 0
	limiter.passedThroughCount = 0
	limiter.suppressedCount = 0
	limiter.summaryCount = 0
	limiter.lastDropTime = time.Time{}
	limiter.lastDropKey = ""
	limiter.lastDropReason = ""
	limiter.mu.Unlock()
}

func (limiter *rateLimiter) isLimitedLevel(level int) bool {
	if len(limiter.limitedLevel) == 0 {
		return true
	}
	_, ok := limiter.limitedLevel[level]
	return ok
}

func (limiter *rateLimiter) isExempt(entry *Entry) bool {
	if entry == nil {
		return false
	}
	if _, ok := limiter.exemptLevel[entry.Level]; ok {
		return true
	}
	if limiter.cfg.ExemptMatcher != nil && limiter.cfg.ExemptMatcher(entry) {
		return true
	}
	return false
}

func (limiter *rateLimiter) resolveKey(entry *Entry) string {
	if limiter.cfg.Scope == RateLimitScopeGlobal {
		return "__global__"
	}
	if limiter.cfg.KeyFunc != nil {
		key := strings.TrimSpace(limiter.cfg.KeyFunc(entry))
		if key != "" {
			return key
		}
	}
	if entry == nil {
		return "__empty__"
	}
	message := strings.TrimSpace(entry.Message)
	if message == "" {
		return strconv.Itoa(entry.Level)
	}
	return strconv.Itoa(entry.Level) + ":" + message
}

func (limiter *rateLimiter) cleanupBucketsLocked(now time.Time) {
	if limiter.cfg.Scope != RateLimitScopeByKey || limiter.cfg.KeyTTL <= 0 {
		return
	}
	if !limiter.lastCleanupTime.IsZero() && now.Sub(limiter.lastCleanupTime) < time.Second {
		return
	}
	for key, bucket := range limiter.buckets {
		if bucket == nil {
			delete(limiter.buckets, key)
			continue
		}
		if now.Sub(bucket.lastSeen) > limiter.cfg.KeyTTL {
			delete(limiter.buckets, key)
		}
	}
	limiter.lastCleanupTime = now
}

func (limiter *rateLimiter) getBucketLocked(key string, now time.Time) *rateLimitBucket {
	if bucket, ok := limiter.buckets[key]; ok && bucket != nil {
		return bucket
	}

	if limiter.cfg.Scope == RateLimitScopeByKey && limiter.cfg.MaxKeys > 0 && len(limiter.buckets) >= limiter.cfg.MaxKeys {
		overflowKey := "__overflow__"
		if bucket, ok := limiter.buckets[overflowKey]; ok && bucket != nil {
			return bucket
		}
		oldestKey := ""
		oldestTime := now
		for existingKey, bucket := range limiter.buckets {
			if bucket == nil {
				oldestKey = existingKey
				break
			}
			if oldestKey == "" || bucket.lastSeen.Before(oldestTime) {
				oldestKey = existingKey
				oldestTime = bucket.lastSeen
			}
		}
		if oldestKey != "" {
			delete(limiter.buckets, oldestKey)
		}
		key = overflowKey
	}

	bucket := &rateLimitBucket{
		tokens:     float64(limiter.cfg.Burst),
		lastRefill: now,
		lastSeen:   now,
	}
	limiter.buckets[key] = bucket
	return bucket
}

func (limiter *rateLimiter) takeTokenLocked(bucket *rateLimitBucket, now time.Time) bool {
	if bucket == nil {
		return true
	}
	if bucket.lastRefill.IsZero() {
		bucket.lastRefill = now
	}
	elapsed := now.Sub(bucket.lastRefill).Seconds()
	if elapsed > 0 && limiter.cfg.Rate > 0 {
		bucket.tokens += elapsed * limiter.cfg.Rate
		maxTokens := float64(limiter.cfg.Burst)
		if bucket.tokens > maxTokens {
			bucket.tokens = maxTokens
		}
		bucket.lastRefill = now
	}
	bucket.lastSeen = now
	if bucket.tokens >= 1 {
		bucket.tokens -= 1
		return true
	}
	return false
}

func cloneEntryForDrop(entry *Entry) Entry {
	if entry == nil {
		return Entry{}
	}
	cloned := *entry
	cloned.Fields = cloneFields(entry.Fields)
	return cloned
}

func (limiter *rateLimiter) Allow(entry *Entry) bool {
	if limiter == nil || entry == nil {
		return true
	}
	now := limiter.now()
	var callback func(RateLimitDropData)
	dropData := RateLimitDropData{}
	allow := true

	limiter.mu.Lock()
	if !limiter.cfg.Enable || limiter.cfg.Rate <= 0 {
		limiter.allowedCount++
		limiter.mu.Unlock()
		return true
	}
	if !limiter.isLimitedLevel(entry.Level) || limiter.isExempt(entry) {
		limiter.allowedCount++
		limiter.mu.Unlock()
		return true
	}

	limiter.cleanupBucketsLocked(now)
	key := limiter.resolveKey(entry)
	bucket := limiter.getBucketLocked(key, now)
	if limiter.takeTokenLocked(bucket, now) {
		limiter.allowedCount++
		limiter.mu.Unlock()
		return true
	}

	limiter.suppressedCount++
	limiter.lastDropTime = now
	limiter.lastDropKey = key
	limiter.lastDropReason = "rate_limit_exceeded"

	if limiter.cfg.DropPolicy == RateLimitPassThrough {
		limiter.passedThroughCount++
		limiter.allowedCount++
		allow = true
	} else {
		limiter.droppedCount++
		allow = false
	}

	dropData = RateLimitDropData{
		Time:               now,
		Key:                key,
		Reason:             limiter.lastDropReason,
		Level:              entry.Level,
		LevelName:          entry.LevelName,
		LoggerName:         entry.LoggerName,
		Message:            entry.Message,
		PassThrough:        limiter.cfg.DropPolicy == RateLimitPassThrough,
		DroppedCount:       limiter.droppedCount,
		PassedThroughCount: limiter.passedThroughCount,
		Suppressed:         limiter.suppressedCount,
	}

	if limiter.cfg.SummaryInterval > 0 {
		if limiter.lastSummaryTime.IsZero() {
			limiter.lastSummaryTime = now
		} else if now.Sub(limiter.lastSummaryTime) >= limiter.cfg.SummaryInterval {
			dropData.Summary = true
			dropData.SummarySuppressed = limiter.suppressedCount
			limiter.summaryCount++
			limiter.suppressedCount = 0
			limiter.lastSummaryTime = now
		}
	}

	callback = limiter.cfg.OnDrop
	limiter.mu.Unlock()

	if callback != nil {
		entryCopy := cloneEntryForDrop(entry)
		dropData.Level = entryCopy.Level
		dropData.LevelName = entryCopy.LevelName
		dropData.LoggerName = entryCopy.LoggerName
		dropData.Message = entryCopy.Message
		func() {
			defer func() {
				recover()
			}()
			callback(dropData)
		}()
	}
	return allow
}

func (logger *starlog) allowByRateLimit(entry *Entry) bool {
	if logger == nil || logger.rateLimiter == nil {
		return true
	}
	return logger.rateLimiter.Allow(entry)
}

func (logger *StarLogger) SetRateLimitConfig(cfg RateLimitConfig) {
	if logger == nil || logger.logcore == nil {
		return
	}
	logger.logcore.mu.Lock()
	if logger.logcore.rateLimiter == nil {
		logger.logcore.rateLimiter = newRateLimiter()
	}
	limiter := logger.logcore.rateLimiter
	logger.logcore.mu.Unlock()
	limiter.SetConfig(cfg)
}

func (logger *StarLogger) GetRateLimitConfig() RateLimitConfig {
	if logger == nil || logger.logcore == nil {
		return normalizeRateLimitConfig(DefaultRateLimitConfig())
	}
	logger.logcore.mu.Lock()
	limiter := logger.logcore.rateLimiter
	logger.logcore.mu.Unlock()
	if limiter == nil {
		return normalizeRateLimitConfig(DefaultRateLimitConfig())
	}
	return limiter.Config()
}

func (logger *StarLogger) EnableRateLimit(enable bool) {
	cfg := logger.GetRateLimitConfig()
	cfg.Enable = enable
	logger.SetRateLimitConfig(cfg)
}

func (logger *StarLogger) SetRateLimitDropHandler(handler func(RateLimitDropData)) {
	cfg := logger.GetRateLimitConfig()
	cfg.OnDrop = handler
	logger.SetRateLimitConfig(cfg)
}

func (logger *StarLogger) GetRateLimitStats() RateLimitStats {
	if logger == nil || logger.logcore == nil {
		return RateLimitStats{}
	}
	logger.logcore.mu.Lock()
	limiter := logger.logcore.rateLimiter
	logger.logcore.mu.Unlock()
	if limiter == nil {
		return RateLimitStats{}
	}
	return limiter.Stats()
}

func (logger *StarLogger) ResetRateLimitStats() {
	if logger == nil || logger.logcore == nil {
		return
	}
	logger.logcore.mu.Lock()
	limiter := logger.logcore.rateLimiter
	logger.logcore.mu.Unlock()
	if limiter == nil {
		return
	}
	limiter.ResetStats()
}