github.com/v2fly/v2ray-core/v4@v4.45.2/transport/internet/kcp/output.go (about)

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