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{