From 1818363d79cc77045d0724ed85146d5446ae925d Mon Sep 17 00:00:00 2001
From: Qingshan Luo <edoger@qq.com>
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")
+}