Merge pull request #1145 from sirupsen/custom_buffer_pool

Add an API to plug a custom buffer free item mangement system

Author: dgsb

buffer_pool.go

@@ -0,0 +1,52 @@
+package logrus
+
+import (
+	"bytes"
+	"sync"
+)
+
+var (
+	bufferPool BufferPool
+)
+
+type BufferPool interface {
+	Put(*bytes.Buffer)
+	Get() *bytes.Buffer
+}
+
+type defaultPool struct {
+	pool *sync.Pool
+}
+
+func (p *defaultPool) Put(buf *bytes.Buffer) {
+	p.pool.Put(buf)
+}
+
+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) {
+	bufferPool = bp
+}
+
+func init() {
+	SetBufferPool(&defaultPool{
+		pool: &sync.Pool{
+			New: func() interface{} {
+				return new(bytes.Buffer)
+			},
+		},
+	})
+}

entry.go

@@ -13,7 +13,6 @@ import (
 )
 
 var (
-	bufferPool *sync.Pool
 
 	// qualified package name, cached at first use
 	logrusPackage string
@@ -31,12 +30,6 @@ const (
 )
 
 func init() {
-	bufferPool = &sync.Pool{
-		New: func() interface{} {
-			return new(bytes.Buffer)
-		},
-	}
-
 	// start at the bottom of the stack before the package-name cache is primed
 	minimumCallerDepth = 1
 }
@@ -243,9 +236,12 @@ func (entry Entry) log(level Level, msg string) {
 
 	entry.fireHooks()
 
-	buffer = bufferPool.Get().(*bytes.Buffer)
+	buffer = getBuffer()
+	defer func() {
+		entry.Buffer = nil
+		putBuffer(buffer)
+	}()
 	buffer.Reset()
-	defer bufferPool.Put(buffer)
 	entry.Buffer = buffer
 
 	entry.write()