From 64a59449f3bfee6956be931c777de21742a8560c Mon Sep 17 00:00:00 2001
From: David Bariod <david.bariod@qonto.eu>
Date: Thu, 28 May 2020 10:47:50 +0200
Subject: [PATCH] Add an API to plug a custom buffer free item mangement system

---
 buffer_pool.go | 52 ++++++++++++++++++++++++++++++++++++++++++++++++++
 entry.go       | 14 +++++---------
 2 files changed, 57 insertions(+), 9 deletions(-)
 create mode 100644 buffer_pool.go

diff --git a/buffer_pool.go b/buffer_pool.go
new file mode 100644
index 0000000..4545dec
--- /dev/null
+++ b/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)
+			},
+		},
+	})
+}
diff --git a/entry.go b/entry.go
index f6e062a..5a5cbfe 100644
--- a/entry.go
+++ b/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()