From 46cab7f1310f6aed97c34049041da2634006dad7 Mon Sep 17 00:00:00 2001
From: Akshay Shah <akshay@akshayshah.org>
Date: Mon, 9 Aug 2021 22:31:33 -0700
Subject: [PATCH] Add get/put helpers for gzip writer pool

---
 gzip.go    | 12 ++++++++++++
 handler.go |  9 ++-------
 2 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/gzip.go b/gzip.go
index 1c588f84..20b26c6c 100644
--- a/gzip.go
+++ b/gzip.go
@@ -13,6 +13,18 @@ var gzWriterPool = sync.Pool{
 	},
 }
 
+func getGzipWriter(w io.Writer) *gzip.Writer {
+	gw := gzWriterPool.Get().(*gzip.Writer)
+	gw.Reset(w)
+	return gw
+}
+
+func putGzipWriter(gw *gzip.Writer) {
+	gw.Close()           // close if we haven't already
+	gw.Reset(io.Discard) // don't keep references
+	gzWriterPool.Put(gw)
+}
+
 // Verify we're implementing these interfaces at compile time.
 var (
 	_ http.ResponseWriter = &gzipResponseWriter{}
diff --git a/handler.go b/handler.go
index 2e874537..27613a56 100644
--- a/handler.go
+++ b/handler.go
@@ -318,14 +318,9 @@ func (h *Handler) writeResultTwirp(ctx context.Context, w http.ResponseWriter, s
 	// ResponseWriter.
 	if spec.ResponseCompression == CompressionGzip && w.Header().Get("Content-Encoding") == "" {
 		w.Header().Set("Content-Encoding", "gzip")
-		gw := gzWriterPool.Get().(*gzip.Writer)
-		gw.Reset(w)
+		gw := getGzipWriter(w)
+		defer putGzipWriter(gw)
 		w = &gzipResponseWriter{ResponseWriter: w, gw: gw}
-		defer func() {
-			gw.Close()           // close if we haven't already
-			gw.Reset(io.Discard) // don't keep references
-			gzWriterPool.Put(gw)
-		}()
 	}
 	if err != nil {
 		// Twirp always writes errors as JSON.