fix(storage): fix append edge cases (#12074)

* fix(storage): fix append edge cases

Fix the validation for generation being included. Fix a case
where the error was not being picked up during append takeover.
Add an integration test for these issues

* fix flush hang

* add another takeover case

Author: tritone

storage/grpc_writer.go

@@ -127,11 +127,17 @@ func (c *grpcStorageClient) OpenWriter(params *openWriterParams, opts ...storage
 					bctx = bucketContext(bctx, gw.bucket)
 				}
 				err = run(bctx, uploadBuff, gw.settings.retry, s.idempotent)
+				offset += int64(recvd)
+				// If this buffer upload was triggered by a flush, reset and
+				// communicate back the result.
+				if gw.flushInProgress {
+					gw.setSize(offset)
+					gw.flushInProgress = false
+					gw.flushComplete <- flushResult{offset: offset, err: err}
+				}
 				if err != nil {
 					return err
 				}
-				offset += int64(recvd)
-
 				// When we are done reading data without errors, set the object and
 				// finish.
 				if doneReading {
@@ -213,7 +219,7 @@ func newGRPCWriter(c *grpcStorageClient, s *settings, params *openWriterParams,
 		append:                params.append,
 		finalizeOnClose:       params.finalizeOnClose,
 		setPipeWriter:         setPipeWriter,
-		flushComplete:         make(chan int64),
+		flushComplete:         make(chan flushResult),
 	}, nil
 }
 
@@ -245,8 +251,13 @@ type gRPCWriter struct {
 	finalizeOnClose       bool
 
 	streamSender    gRPCBidiWriteBufferSender
-	flushInProgress bool       // true when the pipe is being recreated for a flush.
-	flushComplete   chan int64 // use to signal back to flush call that flush to server was completed.
+	flushInProgress bool             // true when the pipe is being recreated for a flush.
+	flushComplete   chan flushResult // use to signal back to flush call that flush to server was completed.
+}
+
+type flushResult struct {
+	err    error
+	offset int64
 }
 
 func bucketContext(ctx context.Context, bucket string) context.Context {
@@ -554,11 +565,6 @@ func (w *gRPCWriter) uploadBuffer(ctx context.Context, recvd int, start int64, d
 			break
 		}
 	}
-	if w.flushInProgress {
-		w.setSize(offset)
-		w.flushInProgress = false
-		w.flushComplete <- offset
-	}
 	return
 }
 
@@ -604,8 +610,8 @@ func (w *gRPCWriter) flush() (int64, error) {
 	w.pw.Close()
 
 	// Wait for flush to complete
-	offset := <-w.flushComplete
-	return offset, nil
+	result := <-w.flushComplete
+	return result.offset, result.err
 }
 
 func checkCanceled(err error) error {
@@ -686,6 +692,11 @@ func (w *gRPCWriter) newGRPCAppendTakeoverWriteBufferSender(ctx context.Context)
 		return nil, err
 	}
 	firstResp := <-s.recvs
+	// Check recvErr after getting the response.
+	if s.recvErr != nil {
+		return nil, s.recvErr
+	}
+
 	// Object resource is returned in the first response on takeover, so capture
 	// this now.
 	s.objResource = firstResp.GetResource()

storage/integration_test.go

@@ -3524,6 +3524,87 @@ func TestIntegration_WriterAppendTakeover(t *testing.T) {
 	})
 }
 
+func TestIntegration_WriterAppendEdgeCases(t *testing.T) {
+	t.Skip("b/402283880")
+	ctx := skipAllButBidi(context.Background(), "ZB test")
+	multiTransportTest(ctx, t, func(t *testing.T, ctx context.Context, _, prefix string, client *Client) {
+		h := testHelper{t}
+		bucketName := prefix + uidSpace.New()
+		bkt := client.Bucket(bucketName)
+
+		h.mustCreateZonalBucket(bkt, testutil.ProjID())
+		defer h.mustDeleteBucket(bkt)
+
+		objName := "object1"
+		obj := bkt.Object(objName)
+		defer h.mustDeleteObject(obj)
+
+		// Takeover Writer to a non-existent object should fail, with or
+		// without generation specified.
+		if _, _, err := obj.NewWriterFromAppendableObject(ctx, &AppendableWriterOpts{}); err == nil {
+			t.Errorf("NewWriterFromAppendableObject: got nil, want error")
+		}
+		_, _, err := obj.Generation(1234).NewWriterFromAppendableObject(ctx, &AppendableWriterOpts{})
+		if status.Code(err) != codes.NotFound {
+			t.Errorf("NewWriterFromAppendableObject: got %v, want NotFound", err)
+		}
+
+		// If a takeover is opened, flush or close to the original writer
+		// should fail.
+		w := obj.NewWriter(ctx)
+		w.Append = true
+		w.ChunkSize = MiB
+		if _, err := w.Write(randomBytes3MiB); err != nil {
+			t.Fatalf("w.Write: %v", err)
+		}
+
+		tw, _, err := obj.Generation(w.Attrs().Generation).NewWriterFromAppendableObject(ctx, nil)
+		if err != nil {
+			t.Fatalf("NewWriterFromAppendableObject: %v", err)
+		}
+		if _, err := tw.Write([]byte("hello world")); err != nil {
+			t.Fatalf("tw.Write: %v", err)
+		}
+		if _, err := tw.Flush(); err != nil {
+			t.Fatalf("tw.Flush: %v", err)
+		}
+
+		// Expect precondition error when writer to orginal Writer.
+		if _, err := w.Write(randomBytes3MiB); status.Code(err) != codes.FailedPrecondition {
+			t.Fatalf("got %v", err)
+		}
+
+		// Another NewWriter to the unfinalized object should also return a
+		// precondition error when data is flushed.
+		w2 := obj.NewWriter(ctx)
+		w2.Append = true
+		if _, err := w2.Write([]byte("hello world")); err != nil {
+			t.Fatalf("w2.Write: %v", err)
+		}
+		if _, err := w2.Flush(); status.Code(err) != codes.FailedPrecondition {
+			t.Fatalf("w2.Flush: %v", err)
+		}
+
+		// If we add yet another takeover writer to finalize and delete the object,
+		// tw should also return an error on flush.
+		tw2, _, err := obj.Generation(w.Attrs().Generation).NewWriterFromAppendableObject(ctx, &AppendableWriterOpts{
+			FinalizeOnClose: true,
+		})
+		if err != nil {
+			t.Fatalf("NewWriterFromAppendableObject: %v", err)
+		}
+		if err := tw2.Close(); err != nil {
+			t.Fatalf("tw2.Close: %v", err)
+		}
+		h.mustDeleteObject(obj)
+		if _, err := tw.Write([]byte("abcde")); err != nil {
+			t.Fatalf("tw.Write: %v", err)
+		}
+		if _, err := tw.Flush(); status.Code(err) != codes.FailedPrecondition {
+			t.Errorf("tw.Flush: got %v, want FailedPrecondition", err)
+		}
+	})
+}
 func TestIntegration_ZeroSizedObject(t *testing.T) {
 	t.Parallel()
 	multiTransportTest(context.Background(), t, func(t *testing.T, ctx context.Context, bucket, _ string, client *Client) {

storage/storage.go

@@ -1267,7 +1267,7 @@ func (o *ObjectHandle) NewWriter(ctx context.Context) *Writer {
 // This feature is in preview and is not yet available for general use.
 func (o *ObjectHandle) NewWriterFromAppendableObject(ctx context.Context, opts *AppendableWriterOpts) (*Writer, int64, error) {
 	ctx = trace.StartSpan(ctx, "cloud.google.com/go/storage.Object.Writer")
-	if o.gen == 0 {
+	if o.gen < 0 {
 		return nil, 0, errors.New("storage: ObjectHandle.Generation must be set to use NewWriterFromAppendableObject")
 	}
 	w := &Writer{