Use sync.Pool to reuse gzip.WriterHEADmaster

1 files changed, 35 insertions, 16 deletions

diff --git a/handler.go b/handler.go
index ab21180..06bade4 100644
--- a/handler.go
+++ b/handler.go
@@ -2,23 +2,30 @@ package gziphandler
 
 import (
 	"compress/gzip"
+	"io/ioutil"
 	"net/http"
 	"strings"
+	"sync"
 )
 
-type gw struct {
+type respWriter struct {
 	http.ResponseWriter
 	w *gzip.Writer
 }
 
-func (w *gw) Write(d []byte) (int, error) {
+func (w *respWriter) Write(d []byte) (int, error) {
 	return w.w.Write(d)
 }
 
-func (w *gw) Close() error {
+func (w *respWriter) Close() error {
 	return w.w.Close()
 }
 
+type handler struct {
+	next http.Handler
+	pool sync.Pool
+}
+
 func acceptsGzip(r *http.Request) bool {
 	ae := r.Header.Get("Accept-Encoding")
 	for _, e := range strings.Split(ae, ",") {
@@ -33,21 +40,33 @@ func acceptsGzip(r *http.Request) bool {
 	return false
 }
 
+func (h *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	if !acceptsGzip(r) {
+		h.next.ServeHTTP(w, r)
+		return
+	}
+	w.Header().Set("Content-Encoding", "gzip")
+	gw := h.pool.Get().(*gzip.Writer)
+	gw.Reset(w)
+	rw := &respWriter{
+		ResponseWriter: w,
+		w:              gw,
+	}
+	h.next.ServeHTTP(rw, r)
+	gw.Close()
+	h.pool.Put(gw)
+}
+
 // Handler returns an http.Handler that compresses the response data written
 // by an existing handler h, using the compress/gzip.Writer with default
 // compression level.
 func Handler(h http.Handler) http.Handler {
-	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		if !acceptsGzip(r) {
-			h.ServeHTTP(w, r)
-			return
-		}
-		w.Header().Set("Content-Encoding", "gzip")
-		gw := &gw{
-			ResponseWriter: w,
-			w:              gzip.NewWriter(w),
-		}
-		h.ServeHTTP(gw, r)
-		gw.Close()
-	})
+	return &handler{
+		next: h,
+		pool: sync.Pool{
+			New: func() interface{} {
+				return gzip.NewWriter(ioutil.Discard)
+			},
+		},
+	}
 }