fix(storage): fix append writer hang (#12201)

tritone · web-flow · commit 7ce2a2ad3ae9 · 2025-05-05T23:33:18.000-04:00
After a Flush call, we weren't correctly tracking the pipe
reader to close to signal to Write that an error had happened.
Fixes that issue and adds another flush to the test, which
triggers the hang without this fix.
diff --git a/storage/grpc_writer.go b/storage/grpc_writer.go
@@ -102,7 +102,7 @@ func (c *grpcStorageClient) OpenWriter(params *openWriterParams, opts ...storage
 	setFlush(func() (int64, error) {
 		return gw.flush()
 	})
-	gw, err := newGRPCWriter(c, s, params, pr, pw, params.setPipeWriter)
+	gw, err := newGRPCWriter(c, s, params, pr, pr, pw, params.setPipeWriter)
 	if err != nil {
 		errorf(err)
 		pr.CloseWithError(err)
@@ -188,14 +188,14 @@ func (c *grpcStorageClient) OpenWriter(params *openWriterParams, opts ...storage
 		// These calls are still valid if err is nil
 		err = checkCanceled(err)
 		errorf(err)
-		pr.CloseWithError(err)
+		gw.pr.CloseWithError(err)
 		close(params.donec)
 	}()
 
 	return pw, nil
 }
 
-func newGRPCWriter(c *grpcStorageClient, s *settings, params *openWriterParams, r io.Reader, pw *io.PipeWriter, setPipeWriter func(*io.PipeWriter)) (*gRPCWriter, error) {
+func newGRPCWriter(c *grpcStorageClient, s *settings, params *openWriterParams, r io.Reader, pr *io.PipeReader, pw *io.PipeWriter, setPipeWriter func(*io.PipeWriter)) (*gRPCWriter, error) {
 	if params.attrs.Retention != nil {
 		// TO-DO: remove once ObjectRetention is available - see b/308194853
 		return nil, status.Errorf(codes.Unimplemented, "storage: object retention is not supported in gRPC")
@@ -241,6 +241,7 @@ func newGRPCWriter(c *grpcStorageClient, s *settings, params *openWriterParams,
 		ctx:                   params.ctx,
 		reader:                r,
 		pw:                    pw,
+		pr:                    pr,
 		bucket:                params.bucket,
 		attrs:                 params.attrs,
 		conds:                 params.conds,
@@ -266,6 +267,7 @@ type gRPCWriter struct {
 	c             *grpcStorageClient
 	buf           []byte
 	reader        io.Reader
+	pr            *io.PipeReader // Keep track of pr and pw to update post-flush
 	pw            *io.PipeWriter
 	setPipeWriter func(*io.PipeWriter) // used to set in parent storage.Writer
 
@@ -628,6 +630,7 @@ func (w *gRPCWriter) read() (int, bool, error) {
 			pr, pw := io.Pipe()
 			w.reader = pr
 			w.pw = pw
+			w.pr = pr
 			w.setPipeWriter(pw)
 		} else {
 			done = true
diff --git a/storage/integration_test.go b/storage/integration_test.go
@@ -3557,6 +3557,9 @@ func TestIntegration_WriterAppendEdgeCases(t *testing.T) {
 		if _, err := w.Write(randomBytes3MiB); err != nil {
 			t.Fatalf("w.Write: %v", err)
 		}
+		if _, err := w.Flush(); err != nil {
+			t.Fatalf("w.Flush: %v", err)
+		}
 
 		tw, _, err := obj.Generation(w.Attrs().Generation).NewWriterFromAppendableObject(ctx, nil)
 		if err != nil {

Original file line number	Diff line number	Diff line change
`@@ -3557,6 +3557,9 @@ func TestIntegration_WriterAppendEdgeCases(t *testing.T) {`
`3557`	`3557`	`if _, err := w.Write(randomBytes3MiB); err != nil {`
`3558`	`3558`	`t.Fatalf("w.Write: %v", err)`
`3559`	`3559`	`}`
	`3560`	`+ if _, err := w.Flush(); err != nil {`
	`3561`	`+ t.Fatalf("w.Flush: %v", err)`
	`3562`	`+ }`
`3560`	`3563`
`3561`	`3564`	`tw, _, err := obj.Generation(w.Attrs().Generation).NewWriterFromAppendableObject(ctx, nil)`
`3562`	`3565`	`if err != nil {`