From ac54e132e409e86ff3a6ec61f7ae3a04a20a8243 Mon Sep 17 00:00:00 2001
From: Pablo Andres Fuente <pablo.andres.fuente@gmail.com>
Date: Sat, 7 Nov 2020 03:52:35 +0000
Subject: [PATCH] Adding sync.Pool to Compress Middleware

Adding a sync.Pool for the *gzip.Writer reduces the allocations of the
Compress middleware in 50% and gives an increase on execution speed of
a 85%
This fix #1643
---
 middleware/compress.go      | 26 +++++++++++++++++++++----
 middleware/compress_test.go | 38 +++++++++++++++++++++++++++++++++++++
 2 files changed, 60 insertions(+), 4 deletions(-)

diff --git a/middleware/compress.go b/middleware/compress.go
index dd97d983..e4f9fc51 100644
--- a/middleware/compress.go
+++ b/middleware/compress.go
@@ -8,6 +8,7 @@ import (
 	"net"
 	"net/http"
 	"strings"
+	"sync"
 
 	"github.com/labstack/echo/v4"
 )
@@ -58,6 +59,8 @@ func GzipWithConfig(config GzipConfig) echo.MiddlewareFunc {
 		config.Level = DefaultGzipConfig.Level
 	}
 
+	pool := gzipPool(config)
+
 	return func(next echo.HandlerFunc) echo.HandlerFunc {
 		return func(c echo.Context) error {
 			if config.Skipper(c) {
@@ -68,11 +71,13 @@ func GzipWithConfig(config GzipConfig) echo.MiddlewareFunc {
 			res.Header().Add(echo.HeaderVary, echo.HeaderAcceptEncoding)
 			if strings.Contains(c.Request().Header.Get(echo.HeaderAcceptEncoding), gzipScheme) {
 				res.Header().Set(echo.HeaderContentEncoding, gzipScheme) // Issue #806
-				rw := res.Writer
-				w, err := gzip.NewWriterLevel(rw, config.Level)
-				if err != nil {
-					return err
+				i := pool.Get()
+				w, ok := i.(*gzip.Writer)
+				if !ok {
+					return echo.NewHTTPError(http.StatusInternalServerError, i.(error).Error())
 				}
+				rw := res.Writer
+				w.Reset(rw)
 				defer func() {
 					if res.Size == 0 {
 						if res.Header().Get(echo.HeaderContentEncoding) == gzipScheme {
@@ -85,6 +90,7 @@ func GzipWithConfig(config GzipConfig) echo.MiddlewareFunc {
 						w.Reset(ioutil.Discard)
 					}
 					w.Close()
+					pool.Put(w)
 				}()
 				grw := &gzipResponseWriter{Writer: w, ResponseWriter: rw}
 				res.Writer = grw
@@ -126,3 +132,15 @@ func (w *gzipResponseWriter) Push(target string, opts *http.PushOptions) error {
 	}
 	return http.ErrNotSupported
 }
+
+func gzipPool(config GzipConfig) sync.Pool {
+	return sync.Pool{
+		New: func() interface{} {
+			w, err := gzip.NewWriterLevel(ioutil.Discard, config.Level)
+			if err != nil {
+				return err
+			}
+			return w
+		},
+	}
+}
diff --git a/middleware/compress_test.go b/middleware/compress_test.go
index ac5b6c3b..d16ffca4 100644
--- a/middleware/compress_test.go
+++ b/middleware/compress_test.go
@@ -120,6 +120,22 @@ func TestGzipErrorReturned(t *testing.T) {
 	assert.Empty(t, rec.Header().Get(echo.HeaderContentEncoding))
 }
 
+func TestGzipErrorReturnedInvalidConfig(t *testing.T) {
+	e := echo.New()
+	// Invalid level
+	e.Use(GzipWithConfig(GzipConfig{Level: 12}))
+	e.GET("/", func(c echo.Context) error {
+		c.Response().Write([]byte("test"))
+		return nil
+	})
+	req := httptest.NewRequest(http.MethodGet, "/", nil)
+	req.Header.Set(echo.HeaderAcceptEncoding, gzipScheme)
+	rec := httptest.NewRecorder()
+	e.ServeHTTP(rec, req)
+	assert.Equal(t, http.StatusInternalServerError, rec.Code)
+	assert.Contains(t, rec.Body.String(), "gzip")
+}
+
 // Issue #806
 func TestGzipWithStatic(t *testing.T) {
 	e := echo.New()
@@ -146,3 +162,25 @@ func TestGzipWithStatic(t *testing.T) {
 		}
 	}
 }
+
+func BenchmarkGzip(b *testing.B) {
+	e := echo.New()
+
+	req := httptest.NewRequest(http.MethodGet, "/", nil)
+	req.Header.Set(echo.HeaderAcceptEncoding, gzipScheme)
+
+	h := Gzip()(func(c echo.Context) error {
+		c.Response().Write([]byte("test")) // For Content-Type sniffing
+		return nil
+	})
+
+	b.ReportAllocs()
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		// Gzip
+		rec := httptest.NewRecorder()
+		c := e.NewContext(req, rec)
+		h(c)
+	}
+}