1
0
mirror of https://github.com/v2fly/v2ray-core.git synced 2025-01-17 23:06:30 -05:00

fix sync issue in buffered writer

This commit is contained in:
Darien Raymond 2018-08-31 15:02:35 +02:00
parent f1374f9caf
commit 7053293347
No known key found for this signature in database
GPG Key ID: 7251FFA14BB18169

View File

@ -3,6 +3,7 @@ package buf
import ( import (
"io" "io"
"net" "net"
"sync"
"v2ray.com/core/common" "v2ray.com/core/common"
"v2ray.com/core/common/errors" "v2ray.com/core/common/errors"
@ -56,6 +57,7 @@ func (w *BufferToBytesWriter) ReadFrom(reader io.Reader) (int64, error) {
// BufferedWriter is a Writer with internal buffer. // BufferedWriter is a Writer with internal buffer.
type BufferedWriter struct { type BufferedWriter struct {
sync.Mutex
writer Writer writer Writer
buffer *Buffer buffer *Buffer
buffered bool buffered bool
@ -77,6 +79,9 @@ func (w *BufferedWriter) WriteByte(c byte) error {
// Write implements io.Writer. // Write implements io.Writer.
func (w *BufferedWriter) Write(b []byte) (int, error) { func (w *BufferedWriter) Write(b []byte) (int, error) {
w.Lock()
defer w.Unlock()
if !w.buffered { if !w.buffered {
if writer, ok := w.writer.(io.Writer); ok { if writer, ok := w.writer.(io.Writer); ok {
return writer.Write(b) return writer.Write(b)
@ -95,7 +100,7 @@ func (w *BufferedWriter) Write(b []byte) (int, error) {
return totalBytes, err return totalBytes, err
} }
if !w.buffered || w.buffer.IsFull() { if !w.buffered || w.buffer.IsFull() {
if err := w.Flush(); err != nil { if err := w.flushInternal(); err != nil {
return totalBytes, err return totalBytes, err
} }
} }
@ -107,6 +112,9 @@ func (w *BufferedWriter) Write(b []byte) (int, error) {
// WriteMultiBuffer implements Writer. It takes ownership of the given MultiBuffer. // WriteMultiBuffer implements Writer. It takes ownership of the given MultiBuffer.
func (w *BufferedWriter) WriteMultiBuffer(b MultiBuffer) error { func (w *BufferedWriter) WriteMultiBuffer(b MultiBuffer) error {
w.Lock()
defer w.Unlock()
if !w.buffered { if !w.buffered {
return w.writer.WriteMultiBuffer(b) return w.writer.WriteMultiBuffer(b)
} }
@ -121,7 +129,7 @@ func (w *BufferedWriter) WriteMultiBuffer(b MultiBuffer) error {
return err return err
} }
if w.buffer.IsFull() { if w.buffer.IsFull() {
if err := w.Flush(); err != nil { if err := w.flushInternal(); err != nil {
return err return err
} }
} }
@ -132,6 +140,13 @@ func (w *BufferedWriter) WriteMultiBuffer(b MultiBuffer) error {
// Flush flushes buffered content into underlying writer. // Flush flushes buffered content into underlying writer.
func (w *BufferedWriter) Flush() error { func (w *BufferedWriter) Flush() error {
w.Lock()
defer w.Unlock()
return w.flushInternal()
}
func (w *BufferedWriter) flushInternal() error {
if w.buffer.IsEmpty() { if w.buffer.IsEmpty() {
return nil return nil
} }
@ -150,9 +165,12 @@ func (w *BufferedWriter) Flush() error {
// SetBuffered sets whether the internal buffer is used. If set to false, Flush() will be called to clear the buffer. // SetBuffered sets whether the internal buffer is used. If set to false, Flush() will be called to clear the buffer.
func (w *BufferedWriter) SetBuffered(f bool) error { func (w *BufferedWriter) SetBuffered(f bool) error {
w.Lock()
defer w.Unlock()
w.buffered = f w.buffered = f
if !f { if !f {
return w.Flush() return w.flushInternal()
} }
return nil return nil
} }