starlog/sample_dedup.go

package starlog

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

type SamplingScope int

const (
	SamplingScopeGlobal SamplingScope = iota
	SamplingScopeByKey
)

type SamplingDropData struct {
	Time        time.Time
	Key         string
	Reason      string
	Level       int
	LevelName   string
	LoggerName  string
	Message     string
	Rate        float64
	Allowed     uint64
	Dropped     uint64
	CurrentKeys int
}

type SamplingStats struct {
	Enabled      bool
	Rate         float64
	Scope        SamplingScope
	Allowed      uint64
	Dropped      uint64
	LastDropTime time.Time
	LastDropKey  string
	LastReason   string
	CurrentKeys  int
}

type SamplingConfig struct {
	Enable bool

	Levels []int
	Rate   float64
	Scope  SamplingScope

	KeyFunc func(*Entry) string
	MaxKeys int
	KeyTTL  time.Duration

	OnDrop func(SamplingDropData)
}

type samplingBucket struct {
	allowance float64
	lastSeen  time.Time
}

type sampler struct {
	mu sync.Mutex

	cfg          SamplingConfig
	limitedLevel map[int]struct{}
	buckets      map[string]*samplingBucket
	nowFunc      func() time.Time

	allowedCount    uint64
	droppedCount    uint64
	lastDropTime    time.Time
	lastDropKey     string
	lastDropReason  string
	lastCleanupTime time.Time
}

func DefaultSamplingConfig() SamplingConfig {
	return SamplingConfig{
		Enable:  false,
		Levels:  nil,
		Rate:    1,
		Scope:   SamplingScopeGlobal,
		KeyFunc: nil,
		MaxKeys: 4096,
		KeyTTL:  10 * time.Minute,
		OnDrop:  nil,
	}
}

func cloneSamplingConfig(cfg SamplingConfig) SamplingConfig {
	cloned := cfg
	cloned.Levels = cloneIntSlice(cfg.Levels)
	return cloned
}

func normalizeSamplingConfig(cfg SamplingConfig) SamplingConfig {
	if cfg.Rate < 0 {
		cfg.Rate = 0
	}
	if cfg.Rate > 1 {
		cfg.Rate = 1
	}
	switch cfg.Scope {
	case SamplingScopeGlobal, SamplingScopeByKey:
	default:
		cfg.Scope = SamplingScopeGlobal
	}
	if cfg.MaxKeys <= 0 {
		cfg.MaxKeys = 4096
	}
	if cfg.KeyTTL <= 0 {
		cfg.KeyTTL = 10 * time.Minute
	}
	return cloneSamplingConfig(cfg)
}

func newSampler() *sampler {
	cfg := normalizeSamplingConfig(DefaultSamplingConfig())
	return &sampler{
		cfg:          cfg,
		limitedLevel: buildLevelSet(cfg.Levels),
		buckets:      make(map[string]*samplingBucket),
		nowFunc:      time.Now,
	}
}

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

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

func (s *sampler) SetConfig(cfg SamplingConfig) {
	if s == nil {
		return
	}
	normalized := normalizeSamplingConfig(cfg)
	s.mu.Lock()
	s.cfg = normalized
	s.limitedLevel = buildLevelSet(normalized.Levels)
	s.buckets = make(map[string]*samplingBucket)
	s.lastCleanupTime = time.Time{}
	s.mu.Unlock()
}

func (s *sampler) Config() SamplingConfig {
	if s == nil {
		return normalizeSamplingConfig(DefaultSamplingConfig())
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	return cloneSamplingConfig(s.cfg)
}

func (s *sampler) Stats() SamplingStats {
	if s == nil {
		return SamplingStats{}
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	return SamplingStats{
		Enabled:      s.cfg.Enable,
		Rate:         s.cfg.Rate,
		Scope:        s.cfg.Scope,
		Allowed:      s.allowedCount,
		Dropped:      s.droppedCount,
		LastDropTime: s.lastDropTime,
		LastDropKey:  s.lastDropKey,
		LastReason:   s.lastDropReason,
		CurrentKeys:  len(s.buckets),
	}
}

func (s *sampler) ResetStats() {
	if s == nil {
		return
	}
	s.mu.Lock()
	s.allowedCount = 0
	s.droppedCount = 0
	s.lastDropTime = time.Time{}
	s.lastDropKey = ""
	s.lastDropReason = ""
	s.mu.Unlock()
}

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

func (s *sampler) resolveKey(entry *Entry) string {
	if s.cfg.Scope == SamplingScopeGlobal {
		return "__global__"
	}
	if s.cfg.KeyFunc != nil {
		key := strings.TrimSpace(s.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 (s *sampler) cleanupBucketsLocked(now time.Time) {
	if s.cfg.Scope != SamplingScopeByKey || s.cfg.KeyTTL <= 0 {
		return
	}
	if !s.lastCleanupTime.IsZero() && now.Sub(s.lastCleanupTime) < time.Second {
		return
	}
	for key, bucket := range s.buckets {
		if bucket == nil {
			delete(s.buckets, key)
			continue
		}
		if now.Sub(bucket.lastSeen) > s.cfg.KeyTTL {
			delete(s.buckets, key)
		}
	}
	s.lastCleanupTime = now
}

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

	if s.cfg.Scope == SamplingScopeByKey && s.cfg.MaxKeys > 0 && len(s.buckets) >= s.cfg.MaxKeys {
		oldestKey := ""
		oldestTime := now
		for existingKey, bucket := range s.buckets {
			if bucket == nil {
				oldestKey = existingKey
				break
			}
			if oldestKey == "" || bucket.lastSeen.Before(oldestTime) {
				oldestKey = existingKey
				oldestTime = bucket.lastSeen
			}
		}
		if oldestKey != "" {
			delete(s.buckets, oldestKey)
		}
	}

	bucket := &samplingBucket{
		allowance: 1,
		lastSeen:  now,
	}
	s.buckets[key] = bucket
	return bucket
}

func (s *sampler) Allow(entry *Entry) bool {
	if s == nil || entry == nil {
		return true
	}
	now := s.now()
	var callback func(SamplingDropData)
	dropData := SamplingDropData{}
	allow := true

	s.mu.Lock()
	if !s.cfg.Enable {
		s.allowedCount++
		s.mu.Unlock()
		return true
	}
	if !s.isLimitedLevel(entry.Level) {
		s.allowedCount++
		s.mu.Unlock()
		return true
	}
	if s.cfg.Rate >= 1 {
		s.allowedCount++
		s.mu.Unlock()
		return true
	}

	key := s.resolveKey(entry)
	if s.cfg.Rate <= 0 {
		s.droppedCount++
		s.lastDropTime = now
		s.lastDropKey = key
		s.lastDropReason = "sampling_rate_zero"
		allow = false
		dropData = SamplingDropData{
			Time:        now,
			Key:         key,
			Reason:      s.lastDropReason,
			Level:       entry.Level,
			LevelName:   entry.LevelName,
			LoggerName:  entry.LoggerName,
			Message:     entry.Message,
			Rate:        s.cfg.Rate,
			Allowed:     s.allowedCount,
			Dropped:     s.droppedCount,
			CurrentKeys: len(s.buckets),
		}
		callback = s.cfg.OnDrop
		s.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
	}

	s.cleanupBucketsLocked(now)
	bucket := s.getBucketLocked(key, now)
	if bucket.allowance >= 1 {
		allow = true
		bucket.allowance -= 1
		s.allowedCount++
	} else {
		allow = false
		s.droppedCount++
		s.lastDropTime = now
		s.lastDropKey = key
		s.lastDropReason = "sampling_rate_exceeded"
		dropData = SamplingDropData{
			Time:        now,
			Key:         key,
			Reason:      s.lastDropReason,
			Level:       entry.Level,
			LevelName:   entry.LevelName,
			LoggerName:  entry.LoggerName,
			Message:     entry.Message,
			Rate:        s.cfg.Rate,
			Allowed:     s.allowedCount,
			Dropped:     s.droppedCount,
			CurrentKeys: len(s.buckets),
		}
		callback = s.cfg.OnDrop
	}
	bucket.allowance += s.cfg.Rate
	if bucket.allowance > 1 {
		bucket.allowance = 1
	}
	bucket.lastSeen = now
	s.mu.Unlock()

	if !allow && 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) allowBySampling(entry *Entry) bool {
	if logger == nil || logger.sampler == nil {
		return true
	}
	return logger.sampler.Allow(entry)
}

func (logger *StarLogger) SetSamplingConfig(cfg SamplingConfig) {
	if logger == nil || logger.logcore == nil {
		return
	}
	logger.logcore.mu.Lock()
	if logger.logcore.sampler == nil {
		logger.logcore.sampler = newSampler()
	}
	s := logger.logcore.sampler
	logger.logcore.mu.Unlock()
	s.SetConfig(cfg)
}

func (logger *StarLogger) GetSamplingConfig() SamplingConfig {
	if logger == nil || logger.logcore == nil {
		return normalizeSamplingConfig(DefaultSamplingConfig())
	}
	logger.logcore.mu.Lock()
	s := logger.logcore.sampler
	logger.logcore.mu.Unlock()
	if s == nil {
		return normalizeSamplingConfig(DefaultSamplingConfig())
	}
	return s.Config()
}

func (logger *StarLogger) EnableSampling(enable bool) {
	cfg := logger.GetSamplingConfig()
	cfg.Enable = enable
	logger.SetSamplingConfig(cfg)
}

func (logger *StarLogger) SetSamplingDropHandler(handler func(SamplingDropData)) {
	cfg := logger.GetSamplingConfig()
	cfg.OnDrop = handler
	logger.SetSamplingConfig(cfg)
}

func (logger *StarLogger) GetSamplingStats() SamplingStats {
	if logger == nil || logger.logcore == nil {
		return SamplingStats{}
	}
	logger.logcore.mu.Lock()
	s := logger.logcore.sampler
	logger.logcore.mu.Unlock()
	if s == nil {
		return SamplingStats{}
	}
	return s.Stats()
}

func (logger *StarLogger) ResetSamplingStats() {
	if logger == nil || logger.logcore == nil {
		return
	}
	logger.logcore.mu.Lock()
	s := logger.logcore.sampler
	logger.logcore.mu.Unlock()
	if s == nil {
		return
	}
	s.ResetStats()
}

type DedupScope int

const (
	DedupScopeGlobal DedupScope = iota
	DedupScopeByKey
)

type DedupDropData struct {
	Time        time.Time
	Key         string
	Reason      string
	Level       int
	LevelName   string
	LoggerName  string
	Message     string
	Window      time.Duration
	Allowed     uint64
	Dropped     uint64
	CurrentKeys int
}

type DedupStats struct {
	Enabled      bool
	Window       time.Duration
	Scope        DedupScope
	Allowed      uint64
	Dropped      uint64
	LastDropTime time.Time
	LastDropKey  string
	LastReason   string
	CurrentKeys  int
}

type DedupConfig struct {
	Enable bool

	Levels []int
	Window time.Duration
	Scope  DedupScope

	KeyFunc func(*Entry) string
	MaxKeys int
	KeyTTL  time.Duration

	OnDrop func(DedupDropData)
}

type dedupItem struct {
	lastSeen time.Time
}

type deduper struct {
	mu sync.Mutex

	cfg          DedupConfig
	limitedLevel map[int]struct{}
	items        map[string]*dedupItem
	nowFunc      func() time.Time

	allowedCount    uint64
	droppedCount    uint64
	lastDropTime    time.Time
	lastDropKey     string
	lastDropReason  string
	lastCleanupTime time.Time
}

func DefaultDedupConfig() DedupConfig {
	return DedupConfig{
		Enable:  false,
		Levels:  nil,
		Window:  2 * time.Second,
		Scope:   DedupScopeByKey,
		KeyFunc: nil,
		MaxKeys: 4096,
		KeyTTL:  10 * time.Second,
		OnDrop:  nil,
	}
}

func cloneDedupConfig(cfg DedupConfig) DedupConfig {
	cloned := cfg
	cloned.Levels = cloneIntSlice(cfg.Levels)
	return cloned
}

func normalizeDedupConfig(cfg DedupConfig) DedupConfig {
	if cfg.Window <= 0 {
		cfg.Window = 2 * time.Second
	}
	switch cfg.Scope {
	case DedupScopeGlobal, DedupScopeByKey:
	default:
		cfg.Scope = DedupScopeByKey
	}
	if cfg.MaxKeys <= 0 {
		cfg.MaxKeys = 4096
	}
	if cfg.KeyTTL <= 0 {
		cfg.KeyTTL = cfg.Window * 4
		if cfg.KeyTTL < 10*time.Second {
			cfg.KeyTTL = 10 * time.Second
		}
	}
	return cloneDedupConfig(cfg)
}

func newDeduper() *deduper {
	cfg := normalizeDedupConfig(DefaultDedupConfig())
	return &deduper{
		cfg:          cfg,
		limitedLevel: buildLevelSet(cfg.Levels),
		items:        make(map[string]*dedupItem),
		nowFunc:      time.Now,
	}
}

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

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

func (d *deduper) SetConfig(cfg DedupConfig) {
	if d == nil {
		return
	}
	normalized := normalizeDedupConfig(cfg)
	d.mu.Lock()
	d.cfg = normalized
	d.limitedLevel = buildLevelSet(normalized.Levels)
	d.items = make(map[string]*dedupItem)
	d.lastCleanupTime = time.Time{}
	d.mu.Unlock()
}

func (d *deduper) Config() DedupConfig {
	if d == nil {
		return normalizeDedupConfig(DefaultDedupConfig())
	}
	d.mu.Lock()
	defer d.mu.Unlock()
	return cloneDedupConfig(d.cfg)
}

func (d *deduper) Stats() DedupStats {
	if d == nil {
		return DedupStats{}
	}
	d.mu.Lock()
	defer d.mu.Unlock()
	return DedupStats{
		Enabled:      d.cfg.Enable,
		Window:       d.cfg.Window,
		Scope:        d.cfg.Scope,
		Allowed:      d.allowedCount,
		Dropped:      d.droppedCount,
		LastDropTime: d.lastDropTime,
		LastDropKey:  d.lastDropKey,
		LastReason:   d.lastDropReason,
		CurrentKeys:  len(d.items),
	}
}

func (d *deduper) ResetStats() {
	if d == nil {
		return
	}
	d.mu.Lock()
	d.allowedCount = 0
	d.droppedCount = 0
	d.lastDropTime = time.Time{}
	d.lastDropKey = ""
	d.lastDropReason = ""
	d.mu.Unlock()
}

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

func (d *deduper) resolveKey(entry *Entry) string {
	if d.cfg.Scope == DedupScopeGlobal {
		return "__global__"
	}
	if d.cfg.KeyFunc != nil {
		key := strings.TrimSpace(d.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 (d *deduper) cleanupItemsLocked(now time.Time) {
	if d.cfg.Scope != DedupScopeByKey || d.cfg.KeyTTL <= 0 {
		return
	}
	if !d.lastCleanupTime.IsZero() && now.Sub(d.lastCleanupTime) < time.Second {
		return
	}
	for key, item := range d.items {
		if item == nil {
			delete(d.items, key)
			continue
		}
		if now.Sub(item.lastSeen) > d.cfg.KeyTTL {
			delete(d.items, key)
		}
	}
	d.lastCleanupTime = now
}

func (d *deduper) getItemLocked(key string, now time.Time) *dedupItem {
	if item, ok := d.items[key]; ok && item != nil {
		return item
	}
	if d.cfg.Scope == DedupScopeByKey && d.cfg.MaxKeys > 0 && len(d.items) >= d.cfg.MaxKeys {
		oldestKey := ""
		oldestTime := now
		for existingKey, item := range d.items {
			if item == nil {
				oldestKey = existingKey
				break
			}
			if oldestKey == "" || item.lastSeen.Before(oldestTime) {
				oldestKey = existingKey
				oldestTime = item.lastSeen
			}
		}
		if oldestKey != "" {
			delete(d.items, oldestKey)
		}
	}
	item := &dedupItem{}
	d.items[key] = item
	return item
}

func (d *deduper) Allow(entry *Entry) bool {
	if d == nil || entry == nil {
		return true
	}
	now := d.now()
	var callback func(DedupDropData)
	dropData := DedupDropData{}
	allow := true

	d.mu.Lock()
	if !d.cfg.Enable {
		d.allowedCount++
		d.mu.Unlock()
		return true
	}
	if !d.isLimitedLevel(entry.Level) {
		d.allowedCount++
		d.mu.Unlock()
		return true
	}

	key := d.resolveKey(entry)
	d.cleanupItemsLocked(now)
	item := d.getItemLocked(key, now)
	if !item.lastSeen.IsZero() && now.Sub(item.lastSeen) < d.cfg.Window {
		item.lastSeen = now
		d.droppedCount++
		d.lastDropTime = now
		d.lastDropKey = key
		d.lastDropReason = "dedup_window"
		allow = false
		dropData = DedupDropData{
			Time:        now,
			Key:         key,
			Reason:      d.lastDropReason,
			Level:       entry.Level,
			LevelName:   entry.LevelName,
			LoggerName:  entry.LoggerName,
			Message:     entry.Message,
			Window:      d.cfg.Window,
			Allowed:     d.allowedCount,
			Dropped:     d.droppedCount,
			CurrentKeys: len(d.items),
		}
		callback = d.cfg.OnDrop
	} else {
		item.lastSeen = now
		d.allowedCount++
	}
	d.mu.Unlock()

	if !allow && 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) allowByDedup(entry *Entry) bool {
	if logger == nil || logger.deduper == nil {
		return true
	}
	return logger.deduper.Allow(entry)
}

func (logger *StarLogger) SetDedupConfig(cfg DedupConfig) {
	if logger == nil || logger.logcore == nil {
		return
	}
	logger.logcore.mu.Lock()
	if logger.logcore.deduper == nil {
		logger.logcore.deduper = newDeduper()
	}
	d := logger.logcore.deduper
	logger.logcore.mu.Unlock()
	d.SetConfig(cfg)
}

func (logger *StarLogger) GetDedupConfig() DedupConfig {
	if logger == nil || logger.logcore == nil {
		return normalizeDedupConfig(DefaultDedupConfig())
	}
	logger.logcore.mu.Lock()
	d := logger.logcore.deduper
	logger.logcore.mu.Unlock()
	if d == nil {
		return normalizeDedupConfig(DefaultDedupConfig())
	}
	return d.Config()
}

func (logger *StarLogger) EnableDedup(enable bool) {
	cfg := logger.GetDedupConfig()
	cfg.Enable = enable
	logger.SetDedupConfig(cfg)
}

func (logger *StarLogger) SetDedupDropHandler(handler func(DedupDropData)) {
	cfg := logger.GetDedupConfig()
	cfg.OnDrop = handler
	logger.SetDedupConfig(cfg)
}

func (logger *StarLogger) GetDedupStats() DedupStats {
	if logger == nil || logger.logcore == nil {
		return DedupStats{}
	}
	logger.logcore.mu.Lock()
	d := logger.logcore.deduper
	logger.logcore.mu.Unlock()
	if d == nil {
		return DedupStats{}
	}
	return d.Stats()
}

func (logger *StarLogger) ResetDedupStats() {
	if logger == nil || logger.logcore == nil {
		return
	}
	logger.logcore.mu.Lock()
	d := logger.logcore.deduper
	logger.logcore.mu.Unlock()
	if d == nil {
		return
	}
	d.ResetStats()
}