/
io.go
98 lines (89 loc) · 1.74 KB
/
io.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
package helper
import (
"bytes"
"io"
"net"
"time"
)
type BufferChunkReader struct {
io.Reader
Err error
}
func (r *BufferChunkReader) Read(p []byte) (int, error) {
n, err := r.Reader.Read(p)
r.Err = err
if nil != err {
return n, err
}
return n, io.EOF
}
type DebugReader struct {
io.Reader
Buf bytes.Buffer
}
func (r *DebugReader) Read(p []byte) (int, error) {
n, err := r.Reader.Read(p)
if n > 0 {
r.Buf.Write(p[0:n])
}
return n, err
}
func IsTimeoutError(err error) bool {
if err, ok := err.(net.Error); ok && err.Timeout() {
return true
}
return false
}
type TimeoutReadWriteCloser struct {
io.ReadWriteCloser
readDeadline time.Time
writeDeadline time.Time
}
func (s *TimeoutReadWriteCloser) SetReadDeadline(t time.Time) error {
s.readDeadline = t
return nil
}
func (s *TimeoutReadWriteCloser) SetWriteDeadline(t time.Time) error {
s.writeDeadline = t
return nil
}
func (s *TimeoutReadWriteCloser) Read(p []byte) (n int, err error) {
var timeout <-chan time.Time
if !s.readDeadline.IsZero() {
delay := s.readDeadline.Sub(time.Now())
timeout = time.After(delay)
} else {
return s.ReadWriteCloser.Read(p)
}
done := make(chan bool, 1)
go func() {
n, err = s.ReadWriteCloser.Read(p)
done <- true
}()
select {
case <-done:
return
case <-timeout:
return 0, ErrReadTimeout
}
}
func (s *TimeoutReadWriteCloser) Write(p []byte) (n int, err error) {
var timeout <-chan time.Time
if !s.writeDeadline.IsZero() {
delay := s.writeDeadline.Sub(time.Now())
timeout = time.After(delay)
} else {
return s.ReadWriteCloser.Write(p)
}
done := make(chan bool, 1)
go func() {
n, err = s.ReadWriteCloser.Write(p)
done <- true
}()
select {
case <-done:
return
case <-timeout:
return 0, ErrWriteTimeout
}
}