Wrap media uploads in TimeoutReader

4 files changed, 30 insertions, 10 deletions

diff --git a/drive/sync_upload.go b/drive/sync_upload.go
index a94e507..7a0833a 100644
--- a/drive/sync_upload.go
+++ b/drive/sync_upload.go
@@ -305,9 +305,12 @@ func (self *Drive) uploadMissingFile(parentId string, lf *LocalFile, args Upload
     chunkSize := googleapi.ChunkSize(int(args.ChunkSize))
 
     // Wrap file in progress reader
-    srcReader := getProgressReader(srcFile, args.Progress, lf.info.Size())
+    progressReader := getProgressReader(srcFile, args.Progress, lf.info.Size())
 
-    _, err = self.service.Files.Create(dstFile).Fields("id", "name", "size", "md5Checksum").Media(srcReader, chunkSize).Do()
+    // Wrap reader in timeout reader
+    reader, ctx := getTimeoutReaderContext(progressReader)
+
+    _, err = self.service.Files.Create(dstFile).Fields("id", "name", "size", "md5Checksum").Context(ctx).Media(reader, chunkSize).Do()
     if err != nil {
         if isBackendError(err) && try < MaxBackendErrorRetries {
             exponentialBackoffSleep(try)
@@ -341,9 +344,12 @@ func (self *Drive) updateChangedFile(cf *changedFile, args UploadSyncArgs, try i
     chunkSize := googleapi.ChunkSize(int(args.ChunkSize))
 
     // Wrap file in progress reader
-    srcReader := getProgressReader(srcFile, args.Progress, cf.local.info.Size())
+    progressReader := getProgressReader(srcFile, args.Progress, cf.local.info.Size())
+
+    // Wrap reader in timeout reader
+    reader, ctx := getTimeoutReaderContext(progressReader)
 
-    _, err = self.service.Files.Update(cf.remote.file.Id, dstFile).Media(srcReader, chunkSize).Do()
+    _, err = self.service.Files.Update(cf.remote.file.Id, dstFile).Context(ctx).Media(reader, chunkSize).Do()
     if err != nil {
         if isBackendError(err) && try < MaxBackendErrorRetries {
             exponentialBackoffSleep(try)
diff --git a/drive/timeout_reader.go b/drive/timeout_reader.go
index e228160..ba2bb83 100644
--- a/drive/timeout_reader.go
+++ b/drive/timeout_reader.go
@@ -10,6 +10,11 @@ import (
 const MaxIdleTimeout = time.Second * 120
 const TimeoutTimerInterval = time.Second * 10
 
+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,
diff --git a/drive/update.go b/drive/update.go
index c4ee341..5bdd040 100644
--- a/drive/update.go
+++ b/drive/update.go
@@ -54,12 +54,15 @@ func (self *Drive) Update(args UpdateArgs) error {
     chunkSize := googleapi.ChunkSize(int(args.ChunkSize))
 
     // Wrap file in progress reader
-    srcReader := getProgressReader(srcFile, args.Progress, srcFileInfo.Size())
+    progressReader := getProgressReader(srcFile, args.Progress, srcFileInfo.Size())
+
+    // Wrap reader in timeout reader
+    reader, ctx := getTimeoutReaderContext(progressReader)
 
     fmt.Fprintf(args.Out, "Uploading %s\n", args.Path)
     started := time.Now()
 
-    f, err := self.service.Files.Update(args.Id, dstFile).Fields("id", "name", "size").Media(srcReader, chunkSize).Do()
+    f, err := self.service.Files.Update(args.Id, dstFile).Fields("id", "name", "size").Context(ctx).Media(reader, chunkSize).Do()
     if err != nil {
         return fmt.Errorf("Failed to upload file: %s", err)
     }
diff --git a/drive/upload.go b/drive/upload.go
index 05f52cd..2b8c7c3 100644
--- a/drive/upload.go
+++ b/drive/upload.go
@@ -158,12 +158,15 @@ func (self *Drive) uploadFile(args UploadArgs) (*drive.File, int64, error) {
     chunkSize := googleapi.ChunkSize(int(args.ChunkSize))
 
     // Wrap file in progress reader
-    srcReader := getProgressReader(srcFile, args.Progress, srcFileInfo.Size())
+    progressReader := getProgressReader(srcFile, args.Progress, srcFileInfo.Size())
+
+    // Wrap reader in timeout reader
+    reader, ctx := getTimeoutReaderContext(progressReader)
 
     fmt.Fprintf(args.Out, "Uploading %s\n", args.Path)
     started := time.Now()
 
-    f, err := self.service.Files.Create(dstFile).Fields("id", "name", "size", "md5Checksum", "webContentLink").Media(srcReader, chunkSize).Do()
+    f, err := self.service.Files.Create(dstFile).Fields("id", "name", "size", "md5Checksum", "webContentLink").Context(ctx).Media(reader, chunkSize).Do()
     if err != nil {
         return nil, 0, fmt.Errorf("Failed to upload file: %s", err)
     }
@@ -205,12 +208,15 @@ func (self *Drive) UploadStream(args UploadStreamArgs) error {
     chunkSize := googleapi.ChunkSize(int(args.ChunkSize))
 
     // Wrap file in progress reader
-    srcReader := getProgressReader(args.In, args.Progress, 0)
+    progressReader := getProgressReader(args.In, args.Progress, 0)
+
+    // Wrap reader in timeout reader
+    reader, ctx := getTimeoutReaderContext(progressReader)
 
     fmt.Fprintf(args.Out, "Uploading %s\n", dstFile.Name)
     started := time.Now()
 
-    f, err := self.service.Files.Create(dstFile).Fields("id", "name", "size", "webContentLink").Media(srcReader, chunkSize).Do()
+    f, err := self.service.Files.Create(dstFile).Fields("id", "name", "size", "webContentLink").Context(ctx).Media(reader, chunkSize).Do()
     if err != nil {
         return fmt.Errorf("Failed to upload file: %s", err)
     }