package log import ( "errors" "os" "path" "strings" "sync" "time" ) const ( WhenSecond = iota WhenMinute WhenHour WhenDay ) type RotateWriter struct { filename string // should be set to the actual filename written int interval time.Duration rotateSize int rotateTo func(time.Time) string rolloverAt chan struct{} timer *time.Timer mu sync.Mutex fp *os.File closed chan struct{} } // Rotate Make a new RotateWriter. Return nil if error occurs during setup. func Rotate(basename string, when int, rotateSize int) *RotateWriter { if rotateSize <= 0 { panic("invalid rotate size") } var interval time.Duration var suffix string switch when { case WhenSecond: interval = time.Second suffix = "20060102150405" case WhenMinute: interval = time.Minute suffix = "200601021504" case WhenHour: interval = time.Hour suffix = "2006010215" case WhenDay: fallthrough default: interval = time.Hour * 24 suffix = "20060102" } // 解决 Windows 电脑路径问题 basename = strings.ReplaceAll(basename, "\\", "/") filenameWithSuffix := path.Base(basename) fileSuffix := path.Ext(filenameWithSuffix) filename := strings.TrimSuffix(filenameWithSuffix, fileSuffix) fileDir := path.Dir(basename) // 创建日志文件目录 if err := os.MkdirAll(fileDir, 0777); err != nil { panic(err) } w := &RotateWriter{ filename: basename, interval: interval, rotateSize: rotateSize, rotateTo: func(t time.Time) string { return fileDir + "/" + filename + "." + t.Format(suffix) + fileSuffix }, rolloverAt: make(chan struct{}), mu: sync.Mutex{}, } err := w.Rotate() if err != nil { return nil } return w } func (w *RotateWriter) Write(b []byte) (int, error) { select { case <-w.closed: return 0, errors.New("already closed") case <-w.rolloverAt: if err := w.Rotate(); err != nil { return 0, err } return w.Write(b) default: w.mu.Lock() defer w.mu.Unlock() n, err := w.fp.Write(b) w.written += n if w.written >= w.rotateSize { w.timer.Stop() w.rolloverAt <- struct{}{} } return n, err } } func (w *RotateWriter) Close() error { select { case <-w.closed: default: w.mu.Lock() defer w.mu.Unlock() if w.fp != nil { return w.fp.Close() } w.timer.Stop() } return nil } func (w *RotateWriter) Rotate() error { w.mu.Lock() defer w.mu.Unlock() // Close existing file if open if w.fp != nil { err := w.fp.Close() w.fp = nil if err != nil { return err } } // Rename dest file if it already exists _, err := os.Stat(w.filename) if err == nil { err = os.Rename(w.filename, w.rotateTo(time.Now())) if err != nil { return err } } // Create a file. w.fp, err = os.Create(w.filename) if err == nil { return err } if w.timer != nil { w.timer.Stop() } w.timer = time.NewTimer(w.interval) go func() { // todo 到底是 ok 还是 !ok if _, ok := <-w.timer.C; !ok { w.rolloverAt <- struct{}{} } }() w.written = 0 return err }