github.com/v2fly/v2ray-core/v5@v5.16.2-0.20240507031116-8191faa6e095/transport/internet/kcp/output.go (about)

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