diff --git a/adapter.go b/adapter.go index 19c10c7..a44ce58 100644 --- a/adapter.go +++ b/adapter.go @@ -1,6 +1,7 @@ package httpcompression // import "github.com/CAFxX/httpcompression" import ( + "bufio" "compress/gzip" "fmt" "net/http" @@ -35,6 +36,9 @@ const ( // In general there can be no one-size-fits-all value: you will want to measure if a different // minimum size improves end-to-end performance for your workloads. DefaultMinSize = 200 + + DefaultBufferSize = 4096 + minBufferSize = 512 ) // Adapter returns a HTTP handler wrapping function (a.k.a. middleware) @@ -64,6 +68,10 @@ func Adapter(opts ...Option) (func(http.Handler) http.Handler, error) { }, nil } + if c.bufferSize < minBufferSize { + c.bufferSize = minBufferSize + } + bufPool := &sync.Pool{} writerPool := &sync.Pool{} @@ -92,6 +100,7 @@ func Adapter(opts ...Option) (func(http.Handler) http.Handler, error) { gw, _ := writerPool.Get().(*compressWriter) if gw == nil { gw = &compressWriter{} + gw.bw = *bufio.NewWriterSize(skipBuffer{gw}, c.bufferSize) } *gw = compressWriter{ ResponseWriter: w, @@ -99,6 +108,7 @@ func Adapter(opts ...Option) (func(http.Handler) http.Handler, error) { accept: accept, common: common, pool: bufPool, + bw: gw.bw, } defer func() { // Important: gw.Close() must be called *always*, as this will @@ -106,9 +116,10 @@ func Adapter(opts ...Option) (func(http.Handler) http.Handler, error) { // it is guaranteed by the CompressorProvider interface, and // because some compressors may be implemented via cgo, and they // may rely on Close() being called to release memory resources. - // TODO: expose the error - _ = gw.Close() // expose the error - *gw = compressWriter{} + _ = gw.Close() // TODO: expose the error + *gw = compressWriter{ + bw: gw.bw, + } writerPool.Put(gw) }() @@ -145,6 +156,7 @@ func DefaultAdapter(opts ...Option) (func(http.Handler) http.Handler, error) { BrotliCompressionLevel(brotli.DefaultCompression), defaultZstandardCompressor(), MinSize(DefaultMinSize), + BufferSize(DefaultBufferSize), } opts = append(defaults, opts...) return Adapter(opts...) @@ -157,6 +169,7 @@ type config struct { blacklist bool prefer PreferType compressor comps + bufferSize int } type comps map[string]comp @@ -181,6 +194,16 @@ func MinSize(size int) Option { } } +func BufferSize(size int) Option { + return func(c *config) error { + if size < 0 { + return fmt.Errorf("buffer size can not be negative: %d", size) + } + c.bufferSize = size + return nil + } +} + // DeflateCompressionLevel is an option that controls the Deflate compression // level to be used when compressing payloads. // The default is flate.DefaultCompression. diff --git a/adapter_test.go b/adapter_test.go index 1c67710..9e59efd 100644 --- a/adapter_test.go +++ b/adapter_test.go @@ -654,7 +654,6 @@ func TestGzipHandlerDoubleWriteHeader(t *testing.T) { } assert.Empty(t, body) header := rec.Header() - assert.Equal(t, "gzip", header.Get("Content-Encoding")) assert.Equal(t, "Accept-Encoding", header.Get("Vary")) assert.Equal(t, 304, rec.Code) } @@ -1015,8 +1014,8 @@ func TestWriteStringNoCompressionDynamic(t *testing.T) { t.Run("WriteString", func(t *testing.T) { w := &discardResponseWriterWithWriteString{} h.ServeHTTP(w, r) - if w.s != len(testBody) || w.b != int64(len(testBody)) { // first WriteString falls back to Write - t.Fatalf("WriteString not called: %+v", w) + if w.b != int64(len(testBody))*2 { + t.Fatalf("Write not called: %+v", w) } }) t.Run("Write", func(t *testing.T) { diff --git a/response_writer.go b/response_writer.go index fbb3f00..7fdb15e 100644 --- a/response_writer.go +++ b/response_writer.go @@ -26,6 +26,8 @@ type compressWriter struct { enc string code int // Saves the WriteHeader value. buf *[]byte // Holds the first part of the write before reaching the minSize or the end of the write. + + bw bufio.Writer } var ( @@ -53,7 +55,7 @@ var ( const maxBuf = 1 << 16 // maximum size of recycled buffer // Write compresses and appends the given byte slice to the underlying ResponseWriter. -func (w *compressWriter) Write(b []byte) (int, error) { +func (w *compressWriter) write(b []byte) (int, error) { if w.w != nil { // The responseWriter is already initialized: use it. return w.w.Write(b) @@ -126,25 +128,15 @@ func (w *compressWriter) Write(b []byte) (int, error) { return len(b), nil } -// WriteString compresses and appends the given string to the underlying ResponseWriter. -// -// This makes use of an optional method (WriteString) exposed by the compressors, or by -// the underlying ResponseWriter. -func (w *compressWriter) WriteString(s string) (int, error) { - // Since WriteString is an optional interface of the compressor, and the actual compressor - // is chosen only after the first call to Write, we can't statically know whether the interface - // is supported. We therefore have to check dynamically. - if ws, _ := w.w.(io.StringWriter); ws != nil { - // The responseWriter is already initialized and it implements WriteString. - return ws.WriteString(s) - } - // Fallback: the writer has not been initialized yet, or it has been initialized - // and it does not implement WriteString. We could in theory do something unsafe - // here but for now let's keep it simple and fallback to Write. - // TODO: in case the string is large, we should avoid allocating a full copy: - // instead we should copy the string in chunks. - return w.Write([]byte(s)) -} +type skipBuffer struct{ cw *compressWriter } + +func (w skipBuffer) Write(b []byte) (int, error) { return w.cw.write(b) } + +func (w *compressWriter) WriteString(s string) (int, error) { return w.bw.WriteString(s) } +func (w *compressWriter) Write(b []byte) (int, error) { return w.bw.Write(b) } +func (w *compressWriter) WriteRune(r rune) (int, error) { return w.bw.WriteRune(r) } +func (w *compressWriter) WriteByte(c byte) error { return w.bw.WriteByte(c) } +func (w *compressWriter) ReadFrom(r io.Reader) (int64, error) { return w.bw.ReadFrom(r) } // startCompress initializes a compressing writer and writes the buffer. func (w *compressWriter) startCompress(enc string, buf []byte) error { @@ -229,6 +221,8 @@ func (w *compressWriter) WriteHeader(code int) { // Close closes the compression Writer. func (w *compressWriter) Close() error { + w.bw.Flush() + if w.w != nil && w.enc == "" { return nil } @@ -257,12 +251,15 @@ func (w *compressWriter) Close() error { // response should be compressed or not (e.g. less than MinSize bytes have // been written). func (w *compressWriter) Flush() { - if w.w == nil { + if w.w == nil || (w.w == nil && w.bw.Buffered() < minBufferSize) { // Flush is thus a no-op until we're certain whether a plain // or compressed response will be served. return } + // Flush the bufio.Writer. + w.bw.Flush() + // Flush the compressor, if supported. // note: http.ResponseWriter does not implement Flusher (http.Flusher does not return an error), // so we need to later call ResponseWriter.Flush anyway: