golib/log/writer.go

package log

import (
	"bufio"
	"io"
	"os"
	"path/filepath"
	"sync"
	"time"
)

type Writer struct {
	pre  string
	suf  string
	path string

	date string

	cur *os.File

	noBuff bool
	buf    *bufio.Writer // buffer

	mu sync.Mutex
}

// NewRawWriter 使用 path 作为目录, pre 作为文件前缀以及 suf 文件后缀
// Writer 内使用了 bufio.Writer, 因此发生异常时需要调用 Close, 否则会最大都丢失 BuffSize 字节数据
func NewRawWriter(pre, suf, path string) (io.WriteCloser, error) {
	if err := handlePath(path); err != nil {
		return nil, err
	}
	w := new(Writer)
	w.pre = pre
	w.suf = suf
	w.path = filepath.Join(path)

	w.date = getDate()
	w.cur = (*os.File)(nil)
	w.buf = (*bufio.Writer)(nil)
	return w, nil
}

func NewWriter(pre, suf, path string) (io.WriteCloser, error) {
	return _socketCache.Get(pre, suf, path)
}

func (w *Writer) Write(p []byte) (n int, err error) {
	if date := getDate(); date != w.date {
		if err = w.Close(); err != nil {
			return 0, err
		}
		if err = w.open(); err != nil {
			return 0, err
		}
		w.date = date
	}

	if w.cur == nil {
		if err = w.open(); err != nil {
			return 0, err
		}
		return w.Write(p)
	}

	w.mu.Lock()
	n, err = w.buf.Write(p)
	w.mu.Unlock()
	return
}

// Close 将 buf 内的缓存数据全部写入硬盘, 然后关闭 socket
// 如果需要弃用 Writer 而不调用 Close 会导致最大丢失 BuffSize 字节数据
func (w *Writer) Close() error {
	w.mu.Lock()
	_ = w.buf.Flush()
	err := w.cur.Close()
	w.cur = (*os.File)(nil)
	w.buf = (*bufio.Writer)(nil)
	w.mu.Unlock()
	return err
}

func (w *Writer) open() error {
	fi, err := os.OpenFile(w.curName(), os.O_WRONLY|os.O_CREATE|os.O_APPEND, os.ModePerm)
	if err != nil {
		return err
	}
	w.mu.Lock()
	w.cur = fi
	w.buf = bufio.NewWriterSize(w.cur, BuffSize)
	w.mu.Unlock()
	return nil
}

func (w *Writer) curName() string {
	return filepath.Join(w.path, w.pre+"_"+w.date+w.suf)
}

func getDate() string {
	return time.Now().Format("2006_01_02")
}

func handlePath(path string) error {
	if _, err := os.Stat(path); err != nil {
		if err = os.MkdirAll(path, os.ModePerm); err != nil {
			return err
		}
		return err
	}
	return nil
}

type socketCache struct {
	cache map[string]io.WriteCloser
	mu    sync.Mutex
}

func (s *socketCache) Get(pre, suf, path string) (io.WriteCloser, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	name := pre + suf + path

	if cache, ok := s.cache[name]; ok {
		return cache, nil
	}

	w, err := NewRawWriter(pre, suf, path)
	if err != nil {
		return nil, err
	}
	s.cache[name] = w

	return w, nil
}

var (
	_socketCache = socketCache{cache: make(map[string]io.WriteCloser)}
)