v2ray-core/common/bufio/writer.go

103 lines
2.0 KiB
Go
Raw Normal View History

2016-12-09 13:17:34 +01:00
package bufio
2016-02-26 23:45:33 +01:00
import (
"io"
2016-12-09 12:08:25 +01:00
2016-12-12 02:14:48 +01:00
"v2ray.com/core/common"
2016-12-09 11:35:27 +01:00
"v2ray.com/core/common/buf"
2016-12-04 09:10:47 +01:00
"v2ray.com/core/common/errors"
2016-02-26 23:45:33 +01:00
)
2016-12-27 21:41:44 +01: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 23:45:33 +01:00
type BufferedWriter struct {
2016-12-27 21:41:44 +01:00
writer io.Writer
buffer *buf.Buffer
buffered bool
2016-02-26 23:45:33 +01:00
}
2016-12-27 21:41:44 +01:00
// NewWriter creates a new BufferedWriter.
2016-12-09 13:17:34 +01:00
func NewWriter(rawWriter io.Writer) *BufferedWriter {
2016-02-26 23:45:33 +01:00
return &BufferedWriter{
2016-12-27 21:41:44 +01:00
writer: rawWriter,
buffer: buf.NewSmall(),
buffered: true,
2016-02-26 23:45:33 +01:00
}
}
2016-11-27 21:39:09 +01:00
func (v *BufferedWriter) ReadFrom(reader io.Reader) (int64, error) {
2016-06-01 21:17:44 +02:00
totalBytes := int64(0)
for {
2016-12-09 12:08:25 +01:00
oriSize := v.buffer.Len()
err := v.buffer.AppendSupplier(buf.ReadFrom(reader))
totalBytes += int64(v.buffer.Len() - oriSize)
2016-06-01 21:17:44 +02:00
if err != nil {
2016-12-04 09:10:47 +01:00
if errors.Cause(err) == io.EOF {
2016-06-01 21:17:44 +02:00
return totalBytes, nil
}
return totalBytes, err
}
if err := v.Flush(); err != nil {
return totalBytes, err
}
2016-06-01 21:17:44 +02:00
}
}
2016-11-27 21:39:09 +01:00
func (v *BufferedWriter) Write(b []byte) (int, error) {
2016-12-27 21:41:44 +01:00
if !v.buffered || v.buffer == nil {
2016-11-27 21:39:09 +01:00
return v.writer.Write(b)
2016-02-26 23:45:33 +01:00
}
2016-12-04 21:34:22 +01:00
nBytes, err := v.buffer.Write(b)
if err != nil {
return 0, err
}
2016-11-27 21:39:09 +01:00
if v.buffer.IsFull() {
err := v.Flush()
2016-12-04 21:34:22 +01: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 23:45:33 +01:00
}
2016-12-04 21:34:22 +01:00
return len(b), nil
2016-02-26 23:45:33 +01:00
}
2016-11-27 21:39:09 +01:00
func (v *BufferedWriter) Flush() error {
defer v.buffer.Clear()
for !v.buffer.IsEmpty() {
2016-12-05 15:19:14 +01:00
nBytes, err := v.writer.Write(v.buffer.Bytes())
2016-05-12 17:20:07 -07:00
if err != nil {
return err
}
2016-11-27 21:39:09 +01:00
v.buffer.SliceFrom(nBytes)
2016-02-26 23:45:33 +01:00
}
2016-05-12 17:20:07 -07:00
return nil
2016-02-26 23:45:33 +01:00
}
2016-12-27 21:41:44 +01:00
func (v *BufferedWriter) Buffered() bool {
return v.buffered
2016-02-26 23:45:33 +01:00
}
2016-12-27 21:41:44 +01:00
func (v *BufferedWriter) SetBuffered(cached bool) {
v.buffered = cached
2016-11-27 21:39:09 +01:00
if !cached && !v.buffer.IsEmpty() {
v.Flush()
2016-02-26 23:45:33 +01:00
}
}
2016-12-27 21:15:32 +01:00
// Release implements common.Releasable.Release().
2016-11-27 21:39:09 +01:00
func (v *BufferedWriter) Release() {
2016-12-27 21:15:32 +01:00
if !v.buffer.IsEmpty() {
v.Flush()
}
2016-12-12 02:14:48 +01:00
2016-12-27 21:15:32 +01:00
if v.buffer != nil {
v.buffer.Release()
v.buffer = nil
}
2016-12-27 21:06:55 +01:00
common.Release(v.writer)
2016-02-26 23:45:33 +01:00
}