v2ray-core/common/io/buffered_writer.go

109 lines
1.8 KiB
Go

package io
import (
"io"
"sync"
"github.com/v2ray/v2ray-core/common/alloc"
)
type BufferedWriter struct {
sync.Mutex
writer io.Writer
buffer *alloc.Buffer
cached bool
}
func NewBufferedWriter(rawWriter io.Writer) *BufferedWriter {
return &BufferedWriter{
writer: rawWriter,
buffer: alloc.NewBuffer().Clear(),
cached: true,
}
}
func (this *BufferedWriter) ReadFrom(reader io.Reader) (int64, error) {
this.Lock()
defer this.Unlock()
if this.writer == nil {
return 0, io.EOF
}
totalBytes := int64(0)
for {
nBytes, err := this.buffer.FillFrom(reader)
totalBytes += int64(nBytes)
if err != nil {
if err == io.EOF {
return totalBytes, nil
}
return totalBytes, err
}
this.FlushWithoutLock()
}
}
func (this *BufferedWriter) Write(b []byte) (int, error) {
this.Lock()
defer this.Unlock()
if this.writer == nil {
return 0, io.EOF
}
if !this.cached {
return this.writer.Write(b)
}
nBytes, _ := this.buffer.Write(b)
if this.buffer.IsFull() {
this.FlushWithoutLock()
}
return nBytes, nil
}
func (this *BufferedWriter) Flush() error {
this.Lock()
defer this.Unlock()
if this.writer == nil {
return io.EOF
}
return this.FlushWithoutLock()
}
func (this *BufferedWriter) FlushWithoutLock() error {
defer this.buffer.Clear()
for !this.buffer.IsEmpty() {
nBytes, err := this.writer.Write(this.buffer.Value)
if err != nil {
return err
}
this.buffer.SliceFrom(nBytes)
}
return nil
}
func (this *BufferedWriter) Cached() bool {
return this.cached
}
func (this *BufferedWriter) SetCached(cached bool) {
this.cached = cached
if !cached && !this.buffer.IsEmpty() {
this.Flush()
}
}
func (this *BufferedWriter) Release() {
this.Flush()
this.Lock()
defer this.Unlock()
this.buffer.Release()
this.buffer = nil
this.writer = nil
}