1
0
mirror of https://github.com/v2fly/v2ray-core.git synced 2024-12-22 18:17:52 -05:00
v2fly/common/bufio/writer.go
Darien Raymond f159d504f2
comments
2016-12-12 02:14:48 +01:00

127 lines
2.1 KiB
Go

package bufio
import (
"io"
"sync"
"v2ray.com/core/common"
"v2ray.com/core/common/buf"
"v2ray.com/core/common/errors"
)
type BufferedWriter struct {
sync.Mutex
writer io.Writer
buffer *buf.Buffer
cached bool
}
func NewWriter(rawWriter io.Writer) *BufferedWriter {
return &BufferedWriter{
writer: rawWriter,
buffer: buf.NewSmall(),
cached: true,
}
}
func (v *BufferedWriter) ReadFrom(reader io.Reader) (int64, error) {
v.Lock()
defer v.Unlock()
if v.writer == nil {
return 0, io.ErrClosedPipe
}
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
}
v.FlushWithoutLock()
}
}
func (v *BufferedWriter) Write(b []byte) (int, error) {
v.Lock()
defer v.Unlock()
if v.writer == nil {
return 0, io.ErrClosedPipe
}
if !v.cached {
return v.writer.Write(b)
}
nBytes, err := v.buffer.Write(b)
if err != nil {
return 0, err
}
if v.buffer.IsFull() {
err := v.FlushWithoutLock()
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
}
func (v *BufferedWriter) Flush() error {
v.Lock()
defer v.Unlock()
if v.writer == nil {
return io.ErrClosedPipe
}
return v.FlushWithoutLock()
}
func (v *BufferedWriter) FlushWithoutLock() 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
}
func (v *BufferedWriter) Cached() bool {
return v.cached
}
func (v *BufferedWriter) SetCached(cached bool) {
v.cached = cached
if !cached && !v.buffer.IsEmpty() {
v.Flush()
}
}
func (v *BufferedWriter) Release() {
v.Flush()
v.Lock()
defer v.Unlock()
v.buffer.Release()
v.buffer = nil
if releasable, ok := v.writer.(common.Releasable); ok {
releasable.Release()
}
v.writer = nil
}