diff options
author | Erik Dubbelboer <erik@dubbelboer.com> | 2020-05-15 15:36:26 +0200 |
---|---|---|
committer | Erik Dubbelboer <erik@dubbelboer.com> | 2020-05-15 15:36:26 +0200 |
commit | 339ad36634b68ab5b4d90fb9991e38d5020a1039 (patch) | |
tree | 167e613ce9b8eef1d3b92c18b00d4039cbc8df91 /brotli.go | |
parent | Don't wrap conn with a TLS Client if it's already a TLS Conn (diff) | |
download | fasthttp-339ad36634b68ab5b4d90fb9991e38d5020a1039.tar.gz fasthttp-339ad36634b68ab5b4d90fb9991e38d5020a1039.tar.bz2 fasthttp-339ad36634b68ab5b4d90fb9991e38d5020a1039.zip |
Add Brotli support
New Functions:
CompressHandlerBrotliLevel(h RequestHandler, brotliLevel, otherLevel int) RequestHandler
Request.BodyUnbrotli() ([]byte, error)
Response.BodyUnbrotli() ([]byte, error)
AppendBrotliBytesLevel(dst, src []byte, level int) []byte
WriteBrotliLevel(w io.Writer, p []byte, level int) (int, error)
WriteBrotli(w io.Writer, p []byte) (int, error)
AppendBrotliBytes(dst, src []byte) []byte
WriteUnbrotli(w io.Writer, p []byte) (int, error)
AppendUnbrotliBytes(dst, src []byte) ([]byte, error)
New Constants:
CompressBrotliNoCompression
CompressBrotliBestSpeed
CompressBrotliBestCompression
CompressBrotliDefaultCompression
Brotli compression levels are different from gzip/flate. Because of this we have separate level constants and CompressHandlerBrotliLevel takes 2 levels.
I didn't add Brotli support to CompressHandler as this could cause a spike in CPU usage when users upgrade fasthttp.
fasthttp.CompressBrotliDefaultCompression is not the same as
brotli.DefaultCompression. brotli.DefaultCompression is more than twice
as slow as fasthttp.CompressBrotliDefaultCompression which I thought was
unreasonable as default.
Diffstat (limited to 'brotli.go')
-rw-r--r-- | brotli.go | 193 |
1 files changed, 193 insertions, 0 deletions
diff --git a/brotli.go b/brotli.go new file mode 100644 index 0000000..a88fdce --- /dev/null +++ b/brotli.go @@ -0,0 +1,193 @@ +package fasthttp + +import ( + "bytes" + "fmt" + "io" + "sync" + + "github.com/andybalholm/brotli" + "github.com/valyala/bytebufferpool" + "github.com/valyala/fasthttp/stackless" +) + +// Supported compression levels. +const ( + CompressBrotliNoCompression = 0 + CompressBrotliBestSpeed = brotli.BestSpeed + CompressBrotliBestCompression = brotli.BestCompression + + // Choose a default brotli compression level comparable to + // CompressDefaultCompression (gzip 6) + // See: https://github.com/valyala/fasthttp/issues/798#issuecomment-626293806 + CompressBrotliDefaultCompression = 4 +) + +func acquireBrotliReader(r io.Reader) (*brotli.Reader, error) { + v := brotliReaderPool.Get() + if v == nil { + return brotli.NewReader(r), nil + } + zr := v.(*brotli.Reader) + if err := zr.Reset(r); err != nil { + return nil, err + } + return zr, nil +} + +func releaseBrotliReader(zr *brotli.Reader) { + brotliReaderPool.Put(zr) +} + +var brotliReaderPool sync.Pool + +func acquireStacklessBrotliWriter(w io.Writer, level int) stackless.Writer { + nLevel := normalizeBrotliCompressLevel(level) + p := stacklessBrotliWriterPoolMap[nLevel] + v := p.Get() + if v == nil { + return stackless.NewWriter(w, func(w io.Writer) stackless.Writer { + return acquireRealBrotliWriter(w, level) + }) + } + sw := v.(stackless.Writer) + sw.Reset(w) + return sw +} + +func releaseStacklessBrotliWriter(sw stackless.Writer, level int) { + sw.Close() + nLevel := normalizeBrotliCompressLevel(level) + p := stacklessBrotliWriterPoolMap[nLevel] + p.Put(sw) +} + +func acquireRealBrotliWriter(w io.Writer, level int) *brotli.Writer { + nLevel := normalizeBrotliCompressLevel(level) + p := realBrotliWriterPoolMap[nLevel] + v := p.Get() + if v == nil { + zw := brotli.NewWriterLevel(w, level) + return zw + } + zw := v.(*brotli.Writer) + zw.Reset(w) + return zw +} + +func releaseRealBrotliWriter(zw *brotli.Writer, level int) { + zw.Close() + nLevel := normalizeBrotliCompressLevel(level) + p := realBrotliWriterPoolMap[nLevel] + p.Put(zw) +} + +var ( + stacklessBrotliWriterPoolMap = newCompressWriterPoolMap() + realBrotliWriterPoolMap = newCompressWriterPoolMap() +) + +// AppendBrotliBytesLevel appends brotlied src to dst using the given +// compression level and returns the resulting dst. +// +// Supported compression levels are: +// +// * CompressBrotliNoCompression +// * CompressBrotliBestSpeed +// * CompressBrotliBestCompression +// * CompressBrotliDefaultCompression +func AppendBrotliBytesLevel(dst, src []byte, level int) []byte { + w := &byteSliceWriter{dst} + WriteBrotliLevel(w, src, level) //nolint:errcheck + return w.b +} + +// WriteBrotliLevel writes brotlied p to w using the given compression level +// and returns the number of compressed bytes written to w. +// +// Supported compression levels are: +// +// * CompressBrotliNoCompression +// * CompressBrotliBestSpeed +// * CompressBrotliBestCompression +// * CompressBrotliDefaultCompression +func WriteBrotliLevel(w io.Writer, p []byte, level int) (int, error) { + switch w.(type) { + case *byteSliceWriter, + *bytes.Buffer, + *bytebufferpool.ByteBuffer: + // These writers don't block, so we can just use stacklessWriteBrotli + ctx := &compressCtx{ + w: w, + p: p, + level: level, + } + stacklessWriteBrotli(ctx) + return len(p), nil + default: + zw := acquireStacklessBrotliWriter(w, level) + n, err := zw.Write(p) + releaseStacklessBrotliWriter(zw, level) + return n, err + } +} + +var stacklessWriteBrotli = stackless.NewFunc(nonblockingWriteBrotli) + +func nonblockingWriteBrotli(ctxv interface{}) { + ctx := ctxv.(*compressCtx) + zw := acquireRealBrotliWriter(ctx.w, ctx.level) + + _, err := zw.Write(ctx.p) + if err != nil { + panic(fmt.Sprintf("BUG: brotli.Writer.Write for len(p)=%d returned unexpected error: %s", len(ctx.p), err)) + } + + releaseRealBrotliWriter(zw, ctx.level) +} + +// WriteBrotli writes brotlied p to w and returns the number of compressed +// bytes written to w. +func WriteBrotli(w io.Writer, p []byte) (int, error) { + return WriteBrotliLevel(w, p, CompressBrotliDefaultCompression) +} + +// AppendBrotliBytes appends brotlied src to dst and returns the resulting dst. +func AppendBrotliBytes(dst, src []byte) []byte { + return AppendBrotliBytesLevel(dst, src, CompressBrotliDefaultCompression) +} + +// WriteUnbrotli writes unbrotlied p to w and returns the number of uncompressed +// bytes written to w. +func WriteUnbrotli(w io.Writer, p []byte) (int, error) { + r := &byteSliceReader{p} + zr, err := acquireBrotliReader(r) + if err != nil { + return 0, err + } + n, err := copyZeroAlloc(w, zr) + releaseBrotliReader(zr) + nn := int(n) + if int64(nn) != n { + return 0, fmt.Errorf("too much data unbrotlied: %d", n) + } + return nn, err +} + +// AppendUnbrotliBytes appends unbrotlied src to dst and returns the resulting dst. +func AppendUnbrotliBytes(dst, src []byte) ([]byte, error) { + w := &byteSliceWriter{dst} + _, err := WriteUnbrotli(w, src) + return w.b, err +} + +// normalizes compression level into [0..11], so it could be used as an index +// in *PoolMap. +func normalizeBrotliCompressLevel(level int) int { + // -2 is the lowest compression level - CompressHuffmanOnly + // 9 is the highest compression level - CompressBestCompression + if level < 0 || level > 11 { + level = CompressBrotliDefaultCompression + } + return level +} |