github.com/Uhtred009/v2ray-core-1@v4.31.2+incompatible/transport/internet/kcp/output.go (about)

     1  // +build !confonly
     2  
     3  package kcp
     4  
     5  import (
     6  	"io"
     7  	"sync"
     8  
     9  	"v2ray.com/core/common/retry"
    10  
    11  	"v2ray.com/core/common/buf"
    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  }