| 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
 | package drive
import (
	"golang.org/x/net/context"
	"io"
	"sync"
	"time"
)
const MaxIdleTimeout = time.Second * 120
const TimeoutTimerInterval = time.Second * 10
type timeoutReaderWrapper func(io.Reader) io.Reader
func getTimeoutReaderWrapperContext() (timeoutReaderWrapper, context.Context) {
	ctx, cancel := context.WithCancel(context.TODO())
	wrapper := func(r io.Reader) io.Reader {
		return getTimeoutReader(r, cancel)
	}
	return wrapper, ctx
}
func getTimeoutReaderContext(r io.Reader) (io.Reader, context.Context) {
	ctx, cancel := context.WithCancel(context.TODO())
	return getTimeoutReader(r, cancel), ctx
}
func getTimeoutReader(r io.Reader, cancel context.CancelFunc) io.Reader {
	return &TimeoutReader{
		reader: r,
		cancel: cancel,
		mutex:  &sync.Mutex{},
	}
}
type TimeoutReader struct {
	reader       io.Reader
	cancel       context.CancelFunc
	lastActivity time.Time
	timer        *time.Timer
	mutex        *sync.Mutex
	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) > MaxIdleTimeout {
		self.cancel()
		self.mutex.Unlock()
		return
	}
	self.mutex.Unlock()
	self.startTimer()
}
 |