mirror of https://github.com/v2ray/v2ray-core
73 lines
1.7 KiB
Go
73 lines
1.7 KiB
Go
package buf
|
|
|
|
import "io"
|
|
|
|
// 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 *Buffer
|
|
buffered bool
|
|
}
|
|
|
|
// NewBufferedWriter creates a new BufferedWriter.
|
|
func NewBufferedWriter(writer io.Writer) *BufferedWriter {
|
|
return NewBufferedWriterSize(writer, 1024)
|
|
}
|
|
|
|
func NewBufferedWriterSize(writer io.Writer, size uint32) *BufferedWriter {
|
|
return &BufferedWriter{
|
|
writer: writer,
|
|
buffer: NewLocal(int(size)),
|
|
buffered: true,
|
|
}
|
|
}
|
|
|
|
// Write implements io.Writer.
|
|
func (w *BufferedWriter) Write(b []byte) (int, error) {
|
|
if !w.buffered || w.buffer == nil {
|
|
return w.writer.Write(b)
|
|
}
|
|
bytesWritten := 0
|
|
for bytesWritten < len(b) {
|
|
nBytes, err := w.buffer.Write(b[bytesWritten:])
|
|
if err != nil {
|
|
return bytesWritten, err
|
|
}
|
|
bytesWritten += nBytes
|
|
if w.buffer.IsFull() {
|
|
if err := w.Flush(); err != nil {
|
|
return bytesWritten, err
|
|
}
|
|
}
|
|
}
|
|
return bytesWritten, nil
|
|
}
|
|
|
|
// Flush writes all buffered content into underlying writer, if any.
|
|
func (w *BufferedWriter) Flush() error {
|
|
defer w.buffer.Clear()
|
|
for !w.buffer.IsEmpty() {
|
|
nBytes, err := w.writer.Write(w.buffer.Bytes())
|
|
if err != nil {
|
|
return err
|
|
}
|
|
w.buffer.SliceFrom(nBytes)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// IsBuffered returns true if this BufferedWriter holds a buffer.
|
|
func (w *BufferedWriter) IsBuffered() bool {
|
|
return w.buffered
|
|
}
|
|
|
|
// SetBuffered controls whether the BufferedWriter holds a buffer for writing. If not buffered, any write() calls into underlying writer directly.
|
|
func (w *BufferedWriter) SetBuffered(cached bool) error {
|
|
w.buffered = cached
|
|
if !cached && !w.buffer.IsEmpty() {
|
|
return w.Flush()
|
|
}
|
|
return nil
|
|
}
|