Configurable timeout as argument #127

7 files changed, 44 insertions, 24 deletions

diff --git a/drive/download.go b/drive/download.go
index 0856258..3cb8310 100644
--- a/drive/download.go
+++ b/drive/download.go
@@ -19,6 +19,7 @@ type DownloadArgs struct {
 	Recursive bool
 	Delete    bool
 	Stdout    bool
+	Timeout   time.Duration
 }
 
 func (self *Drive) Download(args DownloadArgs) error {
@@ -120,7 +121,7 @@ func (self *Drive) downloadRecursive(args DownloadArgs) error {
 
 func (self *Drive) downloadBinary(f *drive.File, args DownloadArgs) (int64, int64, error) {
 	// Get timeout reader wrapper and context
-	timeoutReaderWrapper, ctx := getTimeoutReaderWrapperContext()
+	timeoutReaderWrapper, ctx := getTimeoutReaderWrapperContext(args.Timeout)
 
 	res, err := self.service.Files.Get(f.Id).Context(ctx).Download()
 	if err != nil {
diff --git a/drive/revision_download.go b/drive/revision_download.go
index 04055fa..57392a5 100644
--- a/drive/revision_download.go
+++ b/drive/revision_download.go
@@ -5,6 +5,7 @@ import (
 	"io"
 	"io/ioutil"
 	"path/filepath"
+	"time"
 )
 
 type DownloadRevisionArgs struct {
@@ -15,6 +16,7 @@ type DownloadRevisionArgs struct {
 	Path       string
 	Force      bool
 	Stdout     bool
+	Timeout    time.Duration
 }
 
 func (self *Drive) DownloadRevision(args DownloadRevisionArgs) (err error) {
@@ -30,7 +32,7 @@ func (self *Drive) DownloadRevision(args DownloadRevisionArgs) (err error) {
 	}
 
 	// Get timeout reader wrapper and context
-	timeoutReaderWrapper, ctx := getTimeoutReaderWrapperContext()
+	timeoutReaderWrapper, ctx := getTimeoutReaderWrapperContext(args.Timeout)
 
 	res, err := getRev.Context(ctx).Download()
 	if err != nil {
diff --git a/drive/sync_download.go b/drive/sync_download.go
index 04b50b9..2494557 100644
--- a/drive/sync_download.go
+++ b/drive/sync_download.go
@@ -19,6 +19,7 @@ type DownloadSyncArgs struct {
 	Path             string
 	DryRun           bool
 	DeleteExtraneous bool
+	Timeout          time.Duration
 	Resolution       ConflictResolution
 	Comparer         FileComparer
 }
@@ -188,7 +189,7 @@ func (self *Drive) downloadRemoteFile(id, fpath string, args DownloadSyncArgs, t
 	}
 
 	// Get timeout reader wrapper and context
-	timeoutReaderWrapper, ctx := getTimeoutReaderWrapperContext()
+	timeoutReaderWrapper, ctx := getTimeoutReaderWrapperContext(args.Timeout)
 
 	res, err := self.service.Files.Get(id).Context(ctx).Download()
 	if err != nil {
diff --git a/drive/sync_upload.go b/drive/sync_upload.go
index 0d5c208..c509c0a 100644
--- a/drive/sync_upload.go
+++ b/drive/sync_upload.go
@@ -20,6 +20,7 @@ type UploadSyncArgs struct {
 	DryRun           bool
 	DeleteExtraneous bool
 	ChunkSize        int64
+	Timeout          time.Duration
 	Resolution       ConflictResolution
 	Comparer         FileComparer
 }
@@ -308,7 +309,7 @@ func (self *Drive) uploadMissingFile(parentId string, lf *LocalFile, args Upload
 	progressReader := getProgressReader(srcFile, args.Progress, lf.info.Size())
 
 	// Wrap reader in timeout reader
-	reader, ctx := getTimeoutReaderContext(progressReader)
+	reader, ctx := getTimeoutReaderContext(progressReader, args.Timeout)
 
 	_, err = self.service.Files.Create(dstFile).Fields("id", "name", "size", "md5Checksum").Context(ctx).Media(reader, chunkSize).Do()
 	if err != nil {
@@ -347,7 +348,7 @@ func (self *Drive) updateChangedFile(cf *changedFile, args UploadSyncArgs, try i
 	progressReader := getProgressReader(srcFile, args.Progress, cf.local.info.Size())
 
 	// Wrap reader in timeout reader
-	reader, ctx := getTimeoutReaderContext(progressReader)
+	reader, ctx := getTimeoutReaderContext(progressReader, args.Timeout)
 
 	_, err = self.service.Files.Update(cf.remote.file.Id, dstFile).Context(ctx).Media(reader, chunkSize).Do()
 	if err != nil {
diff --git a/drive/timeout_reader.go b/drive/timeout_reader.go
index 9930c12..67fd5b0 100644
--- a/drive/timeout_reader.go
+++ b/drive/timeout_reader.go
@@ -7,39 +7,51 @@ import (
 	"time"
 )
 
-const MaxIdleTimeout = time.Second * 120
 const TimeoutTimerInterval = time.Second * 10
 
 type timeoutReaderWrapper func(io.Reader) io.Reader
 
-func getTimeoutReaderWrapperContext() (timeoutReaderWrapper, context.Context) {
+func getTimeoutReaderWrapperContext(timeout time.Duration) (timeoutReaderWrapper, context.Context) {
 	ctx, cancel := context.WithCancel(context.TODO())
 	wrapper := func(r io.Reader) io.Reader {
-		return getTimeoutReader(r, cancel)
+		// 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) (io.Reader, context.Context) {
+func getTimeoutReaderContext(r io.Reader, timeout time.Duration) (io.Reader, context.Context) {
 	ctx, cancel := context.WithCancel(context.TODO())
-	return getTimeoutReader(r, cancel), ctx
+
+	// 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) io.Reader {
+func getTimeoutReader(r io.Reader, cancel context.CancelFunc, timeout time.Duration) io.Reader {
 	return &TimeoutReader{
-		reader: r,
-		cancel: cancel,
-		mutex:  &sync.Mutex{},
+		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
-	done         bool
+	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) {
@@ -90,7 +102,7 @@ func (self *TimeoutReader) timeout() {
 		return
 	}
 
-	if time.Since(self.lastActivity) > MaxIdleTimeout {
+	if time.Since(self.lastActivity) > self.maxIdleTimeout {
 		self.cancel()
 		self.mutex.Unlock()
 		return
diff --git a/drive/update.go b/drive/update.go
index 156eb2f..7af403f 100644
--- a/drive/update.go
+++ b/drive/update.go
@@ -20,6 +20,7 @@ type UpdateArgs struct {
 	Mime      string
 	Recursive bool
 	ChunkSize int64
+	Timeout   time.Duration
 }
 
 func (self *Drive) Update(args UpdateArgs) error {
@@ -57,7 +58,7 @@ func (self *Drive) Update(args UpdateArgs) error {
 	progressReader := getProgressReader(srcFile, args.Progress, srcFileInfo.Size())
 
 	// Wrap reader in timeout reader
-	reader, ctx := getTimeoutReaderContext(progressReader)
+	reader, ctx := getTimeoutReaderContext(progressReader, args.Timeout)
 
 	fmt.Fprintf(args.Out, "Uploading %s\n", args.Path)
 	started := time.Now()
diff --git a/drive/upload.go b/drive/upload.go
index c42bebd..6f5edd5 100644
--- a/drive/upload.go
+++ b/drive/upload.go
@@ -22,6 +22,7 @@ type UploadArgs struct {
 	Share     bool
 	Delete    bool
 	ChunkSize int64
+	Timeout   time.Duration
 }
 
 func (self *Drive) Upload(args UploadArgs) error {
@@ -173,7 +174,7 @@ func (self *Drive) uploadFile(args UploadArgs) (*drive.File, int64, error) {
 	progressReader := getProgressReader(srcFile, args.Progress, srcFileInfo.Size())
 
 	// Wrap reader in timeout reader
-	reader, ctx := getTimeoutReaderContext(progressReader)
+	reader, ctx := getTimeoutReaderContext(progressReader, args.Timeout)
 
 	fmt.Fprintf(args.Out, "Uploading %s\n", args.Path)
 	started := time.Now()
@@ -198,6 +199,7 @@ type UploadStreamArgs struct {
 	Share     bool
 	ChunkSize int64
 	Progress  io.Writer
+	Timeout   time.Duration
 }
 
 func (self *Drive) UploadStream(args UploadStreamArgs) error {
@@ -223,7 +225,7 @@ func (self *Drive) UploadStream(args UploadStreamArgs) error {
 	progressReader := getProgressReader(args.In, args.Progress, 0)
 
 	// Wrap reader in timeout reader
-	reader, ctx := getTimeoutReaderContext(progressReader)
+	reader, ctx := getTimeoutReaderContext(progressReader, args.Timeout)
 
 	fmt.Fprintf(args.Out, "Uploading %s\n", dstFile.Name)
 	started := time.Now()