Permalink
Switch branches/tags
Nothing to show
Find file
Fetching contributors…
Cannot retrieve contributors at this time
182 lines (159 sloc) 3.58 KB
package timber
import (
"bytes"
"compress/gzip"
"fmt"
"io"
"math/rand"
"os"
"strings"
"sync"
"text/template"
"time"
)
type FilenameFields struct {
Hostname string
Date time.Time
Pid int
Random int64
}
func GetFilenameFields() *FilenameFields {
h, _ := os.Hostname()
return &FilenameFields{
Hostname: h,
Date: time.Now(),
Pid: os.Getpid(),
Random: rand.Int63(),
}
}
func preprocessFilename(name string) string {
t := template.Must(template.New("filename").Parse(name))
buf := new(bytes.Buffer)
t.Execute(buf, GetFilenameFields())
return buf.String()
}
type FileWriter struct {
wr *BufferedWriter
cwr *countingWriter
BaseFilename string
currentFilename string
mutex *sync.RWMutex
RotateChan chan string // defaults to nil. receives previous filename on rotate
RotateSize int64 // rotate after RotateSize bytes have been written to the file
rotateTicker *time.Ticker
rotateReset chan int
}
// This writer has a buffer that I don't ever bother to flush, so it may take a while
// to see messages. Filenames ending in .gz will automatically be compressed on write.
// Filename string is proccessed through the template library using the FilenameFields
// struct.
func NewFileWriter(name string) (*FileWriter, error) {
w := &FileWriter{
BaseFilename: name,
mutex: new(sync.RWMutex),
}
if err := w.open(); err != nil {
return nil, err
}
return w, nil
}
func (w *FileWriter) open() error {
// No lock here
name := preprocessFilename(w.BaseFilename)
file, err := os.OpenFile(name, os.O_WRONLY|os.O_APPEND|os.O_CREATE, 0666)
if err != nil {
return fmt.Errorf("TIMBER! Can't open %v: %v", name, err)
}
var cwr = &countingWriter{file, 0}
var output io.WriteCloser = cwr
// Wrap in gz writer
if strings.HasSuffix(name, ".gz") {
output = &gzFileWriter{
gzip.NewWriter(output),
output,
}
}
// Locked from here
w.mutex.Lock()
defer w.mutex.Unlock()
if w.wr != nil {
w.wr.Close()
// send previous filename on rotate chan
if c := w.RotateChan; c != nil {
c <- w.currentFilename
}
}
w.currentFilename = name
w.cwr = cwr
w.wr, _ = NewBufferedWriter(output)
return nil
}
func (w *FileWriter) LogWrite(m string) {
w.mutex.RLock()
defer w.mutex.RUnlock()
w.wr.LogWrite(m)
w.checkSize()
}
func (w *FileWriter) Flush() error {
w.mutex.RLock()
defer w.mutex.RUnlock()
e := w.wr.Flush()
w.checkSize()
return e
}
// Close and re-open the file.
// You should use the timestamp in the filename if you're going to use rotation
func (w *FileWriter) Rotate() error {
return w.open()
}
// Automatically rotate every `d`
func (w *FileWriter) RotateEvery(d time.Duration) {
// reset ticker
w.mutex.Lock()
if w.rotateTicker != nil {
w.rotateTicker.Stop()
w.rotateReset <- 1
}
w.rotateTicker = time.NewTicker(d)
w.mutex.Unlock()
// trigger a rotate every X
go func() {
for {
select {
case <-w.rotateReset:
return
case <-w.rotateTicker.C:
w.Rotate()
}
}
}()
}
func (w *FileWriter) checkSize() {
if w.RotateSize > 0 && w.cwr.bytesWritten >= w.RotateSize {
go w.Rotate()
}
}
func (w *FileWriter) Close() {
w.mutex.Lock()
defer w.mutex.Unlock()
w.wr.Flush()
w.wr.Close()
w.wr = nil
}
type countingWriter struct {
io.WriteCloser
bytesWritten int64
}
func (w *countingWriter) Write(b []byte) (int, error) {
i, e := w.WriteCloser.Write(b)
w.bytesWritten += int64(i)
return i, e
}
type gzFileWriter struct {
*gzip.Writer // the compressor
file io.WriteCloser
}
func (w *gzFileWriter) Close() error {
w.Writer.Close()
return w.file.Close()
}