github.com/sagernet/sing@v0.4.0-beta.19.0.20240518125136-f67a0988a636/common/bufio/buffer.go (about)

     1  package bufio
     2  
     3  import (
     4  	"io"
     5  	"sync"
     6  
     7  	"github.com/sagernet/sing/common/buf"
     8  )
     9  
    10  type BufferedWriter struct {
    11  	upstream io.Writer
    12  	buffer   *buf.Buffer
    13  	access   sync.Mutex
    14  }
    15  
    16  func NewBufferedWriter(upstream io.Writer, buffer *buf.Buffer) *BufferedWriter {
    17  	return &BufferedWriter{
    18  		upstream: upstream,
    19  		buffer:   buffer,
    20  	}
    21  }
    22  
    23  func (w *BufferedWriter) Write(p []byte) (n int, err error) {
    24  	w.access.Lock()
    25  	defer w.access.Unlock()
    26  	if w.buffer == nil {
    27  		return w.upstream.Write(p)
    28  	}
    29  	for {
    30  		var writeN int
    31  		writeN, err = w.buffer.Write(p[n:])
    32  		n += writeN
    33  		if n == len(p) {
    34  			return
    35  		}
    36  		_, err = w.upstream.Write(w.buffer.Bytes())
    37  		if err != nil {
    38  			return
    39  		}
    40  		w.buffer.Reset()
    41  	}
    42  }
    43  
    44  func (w *BufferedWriter) Fallthrough() error {
    45  	w.access.Lock()
    46  	defer w.access.Unlock()
    47  	if w.buffer == nil {
    48  		return nil
    49  	}
    50  	if !w.buffer.IsEmpty() {
    51  		_, err := w.upstream.Write(w.buffer.Bytes())
    52  		if err != nil {
    53  			return err
    54  		}
    55  	}
    56  	w.buffer.Release()
    57  	w.buffer = nil
    58  	return nil
    59  }
    60  
    61  func (w *BufferedWriter) WriterReplaceable() bool {
    62  	return w.buffer == nil
    63  }