From c9605ac81e18d45b2ca7c2ba0f37f02ba9908feb Mon Sep 17 00:00:00 2001
From: dianale31 <31860231+dianale31@users.noreply.github.com>
Date: Mon, 28 Jun 2021 18:33:01 -0400
Subject: [PATCH] Expose stream protocol method to close Writer (#506)

Provides a stream protocol method to close out the resources of the given Writer.

The corresponding Close() method on Reader will be implemented in a follow-up PR.

Refs DR-609
---
 gen/benchmark_test.go            | 1 +
 protocol/binary/stream_writer.go | 7 +++++++
 protocol/stream/stream.go        | 1 +
 3 files changed, 9 insertions(+)

diff --git a/gen/benchmark_test.go b/gen/benchmark_test.go
index 25043984..d93b065c 100644
--- a/gen/benchmark_test.go
+++ b/gen/benchmark_test.go
@@ -138,6 +138,7 @@ func BenchmarkRoundTrip(b *testing.B) {
 			give, ok := bb.give.(streamingThriftType)
 			require.True(b, ok)
 			require.NoError(b, give.Encode(writer), "StreamEncode")
+			require.NoError(b, writer.Close(), "Close StreamWriter")
 		}
 	}
 
diff --git a/protocol/binary/stream_writer.go b/protocol/binary/stream_writer.go
index c18dfb43..8bc9b9ca 100644
--- a/protocol/binary/stream_writer.go
+++ b/protocol/binary/stream_writer.go
@@ -231,3 +231,10 @@ 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() error {
+	ReturnStreamWriter(sw)
+	return nil
+}
diff --git a/protocol/stream/stream.go b/protocol/stream/stream.go
index 45b5b678..ed0b9e01 100644
--- 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() error
 }
 
 // Reader defines an decoder for a Thrift value, implemented in a streaming