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