github.com/castai/kvisor@v1.7.1-0.20240516114728-b3572a2607b5/pkg/logging/logging.go (about)

     1  package logging
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"io"
     7  	"log/slog"
     8  	"os"
     9  	"path/filepath"
    10  	"runtime"
    11  	"time"
    12  
    13  	"golang.org/x/time/rate"
    14  )
    15  
    16  type Config struct {
    17  	Ctx         context.Context
    18  	RateLimiter RateLimiterConfig
    19  	Level       slog.Level
    20  	AddSource   bool
    21  	Output      io.Writer
    22  	Export      ExportConfig
    23  }
    24  
    25  type RateLimiterConfig struct {
    26  	Limit  rate.Limit
    27  	Burst  int
    28  	Inform bool
    29  }
    30  
    31  type ExportConfig struct {
    32  	ExportFunc ExportFunc
    33  	MinLevel   slog.Level
    34  }
    35  
    36  func MustParseLevel(lvlStr string) slog.Level {
    37  	var lvl slog.Level
    38  	err := lvl.UnmarshalText([]byte(lvlStr))
    39  	if err != nil {
    40  		panic("parsing log level from level string " + lvlStr)
    41  	}
    42  	return lvl
    43  }
    44  
    45  func init() {
    46  	slog.SetDefault(slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{})))
    47  }
    48  
    49  func New(cfg *Config) *Logger {
    50  	out := cfg.Output
    51  	if out == nil {
    52  		out = os.Stdout
    53  	}
    54  	if cfg.Ctx == nil {
    55  		cfg.Ctx = context.Background()
    56  	}
    57  	var replace func(groups []string, a slog.Attr) slog.Attr
    58  	if cfg.AddSource {
    59  		replace = func(groups []string, a slog.Attr) slog.Attr {
    60  			// Remove the directory from the source's filename.
    61  			if a.Key == slog.SourceKey {
    62  				source := a.Value.Any().(*slog.Source)
    63  				source.File = filepath.Base(source.File)
    64  			}
    65  			return a
    66  		}
    67  	}
    68  
    69  	// Initial logger.
    70  	var handler slog.Handler = slog.NewTextHandler(out, &slog.HandlerOptions{
    71  		AddSource:   cfg.AddSource,
    72  		Level:       cfg.Level,
    73  		ReplaceAttr: replace,
    74  	})
    75  
    76  	// Export logs handler.
    77  	if cfg.Export.ExportFunc != nil {
    78  		handler = NewExportHandler(cfg.Ctx, handler, cfg.Export)
    79  	}
    80  
    81  	// Rate limiter handler.
    82  	if cfg.RateLimiter.Limit != 0 {
    83  		handler = NewRateLimiterHandler(cfg.Ctx, handler, cfg.RateLimiter)
    84  	}
    85  
    86  	log := slog.New(handler)
    87  	return &Logger{log: log}
    88  }
    89  
    90  func NewTestLog() *Logger {
    91  	return New(&Config{Level: slog.LevelDebug})
    92  }
    93  
    94  type Logger struct {
    95  	log *slog.Logger
    96  }
    97  
    98  func (l *Logger) Error(msg string) {
    99  	l.doLog(slog.LevelError, msg)
   100  }
   101  
   102  func (l *Logger) Errorf(format string, a ...any) {
   103  	l.doLog(slog.LevelError, format, a...)
   104  }
   105  
   106  func (l *Logger) Infof(format string, a ...any) {
   107  	l.doLog(slog.LevelInfo, format, a...)
   108  }
   109  
   110  func (l *Logger) Info(msg string) {
   111  	l.doLog(slog.LevelInfo, msg)
   112  }
   113  
   114  func (l *Logger) Debug(msg string) {
   115  	l.doLog(slog.LevelDebug, msg)
   116  }
   117  
   118  func (l *Logger) Debugf(format string, a ...any) {
   119  	l.doLog(slog.LevelDebug, format, a...)
   120  }
   121  
   122  func (l *Logger) Warn(msg string) {
   123  	l.doLog(slog.LevelWarn, msg)
   124  }
   125  
   126  func (l *Logger) Warnf(format string, a ...any) {
   127  	l.doLog(slog.LevelWarn, format, a...)
   128  }
   129  
   130  func (l *Logger) Fatal(msg string) {
   131  	l.doLog(slog.LevelError, msg)
   132  	os.Exit(1)
   133  }
   134  
   135  func (l *Logger) IsEnabled(lvl slog.Level) bool {
   136  	ctx := context.Background()
   137  	return l.log.Handler().Enabled(ctx, lvl)
   138  }
   139  
   140  func (l *Logger) doLog(lvl slog.Level, format string, args ...any) {
   141  	ctx := context.Background()
   142  	if !l.log.Handler().Enabled(ctx, lvl) {
   143  		return
   144  	}
   145  	var pcs [1]uintptr
   146  	runtime.Callers(3, pcs[:])
   147  	r := slog.NewRecord(time.Now(), lvl, fmt.Sprintf(format, args...), pcs[0])
   148  	_ = l.log.Handler().Handle(ctx, r) //nolint:contextcheck
   149  }
   150  
   151  func (l *Logger) With(args ...any) *Logger {
   152  	return &Logger{log: l.log.With(args...)}
   153  }
   154  
   155  func (l *Logger) WithField(k, v string) *Logger {
   156  	return &Logger{log: l.log.With(slog.String(k, v))}
   157  }