refactor: compress middleware support dynamic level

vicanso · Oct 17, 2021 · 35b5cd5 · 35b5cd5
1 parent 14e8224
commit 35b5cd5
Show file tree

Hide file tree

Showing 6 changed files with 29 additions and 6 deletions.
diff --git a/go.mod b/go.mod
@@ -4,7 +4,7 @@ go 1.16
 
 require (
 	github.com/stretchr/testify v1.7.0
-	github.com/tidwall/gjson v1.9.3
+	github.com/tidwall/gjson v1.9.4
 	github.com/vicanso/hes v0.3.9
 	github.com/vicanso/intranet-ip v0.0.1
 	github.com/vicanso/keygrip v1.2.1

diff --git a/go.sum b/go.sum
@@ -14,6 +14,8 @@ github.com/tidwall/gjson v1.9.0 h1:+Od7AE26jAaMgVC31cQV/Ope5iKXulNMflrlB7k+F9E=
 github.com/tidwall/gjson v1.9.0/go.mod h1:5/xDoumyyDNerp2U36lyolv46b3uF/9Bu6OfyQ9GImk=
 github.com/tidwall/gjson v1.9.3 h1:hqzS9wAHMO+KVBBkLxYdkEeeFHuqr95GfClRLKlgK0E=
 github.com/tidwall/gjson v1.9.3/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk=
+github.com/tidwall/gjson v1.9.4 h1:oNis7dk9Rs3dKJNNigXZT1MTOiJeBtpurn+IpCB75MY=
+github.com/tidwall/gjson v1.9.4/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk=
 github.com/tidwall/match v1.0.3 h1:FQUVvBImDutD8wJLN6c5eMzWtjgONK9MwIBCOrUJKeE=
 github.com/tidwall/match v1.0.3/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM=
 github.com/tidwall/match v1.1.1 h1:+Ho715JplO36QYgwN9PGYNhgZvoUSc9X2c80KVTi+GA=

diff --git a/middleware/compress.go b/middleware/compress.go
@@ -47,7 +47,7 @@ type (
 		// Accept accept check function
 		Accept(c *elton.Context, bodySize int) (acceptable bool, encoding string)
 		// Compress compress function
-		Compress([]byte) (*bytes.Buffer, error)
+		Compress([]byte, ...int) (*bytes.Buffer, error)
 		// Pipe pipe function
 		Pipe(*elton.Context) error
 	}
@@ -59,6 +59,8 @@ type (
 		Compressors []Compressor
 		// Skipper skipper function
 		Skipper elton.Skipper
+		// DynamicLevel return dynamic level
+		DynamicLevel func(c *elton.Context, bodySize int, encoding string) int
 	}
 )
 
@@ -110,6 +112,7 @@ func NewCompress(config CompressConfig) elton.Handler {
 	if len(compressorList) == 0 {
 		panic(errors.New("compressor can't be empty"))
 	}
+	dynamicLevel := config.DynamicLevel
 	return func(c *elton.Context) error {
 		if skipper(c) {
 			return c.Next()
@@ -178,7 +181,13 @@ func NewCompress(config CompressConfig) elton.Handler {
 				break
 			}
 
-			newBuf, e := compressor.Compress(body)
+			levels := make([]int, 0)
+			// 如果获取压缩级别函数有设置
+			if dynamicLevel != nil {
+				levels = append(levels, dynamicLevel(c, len(body), encoding))
+			}
+
+			newBuf, e := compressor.Compress(body, levels...)
 			// 如果压缩成功，则使用压缩数据
 			// 失败则忽略
 			if e != nil {

diff --git a/middleware/compress_test.go b/middleware/compress_test.go
@@ -25,6 +25,7 @@ package middleware
 import (
 	"bytes"
 	"errors"
+	"math"
 	"math/rand"
 	"net/http/httptest"
 	"testing"
@@ -42,7 +43,7 @@ func (t *testCompressor) Accept(c *elton.Context, bodySize int) (acceptable bool
 	return AcceptEncoding(c, "test")
 }
 
-func (t *testCompressor) Compress(buf []byte) (*bytes.Buffer, error) {
+func (t *testCompressor) Compress(buf []byte, levels ...int) (*bytes.Buffer, error) {
 	return bytes.NewBufferString("abcd"), nil
 }
 
@@ -99,6 +100,10 @@ func TestCompressMiddleware(t *testing.T) {
 		Compressors: []Compressor{
 			new(testCompressor),
 		},
+		DynamicLevel: func(c *elton.Context, bodySize int, encoding string) int {
+			// 不指定压缩级别
+			return math.MinInt
+		},
 	})
 
 	tests := []struct {

diff --git a/middleware/gzip.go b/middleware/gzip.go
@@ -26,6 +26,7 @@ import (
 	"bytes"
 	"compress/gzip"
 	"io"
+	"math"
 
 	"github.com/vicanso/elton"
 )
@@ -53,8 +54,11 @@ func (g *GzipCompressor) Accept(c *elton.Context, bodySize int) (bool, string) {
 }
 
 // Compress compress data by gzip
-func (g *GzipCompressor) Compress(buf []byte) (*bytes.Buffer, error) {
+func (g *GzipCompressor) Compress(buf []byte, levels ...int) (*bytes.Buffer, error) {
 	level := g.getLevel()
+	if len(levels) != 0 && levels[0] != math.MinInt {
+		level = levels[0]
+	}
 	buffer := new(bytes.Buffer)
 
 	w, err := gzip.NewWriterLevel(buffer, level)

diff --git a/middleware/gzip_test.go b/middleware/gzip_test.go
@@ -26,6 +26,7 @@ import (
 	"bytes"
 	"compress/gzip"
 	"io/ioutil"
+	"math"
 	"net/http/httptest"
 	"testing"
 
@@ -49,7 +50,9 @@ func TestGzipCompress(t *testing.T) {
 	assert.True(acceptable)
 	assert.Equal(GzipEncoding, encoding)
 
-	buf, err := g.Compress([]byte(originalData))
+	_, err := g.Compress([]byte(originalData), 1)
+	assert.Nil(err)
+	buf, err := g.Compress([]byte(originalData), math.MinInt)
 	assert.Nil(err)
 
 	r, err := gzip.NewReader(bytes.NewReader(buf.Bytes()))