Merge pull request #1253 from edoger/buffer-pool

Add support for the logger private buffer pool.

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

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 = bufPool.Get()
-	buffer = getBuffer()
 	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()

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

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