package log import ( "bytes" "io" "sync" "time" "go.uber.org/atomic" ) // LineBufferedLogger buffers log lines to be flushed periodically. Without a line buffer, Log() will call the write // syscall for every log line which is expensive if logging thousands of lines per second. type LineBufferedLogger struct { buf *threadsafeBuffer entries atomic.Uint32 cap uint32 w io.Writer onFlush func(entries uint32) } // Size returns the number of entries in the buffer. func (l *LineBufferedLogger) Size() uint32 { return l.entries.Load() } // Write writes the given bytes to the line buffer, and increments the entries counter. // If the buffer is full (entries == cap), it will be flushed, and the entries counter reset. func (l *LineBufferedLogger) Write(p []byte) (n int, err error) { // when we've filled the buffer, flush it if l.Size() >= l.cap { // Flush resets the size to 0 if err := l.Flush(); err != nil { l.buf.Reset() return 0, err } } l.entries.Inc() return l.buf.Write(p) } // Flush forces the buffer to be written to the underlying writer. func (l *LineBufferedLogger) Flush() error { sz := l.Size() if sz <= 0 { return nil } // reset the counter l.entries.Store(0) // WriteTo() calls Reset() on the underlying buffer, so it's not needed here _, err := l.buf.WriteTo(l.w) // only call OnFlush callback if write was successful if err == nil && l.onFlush != nil { l.onFlush(sz) } return err } type LineBufferedLoggerOption func(*LineBufferedLogger) // WithFlushPeriod creates a new LineBufferedLoggerOption that sets the flush period for the LineBufferedLogger. func WithFlushPeriod(d time.Duration) LineBufferedLoggerOption { return func(l *LineBufferedLogger) { go func() { tick := time.NewTicker(d) defer tick.Stop() for range tick.C { l.Flush() } }() } } // WithFlushCallback allows for a callback function to be executed when Flush() is called. // The length of the buffer at the time of the Flush() will be passed to the function. func WithFlushCallback(fn func(entries uint32)) LineBufferedLoggerOption { return func(l *LineBufferedLogger) { l.onFlush = fn } } // WithPrellocatedBuffer preallocates a buffer to reduce GC cycles and slice resizing. func WithPrellocatedBuffer(size uint32) LineBufferedLoggerOption { return func(l *LineBufferedLogger) { l.buf = newThreadsafeBuffer(bytes.NewBuffer(make([]byte, 0, size))) } } // NewLineBufferedLogger creates a new LineBufferedLogger with a configured capacity. // Lines are flushed when the context is done, the buffer is full, or the flush period is reached. func NewLineBufferedLogger(w io.Writer, cap uint32, opts ...LineBufferedLoggerOption) *LineBufferedLogger { l := &LineBufferedLogger{ w: w, buf: newThreadsafeBuffer(bytes.NewBuffer([]byte{})), cap: cap, } for _, opt := range opts { opt(l) } return l } // threadsafeBuffer wraps the non-threadsafe bytes.Buffer. type threadsafeBuffer struct { sync.RWMutex buf *bytes.Buffer } // Read returns the contents of the buffer. func (t *threadsafeBuffer) Read(p []byte) (n int, err error) { t.RLock() defer t.RUnlock() return t.buf.Read(p) } // Write writes the given bytes to the underlying writer. func (t *threadsafeBuffer) Write(p []byte) (n int, err error) { t.Lock() defer t.Unlock() return t.buf.Write(p) } // WriteTo writes the buffered lines to the given writer. func (t *threadsafeBuffer) WriteTo(w io.Writer) (n int64, err error) { t.Lock() defer t.Unlock() return t.buf.WriteTo(w) } // Reset resets the buffer to be empty, // but it retains the underlying storage for use by future writes. // Reset is the same as Truncate(0). func (t *threadsafeBuffer) Reset() { t.Lock() defer t.Unlock() t.buf.Reset() } // newThreadsafeBuffer returns a new threadsafeBuffer wrapping the given bytes.Buffer. func newThreadsafeBuffer(buf *bytes.Buffer) *threadsafeBuffer { return &threadsafeBuffer{ buf: buf, } }