Merge pull request #6 from vimeo/cancel-gcs-write-obj-ctx

Cancel WriteObject context on GCS copy error
vimeo · May 4, 2020 · 39e21d0 · 39e21d0
2 parents 5dafb74 + a0a5ea7
commit 39e21d0
Show file tree

Hide file tree

Showing 3 changed files with 134 additions and 69 deletions.
diff --git a/pkg/gcsstore/gcsstore.go b/pkg/gcsstore/gcsstore.go
@@ -123,8 +123,12 @@ func (upload gcsUpload) WriteChunk(ctx context.Context, offset int64, src io.Rea
 		ID:     cid,
 	}
 
-	n, err := store.Service.WriteObject(ctx, objectParams, src)
+	ctxCancel, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	n, err := store.Service.WriteObject(ctxCancel, objectParams, src)
 	if err != nil {
+		cancel()
 		return 0, err
 	}
 

diff --git a/pkg/gcsstore/gcsstore_mock_test.go b/pkg/gcsstore/gcsstore_mock_test.go
diff --git a/pkg/gcsstore/gcsstore_test.go b/pkg/gcsstore/gcsstore_test.go
@@ -421,9 +421,12 @@ func TestWriteChunk(t *testing.T) {
 	rGet := bytes.NewReader([]byte(mockReaderData))
 
 	ctx := context.Background()
+	ctxCancel, cancel := context.WithCancel(ctx)
+	defer cancel()
+
 	gomock.InOrder(
 		service.EXPECT().FilterObjects(ctx, filterParams).Return(partials, nil),
-		service.EXPECT().WriteObject(ctx, writeObjectParams, rGet).Return(int64(len(mockReaderData)), nil),
+		service.EXPECT().WriteObject(ctxCancel, writeObjectParams, rGet).Return(int64(len(mockReaderData)), nil),
 	)
 
 	upload, err := store.GetUpload(context.Background(), mockID)