Repository: fxamacker/cbor
Pull request #392 (2023-02-28 17:25:20)

[stream-mode branch] Add StreamEncoder.Close() to recycle buffer
`StreamEncoder` uses a buffer from `sync.Pool`, but that buffer was never returned to the pool for reuse.

This PR adds `StreamEncoder.Close()` to return buffer to the pool.

After calling `StreamEncoder.Close()`, all subsequent calls to encode return `ErrStreamClosed` error.



Patch:
diff --git a/stream_encode.go b/stream_encode.go
index 8cc0fde..c15a0df 100644
--- a/stream_encode.go
+++ b/stream_encode.go
@@ -4,10 +4,15 @@
 package cbor
 
 import (
+	"errors"
 	"io"
 	"math/big"
 )
 
+// ErrStreamClosed is the error indicating operations on a closed stream.
+// E.g. any call to encode after a call to Close will return ErrStreamClosed.
+var ErrStreamClosed = errors.New("cbor: operation on closed stream")
+
 // StreamEncoder provides low-level API for sequential encoding.
 //
 // Users of StreamEncoder should be familiar with CBOR data models.
@@ -17,6 +22,7 @@ import (
 // complete and well-formed.
 type StreamEncoder struct {
 	*Encoder
+	closed bool
 }
 
 // NewStreamEncoder returns a new StreamEncoder for sequential encoding.
@@ -31,26 +37,49 @@ func NewStreamEncoder(w io.Writer) *StreamEncoder {
 	}
 }
 
+// Close closes StreamEncoder and subsequence operations (except for Close)
+// will return ErrStreamClosed.
+func (se *StreamEncoder) Close() {
+	if se.closed {
+		return
+	}
+	putEncoderBuffer(se.Encoder.e)
+	se.Encoder = nil
+	se.closed = true
+}
+
 // Flush writes streamed data to underlying io.Writer.
 func (se *StreamEncoder) Flush() error {
+	if se.closed {
+		return ErrStreamClosed
+	}
 	_, err := se.Encoder.e.WriteTo(se.Encoder.w)
 	return err
 }
 
 // EncodeMapHead encodes CBOR map head of specified size.
 func (se *StreamEncoder) EncodeMapHead(size uint64) error {
+	if se.closed {
+		return ErrStreamClosed
+	}
 	encodeHead(se.Encoder.e, byte(cborTypeMap), size)
 	return nil
 }
 
 // EncodeArrayHead encodes CBOR array head of specified size.
 func (se *StreamEncoder) EncodeArrayHead(size uint64) error {
+	if se.closed {
+		return ErrStreamClosed
+	}
 	encodeHead(se.Encoder.e, byte(cborTypeArray), size)
 	return nil
 }
 
 // EncodeTagHead encodes CBOR tag head with num as tag number.
 func (se *StreamEncoder) EncodeTagHead(num uint64) error {
+	if se.closed {
+		return ErrStreamClosed
+	}
 	encodeHead(se.Encoder.e, byte(cborTypeTag), num)
 	return nil
 }
@@ -58,6 +87,9 @@ func (se *StreamEncoder) EncodeTagHead(num uint64) error {
 // EncodeRawBytes writes b to the underlying writer.
 // If b is an empty or nil byte slice, it is a no-op.
 func (se *StreamEncoder) EncodeRawBytes(b []byte) error {
+	if se.closed {
+		return ErrStreamClosed
+	}
 	if len(b) == 0 {
 		return nil
 	}
@@ -67,12 +99,18 @@ func (se *StreamEncoder) EncodeRawBytes(b []byte) error {
 
 // EncodeNil encodes CBOR nil.
 func (se *StreamEncoder) EncodeNil() error {
+	if se.closed {
+		return ErrStreamClosed
+	}
 	se.e.Write(cborNil)
 	return nil
 }
 
 // EncodeBool encodes bool as CBOR bool.
 func (se *StreamEncoder) EncodeBool(b bool) error {
+	if se.closed {
+		return ErrStreamClosed
+	}
 	bytes := cborTrue
 	if !b {
 		bytes = cborFalse
@@ -83,30 +121,45 @@ func (se *StreamEncoder) EncodeBool(b bool) error {
 
 // EncodeUint encodes uint as CBOR positive integer.
 func (se *StreamEncoder) EncodeUint(i uint) error {
+	if se.closed {
+		return ErrStreamClosed
+	}
 	encodeHead(se.e, byte(cborTypePositiveInt), uint64(i))
 	return nil
 }
 
 // EncodeUint8 encodes uint8 as CBOR positive integer.
 func (se *StreamEncoder) EncodeUint8(i uint8) error {
+	if se.closed {
+		return ErrStreamClosed
+	}
 	encodeHead(se.e, byte(cborTypePositiveInt), uint64(i))
 	return nil
 }
 
 // EncodeUint16 encodes uint16 as CBOR positive integer.
 func (se *StreamEncoder) EncodeUint16(i uint16) error {
+	if se.closed {
+		return ErrStreamClosed
+	}
 	encodeHead(se.e, byte(cborTypePositiveInt), uint64(i))
 	return nil
 }
 
 // EncodeUint32 encodes uint32 as CBOR positive integer.
 func (se *StreamEncoder) EncodeUint32(i uint32) error {
+	if se.closed {
+		return ErrStreamClosed
+	}
 	encodeHead(se.e, byte(cborTypePositiveInt), uint64(i))
 	return nil
 }
 
 // EncodeUint64 encodes uint64 as CBOR positive integer.
 func (se *StreamEncoder) EncodeUint64(i uint64) error {
+	if se.closed {
+		return ErrStreamClosed
+	}
 	encodeHead(se.e, byte(cborTypePositiveInt), i)
 	return nil
 }
@@ -133,6 +186,9 @@ func (se *StreamEncoder) EncodeInt32(i int32) error {
 
 // EncodeInt64 encodes int64 as CBOR positive or negtive integer.
 func (se *StreamEncoder) EncodeInt64(i int64) error {
+	if se.closed {
+		return ErrStreamClosed
+	}
 	t := cborTypePositiveInt
 	if i < 0 {
 		t = cborTypeNegativeInt
@@ -144,6 +200,9 @@ func (se *StreamEncoder) EncodeInt64(i int64) error {
 
 // EncodeBytes encodes byte slice as CBOR byte string.
 func (se *StreamEncoder) EncodeBytes(b []byte) error {
+	if se.closed {
+		return ErrStreamClosed
+	}
 	if b == nil {
 		se.e.Write(cborNil)
 		return nil
@@ -156,6 +215,9 @@ func (se *StreamEncoder) EncodeBytes(b []byte) error {
 
 // EncodeString encodes string as CBOR string.
 func (se *StreamEncoder) EncodeString(s string) error {
+	if se.closed {
+		return ErrStreamClosed
+	}
 	encodeHead(se.e, byte(cborTypeTextString), uint64(len(s)))
 	se.e.WriteString(s)
 	return nil
@@ -165,6 +227,9 @@ var bigOne = big.NewInt(1)
 
 // EncodeBigInt encodes big.Int as CBOR bignum (tag number 2 and 3).
 func (se *StreamEncoder) EncodeBigInt(v *big.Int) error {
+	if se.closed {
+		return ErrStreamClosed
+	}
 	if se.em.bigIntConvert == BigIntConvertShortest {
 		if v.IsUint64() {
 			encodeHead(se.e, byte(cborTypePositiveInt), v.Uint64())


Test patch:
diff --git a/stream_encode_test.go b/stream_encode_test.go
index 02a7c65..583858d 100644
--- a/stream_encode_test.go
+++ b/stream_encode_test.go
@@ -5,6 +5,7 @@ package cbor
 
 import (
 	"bytes"
+	"errors"
 	"math/big"
 	"testing"
 )
@@ -36,6 +37,7 @@ func TestStreamEncodeMap(t *testing.T) {
 
 		var buf bytes.Buffer
 		se := NewStreamEncoder(&buf)
+		defer se.Close()
 
 		err := se.EncodeMapHead(2)
 		if err != nil {
@@ -117,6 +119,7 @@ func TestStreamEncodeMap(t *testing.T) {
 
 		var buf bytes.Buffer
 		se := em.NewStreamEncoder(&buf)
+		defer se.Close()
 
 		err = se.EncodeMapHead(2)
 		if err != nil {
@@ -183,6 +186,7 @@ func TestStreamEncodeArray(t *testing.T) {
 
 		var buf bytes.Buffer
 		se := NewStreamEncoder(&buf)
+		defer se.Close()
 
 		err := se.EncodeArrayHead(2)
 		if err != nil {
@@ -239,6 +243,7 @@ func TestStreamEncodeArray(t *testing.T) {
 
 		var buf bytes.Buffer
 		se := em.NewStreamEncoder(&buf)
+		defer se.Close()
 
 		err = se.EncodeArrayHead(2)
 		if err != nil {
@@ -290,6 +295,7 @@ func TestStreamEncodeTag(t *testing.T) {
 
 		var buf bytes.Buffer
 		se := NewStreamEncoder(&buf)
+		defer se.Close()
 
 		err := se.EncodeTagHead(128)
 		if err != nil {
@@ -346,6 +352,7 @@ func TestStreamEncodeTag(t *testing.T) {
 
 		var buf bytes.Buffer
 		se := em.NewStreamEncoder(&buf)
+		defer se.Close()
 
 		err = se.EncodeTagHead(128)
 		if err != nil {
@@ -383,6 +390,7 @@ func TestStreamEncodeNil(t *testing.T) {
 
 	var buf bytes.Buffer
 	se := NewStreamEncoder(&buf)
+	defer se.Close()
 
 	err := se.EncodeNil()
 	if err != nil {
@@ -416,6 +424,7 @@ func TestStreamEncodeRawBytes(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			var buf bytes.Buffer
 			se := NewStreamEncoder(&buf)
+			defer se.Close()
 
 			err := se.EncodeRawBytes(tc.value)
 			if err != nil {
@@ -450,6 +459,7 @@ func TestStreamEncodeBool(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			var buf bytes.Buffer
 			se := NewStreamEncoder(&buf)
+			defer se.Close()
 
 			err := se.EncodeBool(tc.value)
 			if err != nil {
@@ -488,6 +498,7 @@ func TestStreamEncodeUint(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			var buf bytes.Buffer
 			se := NewStreamEncoder(&buf)
+			defer se.Close()
 
 			err := se.EncodeUint(tc.value)
 			if err != nil {
@@ -524,6 +535,7 @@ func TestStreamEncodeUint8(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			var buf bytes.Buffer
 			se := NewStreamEncoder(&buf)
+			defer se.Close()
 
 			err := se.EncodeUint8(tc.value)
 			if err != nil {
@@ -560,6 +572,7 @@ func TestStreamEncodeUint16(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			var buf bytes.Buffer
 			se := NewStreamEncoder(&buf)
+			defer se.Close()
 
 			err := se.EncodeUint16(tc.value)
 			if err != nil {
@@ -597,6 +610,7 @@ func TestStreamEncodeUint32(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			var buf bytes.Buffer
 			se := NewStreamEncoder(&buf)
+			defer se.Close()
 
 			err := se.EncodeUint32(tc.value)
 			if err != nil {
@@ -635,6 +649,7 @@ func TestStreamEncodeUint64(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			var buf bytes.Buffer
 			se := NewStreamEncoder(&buf)
+			defer se.Close()
 
 			err := se.EncodeUint64(tc.value)
 			if err != nil {
@@ -683,6 +698,7 @@ func TestStreamEncodeInt(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			var buf bytes.Buffer
 			se := NewStreamEncoder(&buf)
+			defer se.Close()
 
 			err := se.EncodeInt(tc.value)
 			if err != nil {
@@ -722,6 +738,7 @@ func TestStreamEncodeInt8(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			var buf bytes.Buffer
 			se := NewStreamEncoder(&buf)
+			defer se.Close()
 
 			err := se.EncodeInt8(tc.value)
 			if err != nil {
@@ -764,6 +781,7 @@ func TestStreamEncodeInt16(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			var buf bytes.Buffer
 			se := NewStreamEncoder(&buf)
+			defer se.Close()
 
 			err := se.EncodeInt16(tc.value)
 			if err != nil {
@@ -809,6 +827,7 @@ func TestStreamEncodeInt32(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			var buf bytes.Buffer
 			se := NewStreamEncoder(&buf)
+			defer se.Close()
 
 			err := se.EncodeInt32(tc.value)
 			if err != nil {
@@ -857,6 +876,7 @@ func TestStreamEncodeInt64(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			var buf bytes.Buffer
 			se := NewStreamEncoder(&buf)
+			defer se.Close()
 
 			err := se.EncodeInt64(tc.value)
 			if err != nil {
@@ -899,6 +919,7 @@ func TestStreamEncodeBigIntConvertNone(t *testing.T) {
 
 			var buf bytes.Buffer
 			se := dm.NewStreamEncoder(&buf)
+			defer se.Close()
 
 			err := se.EncodeBigInt(tc.value)
 			if err != nil {
@@ -941,6 +962,7 @@ func TestStreamEncodeBigIntConvertShortest(t *testing.T) {
 
 			var buf bytes.Buffer
 			se := dm.NewStreamEncoder(&buf)
+			defer se.Close()
 
 			err := se.EncodeBigInt(tc.value)
 			if err != nil {
@@ -976,6 +998,7 @@ func TestStreamEncodeBytes(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			var buf bytes.Buffer
 			se := NewStreamEncoder(&buf)
+			defer se.Close()
 
 			err := se.EncodeBytes(tc.value)
 			if err != nil {
@@ -1010,6 +1033,7 @@ func TestStreamEncodeString(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			var buf bytes.Buffer
 			se := NewStreamEncoder(&buf)
+			defer se.Close()
 
 			err := se.EncodeString(tc.value)
 			if err != nil {
@@ -1027,3 +1051,114 @@ func TestStreamEncodeString(t *testing.T) {
 		})
 	}
 }
+
+func TestStreamEncodeCloseError(t *testing.T) {
+	var buf bytes.Buffer
+	se := NewStreamEncoder(&buf)
+
+	se.Close()
+
+	var err error
+
+	err = se.EncodeMapHead(0)
+	if !errors.Is(err, ErrStreamClosed) {
+		t.Errorf("EncodeMapHead() returned error %v, want %v", err, ErrStreamClosed)
+	}
+
+	err = se.EncodeArrayHead(0)
+	if !errors.Is(err, ErrStreamClosed) {
+		t.Errorf("EncodeArrayHead() returned error %v, want %v", err, ErrStreamClosed)
+	}
+
+	err = se.EncodeTagHead(0)
+	if !errors.Is(err, ErrStreamClosed) {
+		t.Errorf("EncodeTagHead() returned error %v, want %v", err, ErrStreamClosed)
+	}
+
+	err = se.EncodeRawBytes(nil)
+	if !errors.Is(err, ErrStreamClosed) {
+		t.Errorf("EncodeRawBytes() returned error %v, want %v", err, ErrStreamClosed)
+	}
+
+	err = se.EncodeNil()
+	if !errors.Is(err, ErrStreamClosed) {
+		t.Errorf("EncodeNil() returned error %v, want %v", err, ErrStreamClosed)
+	}
+
+	err = se.EncodeBool(true)
+	if !errors.Is(err, ErrStreamClosed) {
+		t.Errorf("EncodeBool() returned error %v, want %v", err, ErrStreamClosed)
+	}
+
+	err = se.EncodeUint(0)
+	if !errors.Is(err, ErrStreamClosed) {
+		t.Errorf("EncodeUint() returned error %v, want %v", err, ErrStreamClosed)
+	}
+
+	err = se.EncodeUint8(0)
+	if !errors.Is(err, ErrStreamClosed) {
+		t.Errorf("EncodeUint8() returned error %v, want %v", err, ErrStreamClosed)
+	}
+
+	err = se.EncodeUint16(0)
+	if !errors.Is(err, ErrStreamClosed) {
+		t.Errorf("EncodeUint16() returned error %v, want %v", err, ErrStreamClosed)
+	}
+
+	err = se.EncodeUint32(0)
+	if !errors.Is(err, ErrStreamClosed) {
+		t.Errorf("EncodeUint32() returned error %v, want %v", err, ErrStreamClosed)
+	}
+
+	err = se.EncodeUint64(0)
+	if !errors.Is(err, ErrStreamClosed) {
+		t.Errorf("EncodeUint64() returned error %v, want %v", err, ErrStreamClosed)
+	}
+
+	err = se.EncodeInt(0)
+	if !errors.Is(err, ErrStreamClosed) {
+		t.Errorf("EncodeInt() returned error %v, want %v", err, ErrStreamClosed)
+	}
+
+	err = se.EncodeInt8(0)
+	if !errors.Is(err, ErrStreamClosed) {
+		t.Errorf("EncodeInt8() returned error %v, want %v", err, ErrStreamClosed)
+	}
+
+	err = se.EncodeInt16(0)
+	if !errors.Is(err, ErrStreamClosed) {
+		t.Errorf("EncodeInt16() returned error %v, want %v", err, ErrStreamClosed)
+	}
+
+	err = se.EncodeInt32(0)
+	if !errors.Is(err, ErrStreamClosed) {
+		t.Errorf("EncodeInt32() returned error %v, want %v", err, ErrStreamClosed)
+	}
+
+	err = se.EncodeInt64(0)
+	if !errors.Is(err, ErrStreamClosed) {
+		t.Errorf("EncodeInt64() returned error %v, want %v", err, ErrStreamClosed)
+	}
+
+	err = se.EncodeBytes(nil)
+	if !errors.Is(err, ErrStreamClosed) {
+		t.Errorf("EncodeBytes() returned error %v, want %v", err, ErrStreamClosed)
+	}
+
+	err = se.EncodeString("")
+	if !errors.Is(err, ErrStreamClosed) {
+		t.Errorf("EncodeString() returned error %v, want %v", err, ErrStreamClosed)
+	}
+
+	err = se.EncodeBigInt(bigOne)
+	if !errors.Is(err, ErrStreamClosed) {
+		t.Errorf("EncodeBigInt() returned error %v, want %v", err, ErrStreamClosed)
+	}
+
+	err = se.Flush()
+	if !errors.Is(err, ErrStreamClosed) {
+		t.Errorf("Flush() returned error %v, want %v", err, ErrStreamClosed)
+	}
+
+	se.Close() // Close on closed stream is no-op.
+}