From 1818363d79cc77045d0724ed85146d5446ae925d Mon Sep 17 00:00:00 2001 From: Qingshan Luo Date: Tue, 20 Apr 2021 10:48:30 +0800 Subject: [PATCH] Add support for the logger private buffer pool. --- buffer_pool.go | 9 --------- entry.go | 14 +++++++++++--- logger.go | 10 ++++++++++ logger_test.go | 28 ++++++++++++++++++++++++++++ 4 files changed, 49 insertions(+), 12 deletions(-) diff --git a/buffer_pool.go b/buffer_pool.go index 4545dec..c7787f7 100644 --- a/buffer_pool.go +++ b/buffer_pool.go @@ -26,15 +26,6 @@ func (p *defaultPool) Get() *bytes.Buffer { return p.pool.Get().(*bytes.Buffer) } -func getBuffer() *bytes.Buffer { - return bufferPool.Get() -} - -func putBuffer(buf *bytes.Buffer) { - buf.Reset() - bufferPool.Put(buf) -} - // SetBufferPool allows to replace the default logrus buffer pool // to better meets the specific needs of an application. func SetBufferPool(bp BufferPool) { diff --git a/entry.go b/entry.go index 07a1e5f..0a5a832 100644 --- a/entry.go +++ b/entry.go @@ -232,6 +232,7 @@ func (entry *Entry) log(level Level, msg string) { newEntry.Logger.mu.Lock() reportCaller := newEntry.Logger.ReportCaller + bufPool := newEntry.getBufferPool() newEntry.Logger.mu.Unlock() if reportCaller { @@ -239,11 +240,11 @@ func (entry *Entry) log(level Level, msg string) { } newEntry.fireHooks() - - buffer = getBuffer() + buffer = bufPool.Get() defer func() { newEntry.Buffer = nil - putBuffer(buffer) + buffer.Reset() + bufPool.Put(buffer) }() buffer.Reset() newEntry.Buffer = buffer @@ -260,6 +261,13 @@ func (entry *Entry) log(level Level, msg string) { } } +func (entry *Entry) getBufferPool() (pool BufferPool) { + if entry.Logger.BufferPool != nil { + return entry.Logger.BufferPool + } + return bufferPool +} + func (entry *Entry) fireHooks() { var tmpHooks LevelHooks entry.Logger.mu.Lock() diff --git a/logger.go b/logger.go index 3377044..91b460e 100644 --- a/logger.go +++ b/logger.go @@ -44,6 +44,9 @@ type Logger struct { entryPool sync.Pool // Function to exit the application, defaults to `os.Exit()` ExitFunc exitFunc + // The buffer pool used to format the log. If it is nil, the default global + // buffer pool will be used. + BufferPool BufferPool } type exitFunc func(int) @@ -402,3 +405,10 @@ func (logger *Logger) ReplaceHooks(hooks LevelHooks) LevelHooks { logger.mu.Unlock() return oldHooks } + +// SetBufferPool sets the logger buffer pool. +func (logger *Logger) SetBufferPool(pool BufferPool) { + logger.mu.Lock() + defer logger.mu.Unlock() + logger.BufferPool = pool +} diff --git a/logger_test.go b/logger_test.go index 8afb2fe..595c682 100644 --- a/logger_test.go +++ b/logger_test.go @@ -67,3 +67,31 @@ func TestWarninglnNotEqualToWarning(t *testing.T) { assert.NotEqual(t, buf.String(), bufln.String(), "Warning() and Wantingln() should not be equal") } + +type testBufferPool struct { + buffers []*bytes.Buffer + get int +} + +func (p *testBufferPool) Get() *bytes.Buffer { + p.get++ + return new(bytes.Buffer) +} + +func (p *testBufferPool) Put(buf *bytes.Buffer) { + p.buffers = append(p.buffers, buf) +} + +func TestLogger_SetBufferPool(t *testing.T) { + out := &bytes.Buffer{} + l := New() + l.SetOutput(out) + + pool := new(testBufferPool) + l.SetBufferPool(pool) + + l.Info("test") + + assert.Equal(t, pool.get, 1, "Logger.SetBufferPool(): The BufferPool.Get() must be called") + assert.Len(t, pool.buffers, 1, "Logger.SetBufferPool(): The BufferPool.Put() must be called") +}