Configurable timeout as argument #127

12 files changed, 104 insertions, 28 deletions

diff --git a/.gitignore b/.gitignore
index ee3fac5..62df61a 100644
--- a/.gitignore
+++ b/.gitignore
@@ -7,3 +7,4 @@ _release/bin
 Session.vim
 .netrwhist
 gdrive
+gdrive.sh
diff --git a/compare.go b/compare.go
index ac68284..35c0369 100644
--- 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
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()
diff --git a/gdrive.go b/gdrive.go
index c68d43d..e7d28cf 100644
--- 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,
+					},
 				),
 			},
 		},
@@ -217,6 +224,12 @@ func main() {
 						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"},
 						Description:  fmt.Sprintf("Set chunk size in bytes, default: %d", DefaultUploadChunkSize),
@@ -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"},
@@ -292,6 +311,12 @@ func main() {
 						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"},
 						Description:  fmt.Sprintf("Set chunk size in bytes, default: %d", DefaultUploadChunkSize),
@@ -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,
+					},
 				),
 			},
 		},
@@ -541,6 +572,12 @@ func main() {
 						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"},
 						Description:  fmt.Sprintf("Set chunk size in bytes, default: %d", DefaultUploadChunkSize),
@@ -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,
+					},
 				),
 			},
 		},
diff --git a/handlers_drive.go b/handlers_drive.go
index f0386ce..8db7329 100644
--- a/handlers_drive.go
+++ b/handlers_drive.go
@@ -1,15 +1,16 @@
 package main
 
 import (
+	"fmt"
 	"github.com/prasmussen/gdrive/auth"
 	"github.com/prasmussen/gdrive/cli"
 	"github.com/prasmussen/gdrive/drive"
-	"fmt"
 	"io"
 	"io/ioutil"
 	"net/http"
 	"os"
 	"path/filepath"
+	"time"
 )
 
 const ClientId = "367116221053-7n0vf5akeru7on6o2fjinrecpdoe99eg.apps.googleusercontent.com"
@@ -57,6 +58,7 @@ func downloadHandler(ctx cli.Context) {
 		Recursive: args.Bool("recursive"),
 		Stdout:    args.Bool("stdout"),
 		Progress:  progressWriter(args.Bool("noProgress")),
+		Timeout:   durationInSeconds(args.Int64("timeout")),
 	})
 	checkErr(err)
 }
@@ -84,6 +86,7 @@ func downloadSyncHandler(ctx cli.Context) {
 		RootId:           args.String("fileId"),
 		DryRun:           args.Bool("dryRun"),
 		DeleteExtraneous: args.Bool("deleteExtraneous"),
+		Timeout:          durationInSeconds(args.Int64("timeout")),
 		Resolution:       conflictResolution(args),
 		Comparer:         NewCachedMd5Comparer(cachePath),
 	})
@@ -100,6 +103,7 @@ func downloadRevisionHandler(ctx cli.Context) {
 		Stdout:     args.Bool("stdout"),
 		Path:       args.String("path"),
 		Progress:   progressWriter(args.Bool("noProgress")),
+		Timeout:    durationInSeconds(args.Int64("timeout")),
 	})
 	checkErr(err)
 }
@@ -118,6 +122,7 @@ func uploadHandler(ctx cli.Context) {
 		Share:     args.Bool("share"),
 		Delete:    args.Bool("delete"),
 		ChunkSize: args.Int64("chunksize"),
+		Timeout:   durationInSeconds(args.Int64("timeout")),
 	})
 	checkErr(err)
 }
@@ -132,6 +137,7 @@ func uploadStdinHandler(ctx cli.Context) {
 		Mime:      args.String("mime"),
 		Share:     args.Bool("share"),
 		ChunkSize: args.Int64("chunksize"),
+		Timeout:   durationInSeconds(args.Int64("timeout")),
 		Progress:  progressWriter(args.Bool("noProgress")),
 	})
 	checkErr(err)
@@ -148,6 +154,7 @@ func uploadSyncHandler(ctx cli.Context) {
 		DryRun:           args.Bool("dryRun"),
 		DeleteExtraneous: args.Bool("deleteExtraneous"),
 		ChunkSize:        args.Int64("chunksize"),
+		Timeout:          durationInSeconds(args.Int64("timeout")),
 		Resolution:       conflictResolution(args),
 		Comparer:         NewCachedMd5Comparer(cachePath),
 	})
@@ -165,6 +172,7 @@ func updateHandler(ctx cli.Context) {
 		Mime:      args.String("mime"),
 		Progress:  progressWriter(args.Bool("noProgress")),
 		ChunkSize: args.Int64("chunksize"),
+		Timeout:   durationInSeconds(args.Int64("timeout")),
 	})
 	checkErr(err)
 }
@@ -385,6 +393,10 @@ func progressWriter(discard bool) io.Writer {
 	return os.Stderr
 }
 
+func durationInSeconds(seconds int64) time.Duration {
+	return time.Second * time.Duration(seconds)
+}
+
 func conflictResolution(args cli.Arguments) drive.ConflictResolution {
 	keepLocal := args.Bool("keepLocal")
 	keepRemote := args.Bool("keepRemote")
diff --git a/handlers_meta.go b/handlers_meta.go
index bfd5b8f..2e1c95c 100644
--- a/handlers_meta.go
+++ b/handlers_meta.go
@@ -1,8 +1,8 @@
 package main
 
 import (
-	"github.com/prasmussen/gdrive/cli"
 	"fmt"
+	"github.com/prasmussen/gdrive/cli"
 	"os"
 	"runtime"
 	"strings"