mirror of https://github.com/v2ray/v2ray-core
94 lines
2.1 KiB
Go
94 lines
2.1 KiB
Go
package bufio
|
|
|
|
import (
|
|
"io"
|
|
|
|
"v2ray.com/core/common/buf"
|
|
"v2ray.com/core/common/errors"
|
|
)
|
|
|
|
// BufferedWriter is an io.Writer with internal buffer. It writes to underlying writer when buffer is full or on demand.
|
|
// This type is not thread safe.
|
|
type BufferedWriter struct {
|
|
writer io.Writer
|
|
buffer *buf.Buffer
|
|
buffered bool
|
|
}
|
|
|
|
// NewWriter creates a new BufferedWriter.
|
|
func NewWriter(rawWriter io.Writer) *BufferedWriter {
|
|
return &BufferedWriter{
|
|
writer: rawWriter,
|
|
buffer: buf.NewLocal(1024),
|
|
buffered: true,
|
|
}
|
|
}
|
|
|
|
// ReadFrom implements io.ReaderFrom.ReadFrom().
|
|
func (v *BufferedWriter) ReadFrom(reader io.Reader) (int64, error) {
|
|
totalBytes := int64(0)
|
|
for {
|
|
oriSize := v.buffer.Len()
|
|
err := v.buffer.AppendSupplier(buf.ReadFrom(reader))
|
|
totalBytes += int64(v.buffer.Len() - oriSize)
|
|
if err != nil {
|
|
if errors.Cause(err) == io.EOF {
|
|
return totalBytes, nil
|
|
}
|
|
return totalBytes, err
|
|
}
|
|
if err := v.Flush(); err != nil {
|
|
return totalBytes, err
|
|
}
|
|
}
|
|
}
|
|
|
|
func (v *BufferedWriter) Write(b []byte) (int, error) {
|
|
if !v.buffered || v.buffer == nil {
|
|
return v.writer.Write(b)
|
|
}
|
|
nBytes, err := v.buffer.Write(b)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
if v.buffer.IsFull() {
|
|
err := v.Flush()
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
if nBytes < len(b) {
|
|
if _, err := v.writer.Write(b[nBytes:]); err != nil {
|
|
return nBytes, err
|
|
}
|
|
}
|
|
}
|
|
return len(b), nil
|
|
}
|
|
|
|
// Flush writes all buffered content into underlying writer, if any.
|
|
func (v *BufferedWriter) Flush() error {
|
|
defer v.buffer.Clear()
|
|
for !v.buffer.IsEmpty() {
|
|
nBytes, err := v.writer.Write(v.buffer.Bytes())
|
|
if err != nil {
|
|
return err
|
|
}
|
|
v.buffer.SliceFrom(nBytes)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// IsBuffered returns true if this BufferedWriter holds a buffer.
|
|
func (v *BufferedWriter) IsBuffered() bool {
|
|
return v.buffered
|
|
}
|
|
|
|
// SetBuffered controls whether the BufferedWriter holds a buffer for writing. If not buffered, any write() calls into underlying writer directly.
|
|
func (v *BufferedWriter) SetBuffered(cached bool) error {
|
|
v.buffered = cached
|
|
if !cached && !v.buffer.IsEmpty() {
|
|
return v.Flush()
|
|
}
|
|
return nil
|
|
}
|