github.com/Asutorufa/yuhaiin@v0.3.6-0.20240502055049-7984da7023a0/pkg/log/writer.go (about)

     1  package log
     2  
     3  import (
     4  	"fmt"
     5  	"io"
     6  	"log/slog"
     7  	"os"
     8  	"path/filepath"
     9  	"sort"
    10  	"strings"
    11  	"sync"
    12  	"sync/atomic"
    13  	"time"
    14  )
    15  
    16  var maxSize = 1024 * 1024
    17  var maxFile = 5
    18  
    19  func init() {
    20  	if os.Getenv("YUHAIIN_LITE") == "true" {
    21  		maxSize = 1024 * 256
    22  		maxFile = 0
    23  	}
    24  }
    25  
    26  var _ io.Writer = new(FileWriter)
    27  
    28  type FileWriter struct {
    29  	path logPath
    30  	w    *os.File
    31  	log  *slog.Logger
    32  
    33  	mu sync.RWMutex
    34  
    35  	savedSize atomic.Uint64
    36  }
    37  
    38  func NewLogWriter(file string) *FileWriter {
    39  	return &FileWriter{
    40  		path: NewPath(file),
    41  		log: slog.New(slog.NewTextHandler(os.Stderr,
    42  			&slog.HandlerOptions{AddSource: true, Level: slog.LevelDebug})),
    43  	}
    44  }
    45  
    46  func (f *FileWriter) Close() error {
    47  	if f.w != nil {
    48  		return f.w.Close()
    49  	}
    50  
    51  	return nil
    52  }
    53  
    54  func (f *FileWriter) Write(p []byte) (n int, err error) {
    55  	if f.w == nil {
    56  		f.mu.Lock()
    57  		f.w, err = os.OpenFile(f.path.FullPath(""), os.O_APPEND|os.O_CREATE|os.O_RDWR, 0644)
    58  		f.mu.Unlock()
    59  
    60  		if err != nil {
    61  			f.log.Error("open file failed:", "err", err)
    62  			return 0, err
    63  		}
    64  
    65  		f.mu.RLock()
    66  		stat, err := f.w.Stat()
    67  		if err == nil {
    68  			f.savedSize.Store(uint64(stat.Size()))
    69  		}
    70  		f.mu.RUnlock()
    71  	}
    72  
    73  	f.mu.RLock()
    74  	n, err = f.w.Write(p)
    75  	f.mu.RUnlock()
    76  
    77  	if int(f.savedSize.Add(uint64(n))) >= maxSize {
    78  		f.mu.Lock()
    79  		defer f.mu.Unlock()
    80  
    81  		f.savedSize.Store(0)
    82  
    83  		f.w.Close()
    84  		f.w = nil
    85  
    86  		err = os.Rename(f.path.FullPath(""),
    87  			f.path.FullPath(strings.ReplaceAll(time.Now().Format(time.RFC3339), ":", ".")))
    88  		if err != nil {
    89  			f.log.Error("rename file failed:", "err", err)
    90  		}
    91  
    92  		f.removeOldFile()
    93  	}
    94  
    95  	return n, err
    96  }
    97  
    98  func (f *FileWriter) removeOldFile() {
    99  	files, err := os.ReadDir(f.path.dir)
   100  	if err != nil {
   101  		f.log.Error("read dir failed:", "err", err)
   102  		return
   103  	}
   104  
   105  	if len(files) <= maxFile {
   106  		return
   107  	}
   108  
   109  	sort.Slice(files, func(i, j int) bool { return files[i].Name() > files[j].Name() })
   110  
   111  	count := 0
   112  	for _, file := range files {
   113  		if !(strings.HasPrefix(file.Name(), f.path.base+"_") && strings.HasSuffix(file.Name(), f.path.ext)) {
   114  			continue
   115  		}
   116  
   117  		count++
   118  
   119  		if count <= maxFile {
   120  			continue
   121  		}
   122  
   123  		name := filepath.Join(f.path.dir, file.Name())
   124  
   125  		err = os.Remove(name)
   126  		if err != nil {
   127  			f.log.Error("remove log file failed:", "file", name, "err", err)
   128  		} else {
   129  			f.log.Debug("remove log file", "file", name)
   130  		}
   131  	}
   132  
   133  }
   134  
   135  type logPath struct {
   136  	ext  string
   137  	base string
   138  	dir  string
   139  }
   140  
   141  func NewPath(path string) logPath {
   142  	dir := filepath.Dir(path)
   143  	ext := filepath.Ext(path)
   144  	base := strings.TrimSuffix(filepath.Base(path), ext)
   145  
   146  	return logPath{ext, base, dir}
   147  }
   148  
   149  func (l logPath) FullPath(suffix string) string {
   150  	if suffix != "" {
   151  		suffix = "_" + suffix
   152  	}
   153  
   154  	return filepath.Join(l.dir, fmt.Sprint(l.base, suffix, l.ext))
   155  }
   156  
   157  func (l *logPath) FullName(suffix string) string { return fmt.Sprint(l.base, suffix, l.ext) }