feat(storage): Add Flush for appendable writes (#11791)

* feat(storage): Add Flush for appendable writes

Adds Writer.Flush, which flushes any data in the internal buffer
to GCS.

Only supported for appendable writes. Preview feature.

* test and error fixes

* fix open writer test

* return non-EOF error.

* Update storage/client_test.go

Co-authored-by: Ashmeen Kaur <[email protected]>

* fix Flush after Close error

* fix data race

* add flush at close test

* fix conditions test

* fix other test issues

---------

Co-authored-by: Ashmeen Kaur <[email protected]>

Author: tritone

storage/client.go

@@ -281,6 +281,11 @@ type openWriterParams struct {
 	// setObj callback for reporting the resulting object - see `Writer.obj`.
 	// Required.
 	setObj func(*ObjectAttrs)
+	// setFlush callback for providing a Flush function implementation - see `Writer.Flush`.
+	// Required.
+	setFlush func(func() (int64, error))
+	// setPipeWriter callback for reseting `Writer.pw` if needed.
+	setPipeWriter func(*io.PipeWriter)
 }
 
 type newMultiRangeDownloaderParams struct {

storage/client_test.go

@@ -26,6 +26,7 @@ import (
 	"net/http"
 	"net/url"
 	"os"
+	"slices"
 	"strconv"
 	"strings"
 	"testing"
@@ -904,6 +905,7 @@ func TestOpenWriterEmulated(t *testing.T) {
 			setError: func(_ error) {}, // no-op
 			progress: func(_ int64) {}, // no-op
 			setObj:   func(o *ObjectAttrs) { gotAttrs = o },
+			setFlush: func(f func() (int64, error)) {},
 		}
 		pw, err := client.OpenWriter(params)
 		if err != nil {
@@ -997,7 +999,7 @@ func TestOpenAppendableWriterEmulated(t *testing.T) {
 	})
 }
 
-func TestOpenAppendableWriterMultipleFlushesEmulated(t *testing.T) {
+func TestOpenAppendableWriterMultipleChunksEmulated(t *testing.T) {
 	transportClientTest(skipHTTP("appends only supported via gRPC"), t, func(t *testing.T, ctx context.Context, project, bucket string, client storageClient) {
 		// Populate test data.
 		_, err := client.CreateBucket(ctx, project, bucket, &BucketAttrs{
@@ -1052,6 +1054,187 @@ func TestOpenAppendableWriterMultipleFlushesEmulated(t *testing.T) {
 	})
 }
 
+func TestWriterFlushEmulated(t *testing.T) {
+	transportClientTest(skipHTTP("appends only supported via gRPC"), t, func(t *testing.T, ctx context.Context, project, bucket string, client storageClient) {
+		// Populate test data.
+		_, err := client.CreateBucket(ctx, project, bucket, &BucketAttrs{
+			Name: bucket,
+		}, nil)
+		if err != nil {
+			t.Fatalf("client.CreateBucket: %v", err)
+		}
+		prefix := time.Now().Nanosecond()
+		objName := fmt.Sprintf("%d-object-%d", prefix, time.Now().Nanosecond())
+
+		vc := &Client{tc: client}
+		w := vc.Bucket(bucket).Object(objName).NewWriter(ctx)
+		w.Append = true
+		w.ChunkSize = 3 * MiB
+		var gotOffsets []int64
+		w.ProgressFunc = func(offset int64) {
+			gotOffsets = append(gotOffsets, offset)
+		}
+		// Flush at 0 and at a range of values, both aligned with ChunkSize and where explicit flushes are expected.
+		wantOffsets := []int64{0, 3 * MiB, 4*MiB + 1024, 4*MiB + 2048, 4*MiB + 3072, 6*MiB + 4096, 9 * MiB}
+
+		// Flush first with no data since this is a special case to test.
+		off, err := w.Flush()
+		if err != nil {
+			t.Errorf("flushing at 0: %v", err)
+		}
+		if off != 0 {
+			t.Errorf("flushing at 0: got %v offset, want 0", off)
+		}
+
+		// Write first 4 MiB data, expect progress every 3 MiB.
+		_, err = w.Write(randomBytes9MiB[0 : 4*MiB])
+		if err != nil {
+			t.Fatalf("writing test data: got %v; want ok", err)
+		}
+
+		// Write another 1KiB three times and do a Flush each time.
+		for i := 0; i < 3; i++ {
+			start := int64(4*MiB + i*1024)
+			end := int64(4*MiB + (i+1)*1024)
+			n, err := w.Write(randomBytes9MiB[start:end])
+			if err != nil {
+				t.Fatalf("writing 1k data: got %v; want ok", err)
+			}
+			if n != 1024 {
+				t.Errorf("writing 1k data: got %v bytes written, want %v", n, 1024)
+			}
+			off, err := w.Flush()
+
+			if err != nil {
+				t.Fatalf("flushing 1k data: got %v; want ok", err)
+			}
+			if off != end {
+				t.Errorf("flushing 1k data: got %v bytes committed, want %v", off, end)
+			}
+		}
+
+		// Do one more Flush that would require multiple messages (over 2 MiB).
+		n, err := w.Write(randomBytes9MiB[4*MiB+3072 : 6*MiB+4096])
+		if err != nil {
+			t.Fatalf("writing 2MiB + 1k data: got %v; want ok", err)
+		}
+		if n != 2*MiB+1024 {
+			t.Errorf("writing 2 MiB + 1k data: got %v bytes written, want %v", n, 2*MiB+1024)
+		}
+		off, err = w.Flush()
+		if err != nil {
+			t.Fatalf("flushing 2 MiB + 1k data: got %v; want ok", err)
+		}
+		if off != 6*MiB+4096 {
+			t.Errorf("flushing 2 MiB + 1k data: got %v bytes committed, want %v", off, 6*MiB+4096)
+		}
+
+		// Do one more zero-byte flush; expect noop for this.
+		off, err = w.Flush()
+		if err != nil {
+			t.Fatalf("flushing 0b data: got %v; want ok", err)
+		}
+		if off != 6*MiB+4096 {
+			t.Errorf("flushing 0b data: got %v bytes committed, want %v", off, 6*MiB+4096)
+		}
+
+		// Write remainder of data and close the writer.
+		if _, err := w.Write(randomBytes9MiB[6*MiB+4096:]); err != nil {
+			t.Fatalf("writing remaining data: got %v; want ok", err)
+		}
+		if err := w.Close(); err != nil {
+			t.Fatalf("closing writer: %v", err)
+		}
+
+		// Check that Flush after close fails.
+		if _, err = w.Flush(); err == nil {
+			t.Errorf("flush: expected error after close, got nil")
+		}
+
+		// Check offsets
+		if !slices.Equal(gotOffsets, wantOffsets) {
+			t.Errorf("progress offsets: got %v, want %v", gotOffsets, wantOffsets)
+		}
+
+		// Download object and check data
+		r, err := veneerClient.Bucket(bucket).Object(objName).NewReader(ctx)
+		defer r.Close()
+		if err != nil {
+			t.Fatalf("opening reading: %v", err)
+		}
+		wantLen := len(randomBytes9MiB)
+		got, err := io.ReadAll(r)
+		if n := len(got); n != wantLen {
+			t.Fatalf("expected to read %d bytes, but got %d (%v)", wantLen, n, err)
+		}
+		if diff := cmp.Diff(got, randomBytes9MiB); diff != "" {
+			t.Fatalf("checking written content: got(-), want(+):\n%s", diff)
+		}
+	})
+}
+
+func TestWriterFlushAtCloseEmulated(t *testing.T) {
+	transportClientTest(skipHTTP("appends only supported via gRPC"), t, func(t *testing.T, ctx context.Context, project, bucket string, client storageClient) {
+		// Populate test data.
+		_, err := client.CreateBucket(ctx, project, bucket, &BucketAttrs{
+			Name: bucket,
+		}, nil)
+		if err != nil {
+			t.Fatalf("client.CreateBucket: %v", err)
+		}
+		prefix := time.Now().Nanosecond()
+		objName := fmt.Sprintf("%d-object-%d", prefix, time.Now().Nanosecond())
+
+		vc := &Client{tc: client}
+		w := vc.Bucket(bucket).Object(objName).NewWriter(ctx)
+		w.Append = true
+		w.ChunkSize = MiB
+		var gotOffsets []int64
+		w.ProgressFunc = func(offset int64) {
+			gotOffsets = append(gotOffsets, offset)
+		}
+		wantOffsets := []int64{MiB, 2 * MiB, 3 * MiB}
+
+		// Test Flush right before close only.
+		n, err := w.Write(randomBytes3MiB)
+		if err != nil {
+			t.Fatalf("writing data: got %v; want ok", err)
+		}
+		if n != 3*MiB {
+			t.Errorf("writing data: got %v bytes written, want %v", n, 3*MiB)
+		}
+		off, err := w.Flush()
+		if err != nil {
+			t.Fatalf("flush: got %v; want ok", err)
+		}
+		if off != 3*MiB {
+			t.Errorf("flushing data: got %v bytes written, want %v", off, 3*MiB)
+		}
+		if err := w.Close(); err != nil {
+			t.Fatalf("closing writer: %v", err)
+		}
+		// Check offsets
+		if !slices.Equal(gotOffsets, wantOffsets) {
+			t.Errorf("progress offsets: got %v, want %v", gotOffsets, wantOffsets)
+		}
+
+		// Download object and check data
+		r, err := veneerClient.Bucket(bucket).Object(objName).NewReader(ctx)
+		defer r.Close()
+		if err != nil {
+			t.Fatalf("opening reading: %v", err)
+		}
+		wantLen := 3 * MiB
+		got, err := io.ReadAll(r)
+		if n := len(got); n != wantLen {
+			t.Fatalf("expected to read %d bytes, but got %d (%v)", wantLen, n, err)
+		}
+		if diff := cmp.Diff(got, randomBytes3MiB); diff != "" {
+			t.Fatalf("checking written content: got(-), want(+):\n%s", diff)
+		}
+	})
+}
+
 func TestListNotificationsEmulated(t *testing.T) {
 	transportClientTest(skipGRPC("notifications not implemented"), t, func(t *testing.T, ctx context.Context, project, bucket string, client storageClient) {
 		// Populate test object.
@@ -1705,6 +1888,7 @@ func TestObjectConditionsEmulated(t *testing.T) {
 						},
 						progress: nil,
 						setObj:   nil,
+						setFlush: func(f func() (int64, error)) {},
 					})
 					return err
 				},
@@ -2074,6 +2258,7 @@ func TestWriterChunkTransferTimeoutEmulated(t *testing.T) {
 					setError:             func(_ error) {}, // no-op
 					progress:             func(_ int64) {}, // no-op
 					setObj:               func(o *ObjectAttrs) { gotAttrs = o },
+					setFlush:             func(func() (int64, error)) {}, // no-op
 				}
 
 				pw, err := client.OpenWriter(params)
@@ -2168,6 +2353,7 @@ func TestWriterChunkRetryDeadlineEmulated(t *testing.T) {
 			setError:           func(_ error) {}, // no-op
 			progress:           func(_ int64) {}, // no-op
 			setObj:             func(_ *ObjectAttrs) {},
+			setFlush:           func(f func() (int64, error)) {},
 		}
 
 		pw, err := client.OpenWriter(params, &idempotentOption{true})

storage/grpc_client.go

@@ -1703,6 +1703,7 @@ func (c *grpcStorageClient) OpenWriter(params *openWriterParams, opts ...storage
 	var offset int64
 	errorf := params.setError
 	setObj := params.setObj
+	setFlush := params.setFlush
 	pr, pw := io.Pipe()
 
 	s := callSettings(c.settings, opts...)
@@ -1728,11 +1729,16 @@ func (c *grpcStorageClient) OpenWriter(params *openWriterParams, opts ...storage
 			}
 
 			var gw *gRPCWriter
-			gw, err := newGRPCWriter(c, s, params, r)
+			gw, err := newGRPCWriter(c, s, params, r, pw, params.setPipeWriter)
 			if err != nil {
 				return err
 			}
 
+			// Set Flush func for use by exported Writer.Flush.
+			setFlush(func() (int64, error) {
+				return gw.flush()
+			})
+
 			// Loop until there is an error or the Object has been finalized.
 			for {
 				// Note: This blocks until either the buffer is full or EOF is read.
@@ -2572,7 +2578,7 @@ func (r *gRPCReader) reopenStream() error {
 	return nil
 }
 
-func newGRPCWriter(c *grpcStorageClient, s *settings, params *openWriterParams, r io.Reader) (*gRPCWriter, error) {
+func newGRPCWriter(c *grpcStorageClient, s *settings, params *openWriterParams, r io.Reader, 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")
@@ -2609,6 +2615,7 @@ func newGRPCWriter(c *grpcStorageClient, s *settings, params *openWriterParams,
 		c:                     c,
 		ctx:                   params.ctx,
 		reader:                r,
+		pw:                    pw,
 		bucket:                params.bucket,
 		attrs:                 params.attrs,
 		conds:                 params.conds,
@@ -2620,15 +2627,19 @@ func newGRPCWriter(c *grpcStorageClient, s *settings, params *openWriterParams,
 		forceOneShot:          params.chunkSize <= 0,
 		forceEmptyContentType: params.forceEmptyContentType,
 		append:                params.append,
+		setPipeWriter:         setPipeWriter,
+		flushComplete:         make(chan int64),
 	}, nil
 }
 
 // gRPCWriter is a wrapper around the the gRPC client-stream API that manages
 // sending chunks of data provided by the user over the stream.
 type gRPCWriter struct {
-	c      *grpcStorageClient
-	buf    []byte
-	reader io.Reader
+	c             *grpcStorageClient
+	buf           []byte
+	reader        io.Reader
+	pw            *io.PipeWriter
+	setPipeWriter func(*io.PipeWriter) // used to set in parent storage.Writer
 
 	ctx context.Context
 
@@ -2645,7 +2656,9 @@ type gRPCWriter struct {
 	forceEmptyContentType bool
 	append                bool
 
-	streamSender gRPCBidiWriteBufferSender
+	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.
 }
 
 func bucketContext(ctx context.Context, bucket string) context.Context {
@@ -2945,11 +2958,17 @@ func (w *gRPCWriter) uploadBuffer(ctx context.Context, recvd int, start int64, d
 			break
 		}
 	}
+	if w.flushInProgress {
+		w.flushInProgress = false
+		w.flushComplete <- offset
+	}
 	return
 }
 
 // read copies the data in the reader to the given buffer and reports how much
 // data was read into the buffer and if there is no more data to read (EOF).
+// read returns when either 1. the buffer is full, 2. Writer.Flush was called,
+// or 3. Writer.Close was called.
 func (w *gRPCWriter) read() (int, bool, error) {
 	// Set n to -1 to start the Read loop.
 	var n, recvd int = -1, 0
@@ -2961,12 +2980,37 @@ func (w *gRPCWriter) read() (int, bool, error) {
 	}
 	var done bool
 	if err == io.EOF {
-		done = true
 		err = nil
+		// EOF can come from Writer.Flush or Writer.Close.
+		if w.flushInProgress {
+			// Reset pipe for additional writes after the flush.
+			pr, pw := io.Pipe()
+			w.reader = pr
+			w.pw = pw
+			w.setPipeWriter(pw)
+		} else {
+			done = true
+		}
 	}
 	return recvd, done, err
 }
 
+// flush flushes the current buffer regardless of whether it is full or not.
+// It's the implementation for Writer.Flush.
+func (w *gRPCWriter) flush() (int64, error) {
+	if !w.append {
+		return 0, errors.New("Flush is supported only if Writer.Append is set to true")
+	}
+
+	// Close PipeWriter to trigger EOF on read side of the stream.
+	w.flushInProgress = true
+	w.pw.Close()
+
+	// Wait for flush to complete
+	offset := <-w.flushComplete
+	return offset, nil
+}
+
 func checkCanceled(err error) error {
 	if status.Code(err) == codes.Canceled {
 		return context.Canceled

storage/http_client.go

@@ -970,6 +970,9 @@ func (c *httpStorageClient) OpenWriter(params *openWriterParams, opts ...storage
 	setObj := params.setObj
 	progress := params.progress
 	attrs := params.attrs
+	params.setFlush(func() (int64, error) {
+		return 0, errors.New("Writer.Flush is only supported for gRPC-based clients")
+	})
 
 	mediaOpts := []googleapi.MediaOption{
 		googleapi.ChunkSize(params.chunkSize),