1
0
mirror of https://github.com/v2fly/v2ray-core.git synced 2024-12-23 02:26:49 -05:00
v2fly/common/bufio/writer.go

94 lines
2.1 KiB
Go
Raw Normal View History

2016-12-09 07:17:34 -05:00
package bufio
2016-02-26 17:45:33 -05:00
import (
"io"
2016-12-09 06:08:25 -05:00
2016-12-09 05:35:27 -05:00
"v2ray.com/core/common/buf"
2016-12-04 03:10:47 -05:00
"v2ray.com/core/common/errors"
2016-02-26 17:45:33 -05:00
)
2016-12-27 15:41:44 -05:00
// 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.
2016-02-26 17:45:33 -05:00
type BufferedWriter struct {
2016-12-27 15:41:44 -05:00
writer io.Writer
buffer *buf.Buffer
buffered bool
2016-02-26 17:45:33 -05:00
}
2016-12-27 15:41:44 -05:00
// NewWriter creates a new BufferedWriter.
2016-12-09 07:17:34 -05:00
func NewWriter(rawWriter io.Writer) *BufferedWriter {
2016-02-26 17:45:33 -05:00
return &BufferedWriter{
2016-12-27 15:41:44 -05:00
writer: rawWriter,
2017-01-04 06:52:24 -05:00
buffer: buf.NewLocal(1024),
2016-12-27 15:41:44 -05:00
buffered: true,
2016-02-26 17:45:33 -05:00
}
}
2016-12-27 17:09:08 -05:00
// ReadFrom implements io.ReaderFrom.ReadFrom().
2016-11-27 15:39:09 -05:00
func (v *BufferedWriter) ReadFrom(reader io.Reader) (int64, error) {
2016-06-01 15:17:44 -04:00
totalBytes := int64(0)
for {
2016-12-09 06:08:25 -05:00
oriSize := v.buffer.Len()
err := v.buffer.AppendSupplier(buf.ReadFrom(reader))
totalBytes += int64(v.buffer.Len() - oriSize)
2016-06-01 15:17:44 -04:00
if err != nil {
2016-12-04 03:10:47 -05:00
if errors.Cause(err) == io.EOF {
2016-06-01 15:17:44 -04:00
return totalBytes, nil
}
return totalBytes, err
}
if err := v.Flush(); err != nil {
return totalBytes, err
}
2016-06-01 15:17:44 -04:00
}
}
2016-11-27 15:39:09 -05:00
func (v *BufferedWriter) Write(b []byte) (int, error) {
2016-12-27 15:41:44 -05:00
if !v.buffered || v.buffer == nil {
2016-11-27 15:39:09 -05:00
return v.writer.Write(b)
2016-02-26 17:45:33 -05:00
}
2016-12-04 15:34:22 -05:00
nBytes, err := v.buffer.Write(b)
if err != nil {
return 0, err
}
2016-11-27 15:39:09 -05:00
if v.buffer.IsFull() {
err := v.Flush()
2016-12-04 15:34:22 -05:00
if err != nil {
return 0, err
}
if nBytes < len(b) {
if _, err := v.writer.Write(b[nBytes:]); err != nil {
return nBytes, err
}
}
2016-02-26 17:45:33 -05:00
}
2016-12-04 15:34:22 -05:00
return len(b), nil
2016-02-26 17:45:33 -05:00
}
2016-12-27 17:09:08 -05:00
// Flush writes all buffered content into underlying writer, if any.
2016-11-27 15:39:09 -05:00
func (v *BufferedWriter) Flush() error {
defer v.buffer.Clear()
for !v.buffer.IsEmpty() {
2016-12-05 09:19:14 -05:00
nBytes, err := v.writer.Write(v.buffer.Bytes())
2016-05-12 20:20:07 -04:00
if err != nil {
return err
}
2016-11-27 15:39:09 -05:00
v.buffer.SliceFrom(nBytes)
2016-02-26 17:45:33 -05:00
}
2016-05-12 20:20:07 -04:00
return nil
2016-02-26 17:45:33 -05:00
}
2016-12-27 17:09:08 -05:00
// IsBuffered returns true if this BufferedWriter holds a buffer.
func (v *BufferedWriter) IsBuffered() bool {
2016-12-27 15:41:44 -05:00
return v.buffered
2016-02-26 17:45:33 -05:00
}
2016-12-27 17:09:08 -05:00
// SetBuffered controls whether the BufferedWriter holds a buffer for writing. If not buffered, any write() calls into underlying writer directly.
2017-01-06 05:40:35 -05:00
func (v *BufferedWriter) SetBuffered(cached bool) error {
2016-12-27 15:41:44 -05:00
v.buffered = cached
2016-11-27 15:39:09 -05:00
if !cached && !v.buffer.IsEmpty() {
2017-01-06 05:40:35 -05:00
return v.Flush()
2016-02-26 17:45:33 -05:00
}
2017-01-06 05:40:35 -05:00
return nil
2016-02-26 17:45:33 -05:00
}