Expose stream protocol method to close Writer

thriftrw · Jun 28, 2021 · 44d42e0 · 44d42e0
1 parent 8405012
commit 44d42e0
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 0 deletions.
diff --git a/gen/benchmark_test.go b/gen/benchmark_test.go
@@ -137,6 +137,7 @@ func BenchmarkRoundTrip(b *testing.B) {
 			give, ok := bb.give.(streamingThriftType)
 			require.True(b, ok)
 			require.NoError(b, give.Encode(writer), "StreamEncode")
+			writer.Close()
 		}
 	}
 

diff --git a/protocol/binary/stream_writer.go b/protocol/binary/stream_writer.go
@@ -231,3 +231,9 @@ func (sw *StreamWriter) WriteMapBegin(m stream.MapHeader) error {
 func (sw *StreamWriter) WriteMapEnd() error {
 	return nil
 }
+
+// Close frees up the resources used by the StreamWriter and returns it back
+// to the pool.
+func (sw *StreamWriter) Close() {
+	ReturnStreamWriter(sw)
+}
diff --git a/protocol/stream/stream.go b/protocol/stream/stream.go
@@ -95,6 +95,7 @@ type Writer interface {
 	WriteSetEnd() error
 	WriteListBegin(l ListHeader) error
 	WriteListEnd() error
+	Close()
 }
 
 // Reader defines an decoder for a Thrift value, implemented in a streaming