Add support for the logger private buffer pool.

This commit is contained in:
Qingshan Luo 2021-04-20 10:48:30 +08:00
parent fdf1618bf7
commit 1818363d79
4 changed files with 49 additions and 12 deletions

View File

@ -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) {

View File

@ -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()

View File

@ -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
}

View File

@ -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")
}