From 3ed797ca1434178ba9998a603bb0086d8db3689d Mon Sep 17 00:00:00 2001
From: Ox Cart <ox.to.a.cart@gmail.com>
Date: Thu, 13 Oct 2022 19:23:57 -0500
Subject: [PATCH] Don't put huge buffers back into the pool

---
 golog.go       | 23 ++++++++++-------------
 json_output.go |  4 ++--
 pool.go        | 29 +++++++++++++++++++++++++++++
 pool_test.go   | 21 +++++++++++++++++++++
 text_output.go |  4 ++--
 5 files changed, 64 insertions(+), 17 deletions(-)
 create mode 100644 pool.go
 create mode 100644 pool_test.go

diff --git a/golog.go b/golog.go
index 7bd5cf9..5cbcd45 100644
--- a/golog.go
+++ b/golog.go
@@ -6,18 +6,18 @@
 //
 // - The following linker flag is set
 //
-//     -X github.com/getlantern/golog.linkerFlagEnableTraceThroughLinker=true
+//		-X github.com/getlantern/golog.linkerFlagEnableTraceThroughLinker=true
 //
-// - Optionally, you can also set a comma-separated list of prefixes to trace
-//   through the following linker flag
+//	  - Optionally, you can also set a comma-separated list of prefixes to trace
+//	    through the following linker flag
 //
-//     -X github.com/getlantern/golog.linkerFlagTracePrefixes=prefix1,prefix2
+//	    -X github.com/getlantern/golog.linkerFlagTracePrefixes=prefix1,prefix2
 //
-// - Or, alternatively, trace logging can also be enable if "TRACE=true"
-//   environment variable is set
+//	  - Or, alternatively, trace logging can also be enable if "TRACE=true"
+//	    environment variable is set
 //
-// - Optionally, you can also set a comma-separated list of prefixes to trace
-//   through the "TRACE" environment variable like this: "TRACE=prefix1,prefix2"
+//	  - Optionally, you can also set a comma-separated list of prefixes to trace
+//	    through the "TRACE" environment variable like this: "TRACE=prefix1,prefix2"
 //
 // A stack dump will be printed after the message if "PRINT_STACK=true".
 package golog
@@ -39,7 +39,6 @@ import (
 	"github.com/getlantern/errors"
 	"github.com/getlantern/hidden"
 	"github.com/getlantern/ops"
-	"github.com/oxtoacart/bpool"
 )
 
 const (
@@ -68,8 +67,6 @@ var (
 	reporters      []ErrorReporter
 	reportersMutex sync.RWMutex
 
-	bufferPool = bpool.NewBufferPool(200)
-
 	onFatal atomic.Value
 
 	// enableTraceThroughLinker is set through a linker flag. It's used to
@@ -470,8 +467,8 @@ func argToString(arg interface{}) string {
 		if ml, isMultiline := arg.(MultiLine); !isMultiline {
 			return fmt.Sprintf("%v", arg)
 		} else {
-			buf := bufferPool.Get()
-			defer bufferPool.Put(buf)
+			buf := getBuffer()
+			defer returnBuffer(buf)
 			mlp := ml.MultiLinePrinter()
 			for {
 				more := mlp(buf)
diff --git a/json_output.go b/json_output.go
index 40a3bc3..d4129ee 100644
--- a/json_output.go
+++ b/json_output.go
@@ -46,8 +46,8 @@ func (o *jsonOutput) print(writer io.Writer, prefix string, skipFrames int, prin
 	cleanPrefix := prefix[0 : len(prefix)-2] // prefix contains ': ' at the end, strip it
 	event := Event{Component: cleanPrefix, Severity: severity, Caller: o.caller(skipFrames), Context: values}
 	if printStack {
-		buf := bufferPool.Get()
-		defer bufferPool.Put(buf)
+		buf := getBuffer()
+		defer returnBuffer(buf)
 		_ = writeStack(buf, o.pc)
 		event.Stack = buf.String()
 	}
diff --git a/pool.go b/pool.go
new file mode 100644
index 0000000..c332b34
--- /dev/null
+++ b/pool.go
@@ -0,0 +1,29 @@
+package golog
+
+import (
+	"bytes"
+
+	"github.com/oxtoacart/bpool"
+)
+
+const (
+	// Together, these bound the buffer pool to 307 KB
+	bufferPoolSize = 200
+	maxBufferSize  = 1024
+)
+
+var (
+	_bufferPool = bpool.NewBufferPool(bufferPoolSize)
+)
+
+func getBuffer() *bytes.Buffer {
+	return _bufferPool.Get()
+}
+
+func returnBuffer(buf *bytes.Buffer) bool {
+	if buf.Cap() <= maxBufferSize {
+		_bufferPool.Put(buf)
+		return true
+	}
+	return false
+}
diff --git a/pool_test.go b/pool_test.go
new file mode 100644
index 0000000..f298bf9
--- /dev/null
+++ b/pool_test.go
@@ -0,0 +1,21 @@
+package golog
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestPool(t *testing.T) {
+	buf := _bufferPool.Get()
+	require.NotNil(t, buf)
+	// Write 768 bytes (the max before the buffer's capacity exceeds )
+	for i := 0; i < maxBufferSize; i++ {
+		buf.WriteByte('C')
+	}
+	require.True(t, returnBuffer(buf))
+	for i := 0; i <= maxBufferSize; i++ {
+		buf.WriteByte('C')
+	}
+	require.False(t, returnBuffer(buf))
+}
diff --git a/text_output.go b/text_output.go
index d1ca8ee..73647b2 100644
--- a/text_output.go
+++ b/text_output.go
@@ -37,8 +37,8 @@ func (o *textOutput) Debug(prefix string, skipFrames int, printStack bool, sever
 }
 
 func (o *textOutput) print(writer io.Writer, prefix string, skipFrames int, printStack bool, severity string, arg interface{}, values map[string]interface{}) {
-	buf := bufferPool.Get()
-	defer bufferPool.Put(buf)
+	buf := getBuffer()
+	defer returnBuffer(buf)
 
 	GetPrepender()(buf)
 	linePrefix := o.linePrefix(prefix, skipFrames)