From b68e8c5e6a34cb8753a8e5556a992e2899c3ce48 Mon Sep 17 00:00:00 2001
From: eric-millin <110051399+eric-millin@users.noreply.github.com>
Date: Wed, 26 Apr 2023 03:46:32 -0400
Subject: [PATCH] Use pooled buffers for serialization (#81)

---
 CHANGELOG.md                      |  6 ++++
 json_serialization_writer.go      | 49 +++++++++++++++++++++++++------
 json_serialization_writer_test.go | 17 +++++++++--
 3 files changed, 60 insertions(+), 12 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 5543bdb..66be8c8 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -12,6 +12,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 
 
+## [0.9.3] - 2023-04-24
+
+### Changed
+
+- Use buffer pool for `JsonSerializationWriter`.
+
 ## [0.9.2] - 2023-04-17
 
 ### Changed
diff --git a/json_serialization_writer.go b/json_serialization_writer.go
index 6b526f5..595b91f 100644
--- a/json_serialization_writer.go
+++ b/json_serialization_writer.go
@@ -6,6 +6,7 @@ import (
 	"encoding/json"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/google/uuid"
@@ -14,6 +15,12 @@ import (
 	absser "github.com/microsoft/kiota-abstractions-go/serialization"
 )
 
+var buffPool = sync.Pool{
+	New: func() interface{} {
+		return new(bytes.Buffer)
+	},
+}
+
 // JsonSerializationWriter implements SerializationWriter for JSON.
 type JsonSerializationWriter struct {
 	writer                     *bytes.Buffer
@@ -26,13 +33,22 @@ type JsonSerializationWriter struct {
 // NewJsonSerializationWriter creates a new instance of the JsonSerializationWriter.
 func NewJsonSerializationWriter() *JsonSerializationWriter {
 	return &JsonSerializationWriter{
-		writer:           new(bytes.Buffer),
+		writer:           buffPool.Get().(*bytes.Buffer),
 		separatorIndices: make([]int, 0),
 	}
 }
+func (w *JsonSerializationWriter) getWriter() *bytes.Buffer {
+	if w.writer == nil {
+		panic("The writer has already been closed. Call Reset instead of Close to reuse it or instantiate a new one.")
+	}
+
+	return w.writer
+}
 func (w *JsonSerializationWriter) writeRawValue(value ...string) {
+	writer := w.getWriter()
+
 	for _, v := range value {
-		w.writer.WriteString(v)
+		writer.WriteString(v)
 	}
 }
 func (w *JsonSerializationWriter) writeStringValue(value string) {
@@ -48,7 +64,7 @@ func (w *JsonSerializationWriter) writePropertyName(key string) {
 	w.writeRawValue("\"", key, "\":")
 }
 func (w *JsonSerializationWriter) writePropertySeparator() {
-	w.separatorIndices = append(w.separatorIndices, w.writer.Len())
+	w.separatorIndices = append(w.separatorIndices, w.getWriter().Len())
 	w.writeRawValue(",")
 }
 func (w *JsonSerializationWriter) writeArrayStart() {
@@ -607,16 +623,16 @@ func (w *JsonSerializationWriter) WriteCollectionOfInt8Values(key string, collec
 
 // GetSerializedContent returns the resulting byte array from the serialization writer.
 func (w *JsonSerializationWriter) GetSerializedContent() ([]byte, error) {
-	trimmed := w.writer.Bytes()
+	trimmed := w.getWriter().Bytes()
 	buffLen := len(trimmed)
 
 	for i := len(w.separatorIndices) - 1; i >= 0; i-- {
 		idx := w.separatorIndices[i]
 
 		if idx == buffLen-1 {
-			trimmed = trimmed[0:idx]
+			trimmed = trimmed[:idx]
 		} else if trimmed[idx+1] == byte(']') || trimmed[idx+1] == byte('}') {
-			trimmed = append(trimmed[0:idx], trimmed[idx+1:]...)
+			trimmed = append(trimmed[:idx], trimmed[idx+1:]...)
 		}
 	}
 
@@ -780,9 +796,24 @@ func (w *JsonSerializationWriter) WriteAdditionalData(value map[string]interface
 	return err
 }
 
-// Close clears the internal buffer.
+// Reset sets the internal buffer to empty, allowing the writer to be reused.
+func (w *JsonSerializationWriter) Reset() error {
+	w.getWriter().Reset()
+	w.separatorIndices = w.separatorIndices[:0]
+	return nil
+}
+
+// Close relases the internal buffer. Subsequent calls to the writer will panic.
 func (w *JsonSerializationWriter) Close() error {
-	w.writer = new(bytes.Buffer)
-	w.separatorIndices = make([]int, 0)
+	if w.writer == nil {
+		return nil
+	}
+
+	w.writer.Reset()
+	buffPool.Put(w.writer)
+
+	w.writer = nil
+	w.separatorIndices = nil
+
 	return nil
 }
diff --git a/json_serialization_writer_test.go b/json_serialization_writer_test.go
index 906b7a3..4254de3 100644
--- a/json_serialization_writer_test.go
+++ b/json_serialization_writer_test.go
@@ -129,14 +129,14 @@ func TestDoubleEscapeFailure(t *testing.T) {
 	assert.Equal(t, fmt.Sprintf("\"key\":%q", value), string(result[:]))
 }
 
-func TestBufferClose(t *testing.T) {
+func TestReset(t *testing.T) {
 	serializer := NewJsonSerializationWriter()
 	value := "W/\"CQAAABYAAAAs+XSiyjZdS4Rhtwk0v1pGAAC5bsJ2\""
 	serializer.WriteStringValue("key", &value)
 	result, err := serializer.GetSerializedContent()
 	assert.Nil(t, err)
 	assert.True(t, len(result) > 0)
-	serializer.Close()
+	serializer.Reset()
 	assert.True(t, len(result) > 0)
 	empty, err := serializer.GetSerializedContent()
 	assert.Nil(t, err)
@@ -148,6 +148,18 @@ func TestBufferClose(t *testing.T) {
 	assert.True(t, len(notEmpty) > 0)
 }
 
+func TestClose(t *testing.T) {
+	serializer := NewJsonSerializationWriter()
+	serializer.Close()
+	assert.Panics(t, func() {
+		serializer.GetSerializedContent()
+	})
+	assert.Panics(t, serializer.writer.Reset)
+	assert.NotPanics(t, func() {
+		serializer.Close()
+	})
+}
+
 func TestJsonSerializationWriterHonoursInterface(t *testing.T) {
 	instance := NewJsonSerializationWriter()
 	assert.Implements(t, (*absser.SerializationWriter)(nil), instance)
@@ -194,7 +206,6 @@ func TestWriteInvalidAdditionalData(t *testing.T) {
 	err := serializer.WriteAdditionalData(adlData)
 	assert.Nil(t, err)
 	result, err := serializer.GetSerializedContent()
-	assert.NoError(t, err)
 
 	stringResult := string(result[:])
 	assert.Contains(t, stringResult, "\"pointer_node\":")