starlog/typed.go

package starlog

import (
	"context"
	"errors"
	"fmt"
	"io"
	"math/rand"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"b612.me/starlog/colorable"
)

const (
	LvDebug = iota
	LvInfo
	LvNotice
	LvWarning
	LvError
	LvCritical
	LvPanic
	LvFatal
)

type ColorMode int

const (
	ColorModeOff ColorMode = iota
	ColorModeFullLine
	ColorModeLevelOnly
)

type PendingDropPolicy int

const (
	PendingDropOldest PendingDropPolicy = iota
	PendingDropNewest
	PendingBlock
)

type RedactFailMode int

const (
	RedactFailMaskAll RedactFailMode = iota
	RedactFailOpen
	RedactFailDrop
)

const (
	FieldTypeString = "string"
	FieldTypeNumber = "number"
	FieldTypeBool   = "bool"
	FieldTypeError  = "error"
	FieldTypeNil    = "nil"
	FieldTypeOther  = "other"
)

const defaultAsyncQueueCapacity uint64 = 1024

var (
	ErrAsyncHandlerPanic   = errors.New("async handler panic")
	ErrAsyncHandlerTimeout = errors.New("async handler timeout")
	ErrAsyncQueueFull      = errors.New("async queue full")
	ErrPendingWriteDropped = errors.New("pending write dropped")
	ErrInvalidLevel        = errors.New("invalid log level")
	ErrRedactionFailed     = errors.New("redaction failed")

	levels = map[int]string{
		LvDebug:    "DEBUG",
		LvInfo:     "INFO",
		LvNotice:   "NOTICE",
		LvWarning:  "WARNING",
		LvError:    "ERROR",
		LvCritical: "CRITICAL",
		LvPanic:    "PANIC",
		LvFatal:    "FATAL",
	}
	levelAliases = map[string]int{
		"debug":    LvDebug,
		"info":     LvInfo,
		"notice":   LvNotice,
		"warn":     LvWarning,
		"warning":  LvWarning,
		"err":      LvError,
		"error":    LvError,
		"critical": LvCritical,
		"crit":     LvCritical,
		"panic":    LvPanic,
		"fatal":    LvFatal,
	}

	stdScreen io.Writer = colorable.NewColorableStdout()
	errScreen io.Writer = colorable.NewColorableStderr()

	defaultAsyncRuntimeOnce     sync.Once
	defaultAsyncRuntimeFallback *asyncRuntime
)

type asyncRuntime struct {
	mu sync.Mutex

	queue    *starChanStack
	started  bool
	stopChan chan struct{}
	doneChan chan struct{}

	dropCount uint64

	asyncAlert   func(error, LogData)
	asyncAlertMu sync.RWMutex
	fallbackSync uint32
	timeout      int64

	writeErrCount   uint64
	writeErrHandler func(error, LogData)
	writeErrMu      sync.RWMutex

	queueCapacity uint64
}

type asyncRuntimeContextKey struct{}

type starlog struct {
	mu                  *sync.Mutex
	runtime             *asyncRuntime
	output              io.Writer
	minLevel            int
	errOutputLevel      int
	showFuncName        bool
	showThread          bool
	showLevel           bool
	showDeatilFile      bool
	showColor           bool
	switching           bool
	showStd             bool
	onlyColorLevel      bool
	autoAppendNewline   bool
	stopWriter          bool
	id                  string
	name                string
	colorList           map[int][]Attr
	colorMe             map[int]*Color
	keywordColors       map[string][]Attr
	keywordOrder        []string
	keywordColorizers   map[string]*Color
	keywordMatcher      *keywordMatcher
	keywordMatchOptions KeywordMatchOptions
	showFieldColor      bool
	fieldKeyColor       []Attr
	fieldTypeColors     map[string][]Attr
	fieldValueColors    map[string][]Attr
	entryHandler        Handler
	redactor            Redactor
	redactRules         []RedactRule
	redactFailMode      RedactFailMode
	redactMaskToken     string
	redactErrorCount    uint64
	formatter           Formatter
	sink                Sink
	pendingCond         *sync.Cond
	pendingWrites       []string
	pendingWriteLimit   int
	pendingDropPolicy   PendingDropPolicy
	pendingDropCount    uint64
	pendingBlockCount   uint64
	pendingPeakLen      uint64
	rateLimiter         *rateLimiter
	sampler             *sampler
	deduper             *deduper
	contextFields       func(context.Context) Fields
	entryHandlerTimeout time.Duration
}

type StarLogger struct {
	thread      string
	handlerFunc func(LogData)
	logcore     *starlog
	isStd       bool
	fields      Fields
	logErr      error
	logCtx      context.Context
}

type logTransfer struct {
	handlerFunc func(LogData)
	LogData
}

type LogData struct {
	Name   string
	Log    string
	Colors []Attr
}

type PendingStats struct {
	Limit      int
	Length     int
	PeakLength int
	DropCount  uint64
	BlockCount uint64
	Policy     PendingDropPolicy
	Switching  bool
}

type KeywordMatchOptions struct {
	IgnoreCase bool
	WholeWord  bool
}

// Config is a logger core snapshot that can be read and applied atomically.
// Prefer GetConfig + UpdateConfig/ApplyConfig for multi-field configuration.
type Config struct {
	Name              string
	Level             int
	StdErrLevel       int
	ShowFuncName      bool
	ShowFlag          bool
	ShowLevel         bool
	ShowOriginFile    bool
	ShowColor         bool
	OnlyColorLevel    bool
	ShowStd           bool
	StopWriter        bool
	AutoAppendNewline bool
	Switching         bool
	LevelColors       map[int][]Attr
	KeywordColors     map[string][]Attr
	KeywordMatch      KeywordMatchOptions
	ShowFieldColor    bool
	FieldKeyColor     []Attr
	FieldTypeColors   map[string][]Attr
	FieldValueColors  map[string][]Attr

	EntryHandler        Handler
	EntryHandlerTimeout time.Duration
	Formatter           Formatter
	Sink                Sink
	Writer              io.Writer

	PendingWriteLimit int
	PendingDropPolicy PendingDropPolicy

	Redactor        Redactor
	RedactRules     []RedactRule
	RedactFailMode  RedactFailMode
	RedactMaskToken string

	RateLimit RateLimitConfig
	Sampling  SamplingConfig
	Dedup     DedupConfig

	ContextFieldExtractor func(context.Context) Fields
}

func newAsyncRuntime(queueCapacity uint64) *asyncRuntime {
	if queueCapacity == 0 {
		queueCapacity = defaultAsyncQueueCapacity
	}
	runtime := &asyncRuntime{
		queueCapacity: queueCapacity,
	}
	atomic.StoreUint32(&runtime.fallbackSync, 1)
	return runtime
}

func defaultAsyncRuntime() *asyncRuntime {
	if Std != nil && Std.logcore != nil && Std.logcore.runtime != nil {
		return Std.logcore.runtime
	}
	defaultAsyncRuntimeOnce.Do(func() {
		defaultAsyncRuntimeFallback = newAsyncRuntime(defaultAsyncQueueCapacity)
	})
	return defaultAsyncRuntimeFallback
}

func withAsyncRuntime(ctx context.Context, runtime *asyncRuntime) context.Context {
	if runtime == nil {
		return ctx
	}
	if ctx == nil {
		ctx = context.Background()
	}
	return context.WithValue(ctx, asyncRuntimeContextKey{}, runtime)
}

func runtimeFromContext(ctx context.Context) *asyncRuntime {
	if ctx == nil {
		return nil
	}
	runtime, _ := ctx.Value(asyncRuntimeContextKey{}).(*asyncRuntime)
	return runtime
}

func reportWriteErrorWithContext(ctx context.Context, err error, data LogData) {
	runtime := runtimeFromContext(ctx)
	if runtime == nil {
		runtime = defaultAsyncRuntime()
	}
	runtime.reportWriteError(err, data)
}

func newLogCore(out io.Writer) *starlog {
	core := &starlog{
		mu:             &sync.Mutex{},
		runtime:        newAsyncRuntime(defaultAsyncQueueCapacity),
		output:         out,
		minLevel:       LvDebug,
		errOutputLevel: LvError,
		showFuncName:   true,
		showThread:     true,
		showLevel:      true,
		showStd:        true,
		showDeatilFile: true,
		switching:      false,
		stopWriter:     false,
		showColor:      true,
		id:             generateId(),
		colorList: map[int][]Attr{
			LvDebug:    []Attr{FgWhite},
			LvInfo:     []Attr{FgGreen},
			LvNotice:   []Attr{FgCyan},
			LvWarning:  []Attr{FgYellow},
			LvError:    []Attr{FgMagenta},
			LvCritical: []Attr{FgRed, Bold},
			LvPanic:    []Attr{FgRed, Bold},
			LvFatal:    []Attr{FgRed},
		},
		colorMe: map[int]*Color{
			LvDebug:    NewColor([]Attr{FgWhite}...),
			LvInfo:     NewColor([]Attr{FgGreen}...),
			LvNotice:   NewColor([]Attr{FgCyan}...),
			LvWarning:  NewColor([]Attr{FgYellow}...),
			LvError:    NewColor([]Attr{FgMagenta}...),
			LvCritical: NewColor([]Attr{FgRed, Bold}...),
			LvPanic:    NewColor([]Attr{FgRed, Bold}...),
			LvFatal:    NewColor([]Attr{FgRed}...),
		},
		keywordColors:     make(map[string][]Attr),
		keywordOrder:      nil,
		keywordColorizers: nil,
		keywordMatcher:    nil,
		showFieldColor:    true,
		fieldKeyColor:     []Attr{FgHiBlue},
		fieldTypeColors: map[string][]Attr{
			FieldTypeString: []Attr{FgGreen},
			FieldTypeNumber: []Attr{FgYellow},
			FieldTypeBool:   []Attr{FgMagenta},
			FieldTypeError:  []Attr{FgRed},
			FieldTypeNil:    []Attr{FgHiBlack},
			FieldTypeOther:  []Attr{FgCyan},
		},
		fieldValueColors:    make(map[string][]Attr),
		redactRules:         make([]RedactRule, 0, 4),
		redactFailMode:      RedactFailMaskAll,
		redactMaskToken:     "[REDACTED]",
		pendingWrites:       make([]string, 0, 16),
		pendingWriteLimit:   1024,
		pendingDropPolicy:   PendingDropOldest,
		rateLimiter:         newRateLimiter(),
		sampler:             newSampler(),
		deduper:             newDeduper(),
		entryHandlerTimeout: 0,
	}
	core.rebuildKeywordCachesLocked()
	core.pendingCond = sync.NewCond(core.mu)
	return core
}

func ParseLevel(level string) (int, error) {
	val := strings.TrimSpace(strings.ToLower(level))
	if val == "" {
		return 0, fmt.Errorf("%w: empty", ErrInvalidLevel)
	}
	if parsed, ok := levelAliases[val]; ok {
		return parsed, nil
	}
	num, err := strconv.Atoi(val)
	if err == nil {
		return num, nil
	}
	return 0, fmt.Errorf("%w: %s", ErrInvalidLevel, level)
}

func NewStarlog(out io.Writer) *StarLogger {
	return &StarLogger{
		handlerFunc: nil,
		thread:      "MAN",
		logcore:     newLogCore(out),
		isStd:       false,
		fields:      nil,
		logErr:      nil,
		logCtx:      nil,
	}
}

func (logger *StarLogger) StdErrLevel() int {
	logger.logcore.mu.Lock()
	defer logger.logcore.mu.Unlock()
	return logger.logcore.errOutputLevel
}

func (logger *StarLogger) SetStdErrLevel(level int) {
	logger.logcore.mu.Lock()
	defer logger.logcore.mu.Unlock()
	logger.logcore.errOutputLevel = level
}

func (logger *StarLogger) NewFlag() *StarLogger {
	return &StarLogger{
		thread:      getRandomFlag(false),
		handlerFunc: logger.handlerFunc,
		logcore:     logger.logcore,
		isStd:       false,
		fields:      cloneFields(logger.fields),
		logErr:      logger.logErr,
		logCtx:      logger.logCtx,
	}
}
func (logger *StarLogger) SetNewRandomFlag() {
	logger.thread = getRandomFlag(false)
}

func (logger *StarLogger) SetName(name string) {
	logger.logcore.mu.Lock()
	defer logger.logcore.mu.Unlock()
	logger.logcore.name = name
}

func (logger *StarLogger) GetName() string {
	logger.logcore.mu.Lock()
	defer logger.logcore.mu.Unlock()
	return logger.logcore.name
}

func getRandomFlag(isMain bool) string {
	rand.Seed(time.Now().UnixNano())
	if isMain {
		return "MAN"
	}
	flag := "MAN"
	for flag == "MAN" {
		flag = string([]byte{uint8(rand.Intn(26) + 65), uint8(rand.Intn(26) + 65), uint8(rand.Intn(26) + 65)})
	}
	return flag
}

func generateId() string {
	rand.Seed(time.Now().UnixNano())
	return fmt.Sprintf("%dstar%db612%d", time.Now().UnixNano(), rand.Intn(1000000), rand.Intn(1000000))
}

func (logger *starlog) asyncRuntime() *asyncRuntime {
	if logger == nil || logger.runtime == nil {
		return defaultAsyncRuntime()
	}
	return logger.runtime
}

func (logger *StarLogger) asyncRuntime() *asyncRuntime {
	if logger == nil || logger.logcore == nil {
		return defaultAsyncRuntime()
	}
	return logger.logcore.asyncRuntime()
}

func (logger *starlog) reportAsyncDrop(err error, data LogData) {
	logger.asyncRuntime().reportAsyncDrop(err, data)
}

func (logger *starlog) reportWriteError(err error, data LogData) {
	logger.asyncRuntime().reportWriteError(err, data)
}

func (runtime *asyncRuntime) snapshot() (*starChanStack, bool) {
	if runtime == nil {
		return nil, false
	}
	runtime.mu.Lock()
	defer runtime.mu.Unlock()
	return runtime.queue, runtime.started
}

func (runtime *asyncRuntime) Start() {
	if runtime == nil {
		return
	}
	runtime.mu.Lock()
	if runtime.started {
		runtime.mu.Unlock()
		return
	}
	queue := newStarChanStack(runtime.queueCapacity)
	stopChan := make(chan struct{})
	doneChan := make(chan struct{})
	runtime.queue = queue
	runtime.stopChan = stopChan
	runtime.doneChan = doneChan
	runtime.started = true
	runtime.mu.Unlock()

	go func(queue *starChanStack, stop <-chan struct{}, done chan struct{}) {
		defer close(done)
		for {
			select {
			case <-stop:
				return
			default:
			}
			popped, err := queue.Pop()
			if err != nil {
				if errors.Is(err, io.EOF) {
					return
				}
				return
			}
			val, ok := popped.(logTransfer)
			if !ok {
				continue
			}
			if val.handlerFunc != nil {
				runtime.invokeAsyncHandlerSafely(val.handlerFunc, val.LogData)
			}
		}
	}(queue, stopChan, doneChan)
}

func (runtime *asyncRuntime) Stop() {
	if runtime == nil {
		return
	}
	runtime.mu.Lock()
	if !runtime.started {
		runtime.mu.Unlock()
		return
	}
	stopChan := runtime.stopChan
	doneChan := runtime.doneChan
	queue := runtime.queue
	runtime.queue = nil
	runtime.stopChan = nil
	runtime.doneChan = nil
	runtime.started = false
	runtime.mu.Unlock()

	if stopChan != nil {
		func() {
			defer func() {
				recover()
			}()
			close(stopChan)
		}()
	}
	if queue != nil {
		_ = queue.Close()
	}
	if doneChan != nil {
		<-doneChan
	}
}

func (runtime *asyncRuntime) WaitDrain(ctx context.Context) error {
	if ctx == nil {
		ctx = context.Background()
	}
	for {
		queue, started := runtime.snapshot()
		if !started || queue == nil || queue.Len() == 0 {
			return nil
		}
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-time.After(5 * time.Millisecond):
		}
	}
}

func (runtime *asyncRuntime) Metrics() AsyncMetrics {
	queue, started := runtime.snapshot()
	snapshot := AsyncMetrics{
		Started:        started,
		Dropped:        runtime.GetAsyncDropCount(),
		FallbackToSync: runtime.GetAsyncFallbackToSync(),
		HandlerTimeout: runtime.GetAsyncHandlerTimeout(),
	}
	if queue != nil {
		snapshot.QueueLength = queue.Len()
		snapshot.QueueCapacity = queue.Cap()
		snapshot.QueueFree = queue.Free()
	}
	return snapshot
}

func (runtime *asyncRuntime) SetAsyncErrorHandler(alert func(error, LogData)) {
	if runtime == nil {
		return
	}
	runtime.asyncAlertMu.Lock()
	defer runtime.asyncAlertMu.Unlock()
	runtime.asyncAlert = alert
}

func (runtime *asyncRuntime) SetAsyncFallbackToSync(enable bool) {
	if runtime == nil {
		return
	}
	if enable {
		atomic.StoreUint32(&runtime.fallbackSync, 1)
		return
	}
	atomic.StoreUint32(&runtime.fallbackSync, 0)
}

func (runtime *asyncRuntime) GetAsyncFallbackToSync() bool {
	if runtime == nil {
		return true
	}
	return atomic.LoadUint32(&runtime.fallbackSync) == 1
}

func (runtime *asyncRuntime) SetAsyncHandlerTimeout(timeout time.Duration) {
	if runtime == nil {
		return
	}
	if timeout < 0 {
		timeout = 0
	}
	atomic.StoreInt64(&runtime.timeout, int64(timeout))
}

func (runtime *asyncRuntime) GetAsyncHandlerTimeout() time.Duration {
	if runtime == nil {
		return 0
	}
	return time.Duration(atomic.LoadInt64(&runtime.timeout))
}

func (runtime *asyncRuntime) GetAsyncDropCount() uint64 {
	if runtime == nil {
		return 0
	}
	return atomic.LoadUint64(&runtime.dropCount)
}

func (runtime *asyncRuntime) reportAsyncDrop(err error, data LogData) {
	if runtime == nil {
		return
	}
	atomic.AddUint64(&runtime.dropCount, 1)
	runtime.asyncAlertMu.RLock()
	alert := runtime.asyncAlert
	runtime.asyncAlertMu.RUnlock()
	if alert != nil {
		func() {
			defer func() {
				recover()
			}()
			alert(err, data)
		}()
	}
}

func (runtime *asyncRuntime) invokeAsyncHandlerSafely(handler func(LogData), data LogData) bool {
	if runtime == nil {
		return invokeAsyncHandlerDirect(defaultAsyncRuntime(), handler, data)
	}
	if handler == nil {
		return true
	}
	timeout := runtime.GetAsyncHandlerTimeout()
	if timeout <= 0 {
		return invokeAsyncHandlerDirect(runtime, handler, data)
	}
	done := make(chan bool, 1)
	go func() {
		done <- invokeAsyncHandlerDirect(runtime, handler, data)
	}()
	select {
	case ok := <-done:
		return ok
	case <-time.After(timeout):
		runtime.reportAsyncDrop(ErrAsyncHandlerTimeout, data)
		return false
	}
}

func invokeAsyncHandlerDirect(runtime *asyncRuntime, handler func(LogData), data LogData) (ok bool) {
	defer func() {
		if panicErr := recover(); panicErr != nil {
			if runtime == nil {
				runtime = defaultAsyncRuntime()
			}
			runtime.reportAsyncDrop(fmt.Errorf("%w: %v", ErrAsyncHandlerPanic, panicErr), data)
			ok = false
		}
	}()
	handler(data)
	return true
}

func (runtime *asyncRuntime) SetWriteErrorHandler(alert func(error, LogData)) {
	if runtime == nil {
		return
	}
	runtime.writeErrMu.Lock()
	defer runtime.writeErrMu.Unlock()
	runtime.writeErrHandler = alert
}

func (runtime *asyncRuntime) GetWriteErrorCount() uint64 {
	if runtime == nil {
		return 0
	}
	return atomic.LoadUint64(&runtime.writeErrCount)
}

func (runtime *asyncRuntime) reportWriteError(err error, data LogData) {
	if runtime == nil || err == nil {
		return
	}
	atomic.AddUint64(&runtime.writeErrCount, 1)
	runtime.writeErrMu.RLock()
	alert := runtime.writeErrHandler
	runtime.writeErrMu.RUnlock()
	if alert != nil {
		func() {
			defer func() {
				recover()
			}()
			alert(err, data)
		}()
	}
}

func (runtime *asyncRuntime) resetForTest() {
	if runtime == nil {
		return
	}
	runtime.Stop()
	atomic.StoreUint64(&runtime.dropCount, 0)
	runtime.SetAsyncErrorHandler(nil)
	runtime.SetAsyncFallbackToSync(true)
	runtime.SetAsyncHandlerTimeout(0)
	atomic.StoreUint64(&runtime.writeErrCount, 0)
	runtime.SetWriteErrorHandler(nil)
}

func StartStacks() {
	defaultAsyncRuntime().Start()
}

func StopStacks() {
	defaultAsyncRuntime().Stop()
}

func Stop() {
	StopStacks()
}

func SetAsyncErrorHandler(alert func(error, LogData)) {
	defaultAsyncRuntime().SetAsyncErrorHandler(alert)
}

func SetAsyncFallbackToSync(enable bool) {
	defaultAsyncRuntime().SetAsyncFallbackToSync(enable)
}

func GetAsyncFallbackToSync() bool {
	return defaultAsyncRuntime().GetAsyncFallbackToSync()
}

func SetAsyncHandlerTimeout(timeout time.Duration) {
	defaultAsyncRuntime().SetAsyncHandlerTimeout(timeout)
}

func GetAsyncHandlerTimeout() time.Duration {
	return defaultAsyncRuntime().GetAsyncHandlerTimeout()
}

func GetAsyncDropCount() uint64 {
	return defaultAsyncRuntime().GetAsyncDropCount()
}

func reportAsyncDrop(err error, data LogData) {
	defaultAsyncRuntime().reportAsyncDrop(err, data)
}

func invokeAsyncHandlerSafely(handler func(LogData), data LogData) bool {
	return defaultAsyncRuntime().invokeAsyncHandlerSafely(handler, data)
}

func SetWriteErrorHandler(alert func(error, LogData)) {
	defaultAsyncRuntime().SetWriteErrorHandler(alert)
}

func GetWriteErrorCount() uint64 {
	return defaultAsyncRuntime().GetWriteErrorCount()
}

func reportWriteError(err error, data LogData) {
	defaultAsyncRuntime().reportWriteError(err, data)
}

func resetAsyncMetricsForTest() {
	defaultAsyncRuntime().resetForTest()
}