1package kcp 2 3import ( 4 "io" 5 "sync" 6 7 "github.com/xtls/xray-core/common/retry" 8 9 "github.com/xtls/xray-core/common/buf" 10) 11 12type SegmentWriter interface { 13 Write(seg Segment) error 14} 15 16type SimpleSegmentWriter struct { 17 sync.Mutex 18 buffer *buf.Buffer 19 writer io.Writer 20} 21 22func NewSegmentWriter(writer io.Writer) SegmentWriter { 23 return &SimpleSegmentWriter{ 24 writer: writer, 25 buffer: buf.New(), 26 } 27} 28 29func (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 40type RetryableWriter struct { 41 writer SegmentWriter 42} 43 44func NewRetryableWriter(writer SegmentWriter) SegmentWriter { 45 return &RetryableWriter{ 46 writer: writer, 47 } 48} 49 50func (w *RetryableWriter) Write(seg Segment) error { 51 return retry.Timed(5, 100).On(func() error { 52 return w.writer.Write(seg) 53 }) 54} 55