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

     1  package zlog
     2  
     3  import (
     4  	"go.uber.org/multierr"
     5  	"go.uber.org/zap/zapcore"
     6  )
     7  
     8  func newMultiFilesCore(cfg *Config, enc zapcore.Encoder, enab zapcore.LevelEnabler) (
     9  	core *multiFilesCore, closers []func(), err error) {
    10  	core = &multiFilesCore{
    11  		LevelEnabler: enab,
    12  		enc:          enc,
    13  	}
    14  	closers, err = core.initFileWriters(cfg)
    15  	if err != nil {
    16  		return nil, nil, err
    17  	}
    18  	return
    19  }
    20  
    21  type multiFilesCore struct {
    22  	zapcore.LevelEnabler
    23  	enc        zapcore.Encoder
    24  	defaultOut zapcore.WriteSyncer
    25  	outFunc    func(string) (zapcore.WriteSyncer, bool)
    26  	outList    []zapcore.WriteSyncer
    27  }
    28  
    29  func (c *multiFilesCore) initFileWriters(cfg *Config) (closers []func(), err error) {
    30  	// Close the opened files in case of error occurs.
    31  	defer func() {
    32  		if err != nil {
    33  			runClosers(closers)
    34  			closers = nil
    35  		}
    36  	}()
    37  
    38  	var closer func()
    39  	c.defaultOut, closer, err = cfg.FileWriterFactory(&cfg.File)
    40  	if err != nil {
    41  		return closers, err
    42  	}
    43  	closers = append(closers, closer)
    44  	if len(cfg.PerLoggerFiles) == 0 {
    45  		return closers, nil
    46  	}
    47  
    48  	tree := &radixTree[zapcore.WriteSyncer]{}
    49  	seenOut := make(map[string]zapcore.WriteSyncer)
    50  	var outList []zapcore.WriteSyncer
    51  	for loggerName, fc := range cfg.PerLoggerFiles {
    52  		if fc.Filename == "" {
    53  			continue
    54  		}
    55  		if fc.Filename == cfg.File.Filename {
    56  			continue
    57  		}
    58  		var out zapcore.WriteSyncer
    59  		if out = seenOut[fc.Filename]; out == nil {
    60  			fc := mergeFileConfig(fc, cfg.File)
    61  			out, closer, err = cfg.FileWriterFactory(&fc)
    62  			if err != nil {
    63  				return closers, err
    64  			}
    65  			seenOut[fc.Filename] = out
    66  			outList = append(outList, out)
    67  			closers = append(closers, closer)
    68  		}
    69  		tree.root.insert(loggerName, out)
    70  	}
    71  	c.outFunc = tree.search
    72  	c.outList = outList
    73  	return closers, nil
    74  }
    75  
    76  func (c *multiFilesCore) clone() *multiFilesCore {
    77  	return &multiFilesCore{
    78  		LevelEnabler: c.LevelEnabler,
    79  		enc:          c.enc.Clone(),
    80  		defaultOut:   c.defaultOut,
    81  		outFunc:      c.outFunc,
    82  		outList:      c.outList,
    83  	}
    84  }
    85  
    86  func (c *multiFilesCore) With(fields []zapcore.Field) zapcore.Core {
    87  	clone := c.clone()
    88  	addFields(clone.enc, fields)
    89  	return clone
    90  }
    91  
    92  func (c *multiFilesCore) Check(ent zapcore.Entry, ce *zapcore.CheckedEntry) *zapcore.CheckedEntry {
    93  	if c.Enabled(ent.Level) {
    94  		return ce.AddCore(ent, c)
    95  	}
    96  	return ce
    97  }
    98  
    99  func (c *multiFilesCore) Write(ent zapcore.Entry, fields []zapcore.Field) error {
   100  	buf, err := c.enc.EncodeEntry(ent, fields)
   101  	if err != nil {
   102  		return err
   103  	}
   104  	out := c.defaultOut
   105  	if c.outFunc != nil {
   106  		tmp, found := c.outFunc(ent.LoggerName)
   107  		if found {
   108  			out = tmp
   109  		}
   110  	}
   111  	_, err = out.Write(buf.Bytes())
   112  	buf.Free()
   113  	if err != nil {
   114  		return err
   115  	}
   116  	if ent.Level > zapcore.ErrorLevel {
   117  		// Since we may be crashing the program, sync the output.
   118  		// Ignore Sync errors, pending a clean solution to issue
   119  		// https://github.com/uber-go/zap/issues/370.
   120  		_ = c.Sync()
   121  		_ = Sync()
   122  	}
   123  	return nil
   124  }
   125  
   126  func (c *multiFilesCore) Sync() error {
   127  	retErr := c.defaultOut.Sync()
   128  	for _, out := range c.outList {
   129  		err := out.Sync()
   130  		if err != nil {
   131  			retErr = multierr.Append(retErr, err)
   132  		}
   133  	}
   134  	return retErr
   135  }