Configurable timeout as argument prasmussen#127

imransayed · Apr 9, 2016 · 0e1057e · 0e1057e
1 parent 401e017
commit 0e1057e
Show file tree

Hide file tree

Showing 12 changed files with 104 additions and 28 deletions.
diff --git a/.gitignore b/.gitignore
@@ -7,3 +7,4 @@ _release/bin
 Session.vim
 .netrwhist
 gdrive
+gdrive.sh
diff --git a/compare.go b/compare.go
@@ -1,8 +1,8 @@
 package main
 
 import (
-	"github.com/prasmussen/gdrive/drive"
 	"encoding/json"
+	"github.com/prasmussen/gdrive/drive"
 	"os"
 )
 

diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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()

diff --git a/gdrive.go b/gdrive.go
@@ -1,8 +1,8 @@
 package main
 
 import (
-	"github.com/prasmussen/gdrive/cli"
 	"fmt"
+	"github.com/prasmussen/gdrive/cli"
 	"os"
 )
 
@@ -14,6 +14,7 @@ const DefaultMaxChanges = 100
 const DefaultNameWidth = 40
 const DefaultPathWidth = 60
 const DefaultUploadChunkSize = 8 * 1024 * 1024
+const DefaultTimeout = 5 * 60
 const DefaultQuery = "trashed = false and 'me' in owners"
 const DefaultShareRole = "reader"
 const DefaultShareType = "anyone"
@@ -134,6 +135,12 @@ func main() {
 						Description: "Write file content to stdout",
 						OmitValue:   true,
 					},
+					cli.IntFlag{
+						Name:         "timeout",
+						Patterns:     []string{"--timeout"},
+						Description:  fmt.Sprintf("Set timeout in seconds, use 0 for no timeout. Timeout is reached when no data is transferred in set amount of seconds, default: %d", DefaultTimeout),
+						DefaultValue: DefaultTimeout,
+					},
 				),
 			},
 		},
@@ -216,6 +223,12 @@ func main() {
 						Description: "Delete local file when upload is successful",
 						OmitValue:   true,
 					},
+					cli.IntFlag{
+						Name:         "timeout",
+						Patterns:     []string{"--timeout"},
+						Description:  fmt.Sprintf("Set timeout in seconds, use 0 for no timeout. Timeout is reached when no data is transferred in set amount of seconds, default: %d", DefaultTimeout),
+						DefaultValue: DefaultTimeout,
+					},
 					cli.IntFlag{
 						Name:         "chunksize",
 						Patterns:     []string{"--chunksize"},
@@ -254,6 +267,12 @@ func main() {
 						Description: "Share file",
 						OmitValue:   true,
 					},
+					cli.IntFlag{
+						Name:         "timeout",
+						Patterns:     []string{"--timeout"},
+						Description:  fmt.Sprintf("Set timeout in seconds, use 0 for no timeout. Timeout is reached when no data is transferred in set amount of seconds, default: %d", DefaultTimeout),
+						DefaultValue: DefaultTimeout,
+					},
 					cli.BoolFlag{
 						Name:        "noProgress",
 						Patterns:    []string{"--no-progress"},
@@ -291,6 +310,12 @@ func main() {
 						Patterns:    []string{"--mime"},
 						Description: "Force mime type",
 					},
+					cli.IntFlag{
+						Name:         "timeout",
+						Patterns:     []string{"--timeout"},
+						Description:  fmt.Sprintf("Set timeout in seconds, use 0 for no timeout. Timeout is reached when no data is transferred in set amount of seconds, default: %d", DefaultTimeout),
+						DefaultValue: DefaultTimeout,
+					},
 					cli.IntFlag{
 						Name:         "chunksize",
 						Patterns:     []string{"--chunksize"},
@@ -494,6 +519,12 @@ func main() {
 						Description: "Hide progress",
 						OmitValue:   true,
 					},
+					cli.IntFlag{
+						Name:         "timeout",
+						Patterns:     []string{"--timeout"},
+						Description:  fmt.Sprintf("Set timeout in seconds, use 0 for no timeout. Timeout is reached when no data is transferred in set amount of seconds, default: %d", DefaultTimeout),
+						DefaultValue: DefaultTimeout,
+					},
 				),
 			},
 		},
@@ -540,6 +571,12 @@ func main() {
 						Description: "Hide progress",
 						OmitValue:   true,
 					},
+					cli.IntFlag{
+						Name:         "timeout",
+						Patterns:     []string{"--timeout"},
+						Description:  fmt.Sprintf("Set timeout in seconds, use 0 for no timeout. Timeout is reached when no data is transferred in set amount of seconds, default: %d", DefaultTimeout),
+						DefaultValue: DefaultTimeout,
+					},
 					cli.IntFlag{
 						Name:         "chunksize",
 						Patterns:     []string{"--chunksize"},
@@ -647,6 +684,12 @@ func main() {
 						Patterns:    []string{"--path"},
 						Description: "Download path",
 					},
+					cli.IntFlag{
+						Name:         "timeout",
+						Patterns:     []string{"--timeout"},
+						Description:  fmt.Sprintf("Set timeout in seconds, use 0 for no timeout. Timeout is reached when no data is transferred in set amount of seconds, default: %d", DefaultTimeout),
+						DefaultValue: DefaultTimeout,
+					},
 				),
 			},
 		},