github.com/EagleQL/Xray-core@v1.4.3/transport/internet/kcp/output.go (about)

     1  package kcp
     2  
     3  import (
     4  	"io"
     5  	"sync"
     6  
     7  	"github.com/xtls/xray-core/common/retry"
     8  
     9  	"github.com/xtls/xray-core/common/buf"
    10  )
    11  
    12  type SegmentWriter interface {
    13  	Write(seg Segment) error
    14  }
    15  
    16  type SimpleSegmentWriter struct {
    17  	sync.Mutex
    18  	buffer *buf.Buffer
    19  	writer io.Writer
    20  }
    21  
    22  func NewSegmentWriter(writer io.Writer) SegmentWriter {
    23  	return &SimpleSegmentWriter{
    24  		writer: writer,
    25  		buffer: buf.New(),
    26  	}
    27  }
    28  
    29  func (w *SimpleSegmentWriter) Write(seg Segment) error {
    30  	w.Lock()
    31  	defer w.Unlock()
    32  
    33  	w.buffer.Clear()
    34  	rawBytes := w.buffer.Extend(seg.ByteSize())
    35  	seg.Serialize(rawBytes)
    36  	_, err := w.writer.Write(w.buffer.Bytes())
    37  	return err
    38  }
    39  
    40  type RetryableWriter struct {
    41  	writer SegmentWriter
    42  }
    43  
    44  func NewRetryableWriter(writer SegmentWriter) SegmentWriter {
    45  	return &RetryableWriter{
    46  		writer: writer,
    47  	}
    48  }
    49  
    50  func (w *RetryableWriter) Write(seg Segment) error {
    51  	return retry.Timed(5, 100).On(func() error {
    52  		return w.writer.Write(seg)
    53  	})
    54  }