Add TimeoutReader

TimeoutReader wraps a reader and takes a cancel function as argument, the cancel function will be called when the reader is idle for too long.
antkss · Feb 20, 2016 · 7eaf0c8 · 7eaf0c8
1 parent 5d56138
commit 7eaf0c8
Showing 1 changed file with 86 additions and 0 deletions.
diff --git a/drive/timeout_reader.go b/drive/timeout_reader.go
@@ -0,0 +1,86 @@
+package drive
+
+import (
+    "io"
+    "time"
+    "sync"
+    "golang.org/x/net/context"
+)
+
+const MaxIdleTimeout = time.Second * 120
+const TimeoutTimerInterval = time.Second * 10
+
+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()
+}