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
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
|
package drive
import (
"golang.org/x/net/context"
"io"
"sync"
"time"
)
const TimeoutTimerInterval = time.Second * 10
type timeoutReaderWrapper func(io.Reader) io.Reader
func getTimeoutReaderWrapperContext(timeout time.Duration) (timeoutReaderWrapper, context.Context) {
ctx, cancel := context.WithCancel(context.TODO())
wrapper := func(r io.Reader) io.Reader {
// Return untouched reader if timeout is 0
if timeout == 0 {
return r
}
return getTimeoutReader(r, cancel, timeout)
}
return wrapper, ctx
}
func getTimeoutReaderContext(r io.Reader, timeout time.Duration) (io.Reader, context.Context) {
ctx, cancel := context.WithCancel(context.TODO())
// Return untouched reader if timeout is 0
if timeout == 0 {
return r, ctx
}
return getTimeoutReader(r, cancel, timeout), ctx
}
func getTimeoutReader(r io.Reader, cancel context.CancelFunc, timeout time.Duration) io.Reader {
return &TimeoutReader{
reader: r,
cancel: cancel,
mutex: &sync.Mutex{},
maxIdleTimeout: timeout,
}
}
type TimeoutReader struct {
reader io.Reader
cancel context.CancelFunc
lastActivity time.Time
timer *time.Timer
mutex *sync.Mutex
maxIdleTimeout time.Duration
done bool
}
func (self *TimeoutReader) Read(p []byte) (int, error) {
if self.timer == nil {
self.startTimer()
}
self.mutex.Lock()
// Read
n, err := self.reader.Read(p)
self.lastActivity = time.Now()
self.done = (err != nil)
self.mutex.Unlock()
if self.done {
self.stopTimer()
}
return n, err
}
func (self *TimeoutReader) startTimer() {
self.mutex.Lock()
defer self.mutex.Unlock()
if !self.done {
self.timer = time.AfterFunc(TimeoutTimerInterval, self.timeout)
}
}
func (self *TimeoutReader) stopTimer() {
self.mutex.Lock()
defer self.mutex.Unlock()
if self.timer != nil {
self.timer.Stop()
}
}
func (self *TimeoutReader) timeout() {
self.mutex.Lock()
if self.done {
self.mutex.Unlock()
return
}
if time.Since(self.lastActivity) > self.maxIdleTimeout {
self.cancel()
self.mutex.Unlock()
return
}
self.mutex.Unlock()
self.startTimer()
}
|
