chore(storage): replace Canceled status with context.Canceled (#4896)

storage/integration_test.go

@@ -1111,7 +1111,6 @@ func TestIntegration_SimpleWriteGRPC(t *testing.T) {
 }
 
 func TestIntegration_CancelWriteGRPC(t *testing.T) {
-	t.Skip("https://github.com/googleapis/google-cloud-go/issues/4772")
 	ctx := context.Background()
 
 	// Create an HTTP client to verify test and a gRPC client to test writing.
@@ -1129,9 +1128,12 @@ func TestIntegration_CancelWriteGRPC(t *testing.T) {
 	}()
 
 	cctx, cancel := context.WithCancel(ctx)
-	content := []byte("Hello, world this is a grpc request")
 
 	w := gobj.NewWriter(cctx)
+	// Set a chunk size and send that amount of data to provoke a network call
+	// on the next Write that would fail due to context cancelation.
+	w.ChunkSize = googleapi.MinUploadChunkSize
+	content := make([]byte, w.ChunkSize)
 	_, err := w.Write(content)
 	if err != nil {
 		t.Fatal(err)

storage/writer.go

@@ -27,6 +27,8 @@ import (
 	"google.golang.org/api/googleapi"
 	raw "google.golang.org/api/storage/v1"
 	storagepb "google.golang.org/genproto/googleapis/storage/v2"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 )
 
 const (
@@ -353,6 +355,7 @@ func (w *Writer) openGRPC() error {
 			// Note: This blocks until either the buffer is full or EOF is read.
 			recvd, doneReading, err := read(pr, buf)
 			if err != nil {
+				err = checkCanceled(err)
 				w.error(err)
 				pr.CloseWithError(err)
 				return
@@ -369,6 +372,7 @@ func (w *Writer) openGRPC() error {
 			if !doneReading && w.upid == "" {
 				err = w.startResumableUpload()
 				if err != nil {
+					err = checkCanceled(err)
 					w.error(err)
 					pr.CloseWithError(err)
 					return
@@ -377,6 +381,7 @@ func (w *Writer) openGRPC() error {
 
 			o, off, finalized, err := w.uploadBuffer(toWrite, recvd, offset, doneReading)
 			if err != nil {
+				err = checkCanceled(err)
 				w.error(err)
 				pr.CloseWithError(err)
 				return
@@ -637,3 +642,11 @@ func read(r io.Reader, buf []byte) (int, bool, error) {
 	}
 	return recvd, done, err
 }
+
+func checkCanceled(err error) error {
+	if status.Code(err) == codes.Canceled {
+		return context.Canceled
+	}
+
+	return err
+}