sorbet/pkg/log/writer.go

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
}