well-goknown/vendor/github.com/valyala/fasthttp/stackless/writer.go

150 lines
2.6 KiB
Go
Raw Normal View History

package stackless
import (
"errors"
"fmt"
"io"
"sync"
"github.com/valyala/bytebufferpool"
)
// Writer is an interface stackless writer must conform to.
//
// The interface contains common subset for Writers from compress/* packages.
type Writer interface {
Write(p []byte) (int, error)
Flush() error
Close() error
Reset(w io.Writer)
}
// NewWriterFunc must return new writer that will be wrapped into
// stackless writer.
type NewWriterFunc func(w io.Writer) Writer
// NewWriter creates a stackless writer around a writer returned
// from newWriter.
//
// The returned writer writes data to dstW.
//
// Writers that use a lot of stack space may be wrapped into stackless writer,
// thus saving stack space for high number of concurrently running goroutines.
func NewWriter(dstW io.Writer, newWriter NewWriterFunc) Writer {
w := &writer{
dstW: dstW,
}
w.zw = newWriter(&w.xw)
return w
}
type writer struct {
dstW io.Writer
zw Writer
xw xWriter
err error
n int
p []byte
op op
}
type op int
const (
opWrite op = iota
opFlush
opClose
opReset
)
func (w *writer) Write(p []byte) (int, error) {
w.p = p
err := w.do(opWrite)
w.p = nil
return w.n, err
}
func (w *writer) Flush() error {
return w.do(opFlush)
}
func (w *writer) Close() error {
return w.do(opClose)
}
func (w *writer) Reset(dstW io.Writer) {
w.xw.Reset()
w.do(opReset) //nolint:errcheck
w.dstW = dstW
}
func (w *writer) do(op op) error {
w.op = op
if !stacklessWriterFunc(w) {
return errHighLoad
}
err := w.err
if err != nil {
return err
}
if w.xw.bb != nil && len(w.xw.bb.B) > 0 {
_, err = w.dstW.Write(w.xw.bb.B)
}
w.xw.Reset()
return err
}
var errHighLoad = errors.New("cannot compress data due to high load")
var (
stacklessWriterFuncOnce sync.Once
stacklessWriterFuncFunc func(ctx interface{}) bool
)
func stacklessWriterFunc(ctx interface{}) bool {
stacklessWriterFuncOnce.Do(func() {
stacklessWriterFuncFunc = NewFunc(writerFunc)
})
return stacklessWriterFuncFunc(ctx)
}
func writerFunc(ctx interface{}) {
w := ctx.(*writer)
switch w.op {
case opWrite:
w.n, w.err = w.zw.Write(w.p)
case opFlush:
w.err = w.zw.Flush()
case opClose:
w.err = w.zw.Close()
case opReset:
w.zw.Reset(&w.xw)
w.err = nil
default:
panic(fmt.Sprintf("BUG: unexpected op: %d", w.op))
}
}
type xWriter struct {
bb *bytebufferpool.ByteBuffer
}
func (w *xWriter) Write(p []byte) (int, error) {
if w.bb == nil {
w.bb = bufferPool.Get()
}
return w.bb.Write(p)
}
func (w *xWriter) Reset() {
if w.bb != nil {
bufferPool.Put(w.bb)
w.bb = nil
}
}
var bufferPool bytebufferpool.Pool