github.com/jxskiss/gopkg/v2@v2.14.9-0.20240514120614-899f3e7952b4/zlog/redirect_std.go (about)

     1  package zlog
     2  
     3  import (
     4  	"bytes"
     5  	"log"
     6  	"os"
     7  	"strings"
     8  
     9  	"go.uber.org/zap"
    10  )
    11  
    12  var replaceSlogDefault func(l *zap.Logger, disableCaller bool) func()
    13  
    14  // redirectStdLog redirects output from the standard library's package-global
    15  // logger to the supplied logger, it detects level from the logging messages,
    16  // or use InfoLevel as default.
    17  // Since zap already handles caller annotations, timestamps, etc.,
    18  // it automatically disables the standard library's annotations and prefixing.
    19  //
    20  // It returns a function to restore the original prefix and flags and reset the
    21  // standard library's output to os.Stderr.
    22  func redirectStdLog(l *zap.Logger, disableCaller bool) func() {
    23  	resetPkgSlog := func() {}
    24  	if replaceSlogDefault != nil {
    25  		resetPkgSlog = replaceSlogDefault(l, disableCaller)
    26  	}
    27  	resetPkgLog := replaceLogDefault(l)
    28  	return func() {
    29  		resetPkgSlog()
    30  		resetPkgLog()
    31  	}
    32  }
    33  
    34  func replaceLogDefault(l *zap.Logger) func() {
    35  	oldFlag := log.Flags()
    36  	oldPrefix := log.Prefix()
    37  	log.SetFlags(0)
    38  	log.SetPrefix("")
    39  	log.SetOutput((*stdLogWriter)(
    40  		l.Named("stdlog").WithOptions(zap.AddCallerSkip(3))))
    41  	return func() {
    42  		log.SetFlags(oldFlag)
    43  		log.SetPrefix(oldPrefix)
    44  		log.SetOutput(os.Stderr)
    45  	}
    46  }
    47  
    48  type stdLogWriter zap.Logger
    49  
    50  func (l *stdLogWriter) Write(p []byte) (int, error) {
    51  	n := len(p)
    52  	p = bytes.TrimSpace(p)
    53  	str := string(p)
    54  	level, ok := detectLevel(str)
    55  	if !ok {
    56  		level = InfoLevel
    57  	}
    58  	(*zap.Logger)(l).Log(level, str)
    59  	return n, nil
    60  }
    61  
    62  // detectLevel guess logging level by checking the begging of a message.
    63  // Note that we don't guess level greater than ErrorLevel
    64  // from this function to avoid crashing a program accidentally.
    65  func detectLevel(message string) (Level, bool) {
    66  	const levelPrefixMinLen = 5
    67  	if len(message) < levelPrefixMinLen {
    68  		return 0, false
    69  	}
    70  	end := uint8(':')
    71  	if message[0] == '[' {
    72  		end, message = ']', message[1:]
    73  	}
    74  	switch message[0] {
    75  	case 'T', 't':
    76  		if len(message) > 5 && message[5] == end &&
    77  			strings.EqualFold("trace", message[:5]) {
    78  			return TraceLevel, true
    79  		}
    80  	case 'D', 'd':
    81  		if len(message) > 5 && message[5] == end &&
    82  			strings.EqualFold("debug", message[:5]) {
    83  			return DebugLevel, true
    84  		}
    85  	case 'I', 'i':
    86  		if len(message) > 4 && message[4] == end &&
    87  			strings.EqualFold("info", message[:4]) {
    88  			return InfoLevel, true
    89  		}
    90  	case 'W', 'w':
    91  		if len(message) > 4 && message[4] == end &&
    92  			strings.EqualFold("warn", message[:4]) {
    93  			return WarnLevel, true
    94  		}
    95  		if len(message) > 7 && message[7] == end &&
    96  			strings.EqualFold("warning", message[:7]) {
    97  			return WarnLevel, true
    98  		}
    99  	case 'E', 'e':
   100  		if len(message) > 5 && message[5] == end &&
   101  			strings.EqualFold("error", message[:5]) {
   102  			return ErrorLevel, true
   103  		}
   104  	case 'P', 'p':
   105  		if len(message) > 5 && message[5] == end &&
   106  			strings.EqualFold("panic", message[:5]) {
   107  			return ErrorLevel, true
   108  		}
   109  	case 'F', 'f':
   110  		if len(message) > 5 && message[5] == end &&
   111  			strings.EqualFold("fatal", message[:5]) {
   112  			return ErrorLevel, true
   113  		}
   114  	}
   115  	return 0, false
   116  }