fix(storage): optimize gRPC writer with zero-copy and lazy allocation (#13481)

Author: cpriti-os

storage/grpc_writer.go

@@ -23,6 +23,7 @@ import (
 	"net/http"
 	"net/url"
 	"strings"
+	"sync"
 	"time"
 
 	gapic "cloud.google.com/go/storage/internal/apiv2"
@@ -185,7 +186,8 @@ func (c *grpcStorageClient) OpenWriter(params *openWriterParams, opts ...storage
 		appendGen:             params.appendGen,
 		finalizeOnClose:       params.finalizeOnClose,
 
-		buf:              make([]byte, 0, chunkSize),
+		buf:              nil, // Allocated lazily on first buffered write.
+		chunkSize:        chunkSize,
 		writeQuantum:     writeQuantum,
 		lastSegmentStart: lastSegmentStart,
 		sendableUnits:    sendableUnits,
@@ -267,7 +269,8 @@ type gRPCWriter struct {
 	appendGen             int64
 	finalizeOnClose       bool
 
-	buf []byte
+	buf       []byte
+	chunkSize int
 	// A writeQuantum is the largest quantity of data which can be sent to the
 	// service in a single message.
 	writeQuantum     int
@@ -384,21 +387,26 @@ func (w *gRPCWriter) gatherFirstBuffer() error {
 	for cmd := range w.writesChan {
 		switch v := cmd.(type) {
 		case *gRPCWriterCommandWrite:
-			if len(w.buf)+len(v.p) <= cap(w.buf) {
-				// We have not started sending yet, and we can stage all data without
-				// starting a send. Compare against cap(w.buf) instead of
-				// w.writeQuantum: that way we can perform a oneshot upload for objects
-				// which fit in one chunk, even though we will cut the request into
-				// w.writeQuantum units when we do start sending.
-				origLen := len(w.buf)
-				w.buf = w.buf[:origLen+len(v.p)]
-				copy(w.buf[origLen:], v.p)
-				close(v.done)
-			} else {
-				// Too large. Handle it in writeLoop.
+			// If zero-copy one-shot is requested, OR the payload is larger than the buffer,
+			// bypass buffering entirely and hand off to the writeLoop immediately.
+			if w.forceOneShot || len(w.buf)+len(v.p) > w.chunkSize {
 				w.currentCommand = cmd
 				return nil
 			}
+
+			// Otherwise, lazily allocate and stage the small write (normal buffered path)
+			if w.buf == nil {
+				w.buf = make([]byte, 0, w.chunkSize)
+			}
+			// We have not started sending yet, and we can stage all data without
+			// starting a send. Compare against w.chunkSize instead of
+			// w.writeQuantum: that way we can perform a oneshot upload for objects
+			// which fit in one chunk, even though we will cut the request into
+			// w.writeQuantum units when we do start sending.
+			origLen := len(w.buf)
+			w.buf = w.buf[:origLen+len(v.p)]
+			copy(w.buf[origLen:], v.p)
+			close(v.done)
 			break
 		case *gRPCWriterCommandClose:
 			// If we get here, data (if any) fits in w.buf, so we can force oneshot.
@@ -568,17 +576,33 @@ type gRPCWriterCommand interface {
 }
 
 type gRPCWriterCommandWrite struct {
-	p    []byte
-	done chan struct{}
+	p             []byte
+	done          chan struct{}
+	initialOffset int64
+	hasStarted    bool
+	closeOnce     sync.Once
 }
 
 func (c *gRPCWriterCommandWrite) handle(w *gRPCWriter, cs gRPCWriterCommandHandleChans) error {
 	if len(c.p) == 0 {
 		// No data to write.
-		close(c.done)
+		c.markDone()
+		return nil
+	}
+
+	// Zero-Copy send.
+	if w.forceOneShot {
+		err := c.zeroCopyWrite(w, cs)
+		if err != nil {
+			return err
+		}
+		// If zeroCopyWrite returns without error, the write is done.
 		return nil
 	}
 
+	if w.buf == nil {
+		w.buf = make([]byte, 0, w.chunkSize)
+	}
 	wblen := len(w.buf)
 	allKnownBytes := wblen + len(c.p)
 	fullBufs := allKnownBytes / cap(w.buf)
@@ -605,7 +629,7 @@ func (c *gRPCWriterCommandWrite) handle(w *gRPCWriter, cs gRPCWriterCommandHandl
 			return w.streamSender.err()
 		}
 		w.bufUnsentIdx = int(sentOffset - w.bufBaseOffset)
-		close(c.done)
+		c.markDone()
 		return nil
 	}
 
@@ -698,10 +722,53 @@ func (c *gRPCWriterCommandWrite) handle(w *gRPCWriter, cs gRPCWriterCommandHandl
 	w.buf = w.buf[:len(toCopyIn)]
 	copy(w.buf, toCopyIn)
 	w.bufUnsentIdx = int(sentOffset - w.bufBaseOffset)
-	close(c.done)
+	c.markDone()
+	return nil
+}
+
+func (c *gRPCWriterCommandWrite) zeroCopyWrite(w *gRPCWriter, cs gRPCWriterCommandHandleChans) error {
+	// Pre-emptively get the context channel to avoid closure overhead in the loop.
+	ctxDone := w.preRunCtx.Done()
+
+	// sendBufferToTarget handles the quantum breakdown.
+	newOffset, ok := w.sendBufferToTarget(cs, c.p, w.bufBaseOffset, len(c.p), w.handleCompletion)
+	if !ok {
+		return w.streamSender.err()
+	}
+
+	// Request an ack from the sender goroutine to ensure the buffer has been
+	// dispatched to gRPC and is safe for the user to reuse.
+	if !cs.deliverRequestUnlessCompleted(gRPCBidiWriteRequest{requestAck: true}, w.handleCompletion) {
+		return w.streamSender.err()
+	}
+
+	ackOutstanding := true
+
+	// Wait for server acknowledgement and sender transmissions to enable incremental progress.
+	for ackOutstanding || w.bufBaseOffset < newOffset {
+		select {
+		case completion, ok := <-cs.completions:
+			if !ok {
+				return w.streamSender.err()
+			}
+			w.handleCompletion(completion)
+		case <-cs.requestAcks:
+			ackOutstanding = false
+		case <-ctxDone:
+			return w.preRunCtx.Err()
+		}
+	}
+
+	c.p = nil
+	c.markDone()
 	return nil
 }
 
+// Helper to ensure we don't close done twice and keep the main logic clean.
+func (c *gRPCWriterCommandWrite) markDone() {
+	c.closeOnce.Do(func() { close(c.done) })
+}
+
 type gRPCWriterCommandFlush struct {
 	done chan int64
 }

storage/grpc_writer_test.go

@@ -15,9 +15,12 @@
 package storage
 
 import (
+	"context"
+	"sync"
 	"testing"
 
 	"cloud.google.com/go/storage/internal/apiv2/storagepb"
+	gax "github.com/googleapis/gax-go/v2"
 	"google.golang.org/protobuf/proto"
 )
 
@@ -121,3 +124,173 @@ func TestGetObjectChecksums(t *testing.T) {
 		})
 	}
 }
+func TestGRPCWriter_MemoryAllocationPaths(t *testing.T) {
+	tests := []struct {
+		name         string
+		chunkSize    int
+		dataSize     int
+		forceOneShot bool
+		wantZeroCopy bool
+	}{
+		{
+			name:         "OneShot_ZeroCopy_1MB",
+			chunkSize:    0,
+			dataSize:     1 * 1024 * 1024, // 1 MiB
+			forceOneShot: true,
+			wantZeroCopy: true,
+		},
+		{
+			name:         "OneShot_ZeroCopy_10MB",
+			chunkSize:    0,
+			dataSize:     10 * 1024 * 1024, // 10 MiB
+			forceOneShot: true,
+			wantZeroCopy: true,
+		},
+		{
+			name:         "Resumable_Buffering",
+			chunkSize:    2 * 1024 * 1024,
+			dataSize:     1 * 1024 * 1024, // 1 MiB
+			forceOneShot: false,
+			wantZeroCopy: false,
+		},
+		{
+			name:         "Resumable_ZeroCopy",
+			chunkSize:    1 * 1024 * 1024, // 1 MiB
+			dataSize:     2 * 1024 * 1024, // 2 MiB
+			forceOneShot: false,
+			wantZeroCopy: true,
+		},
+		{
+			name:         "Resumable_Hybrid",
+			chunkSize:    2 * 1024 * 1024, // 2 MiB
+			dataSize:     3 * 1024 * 1024, // 3 MiB
+			forceOneShot: false,
+			wantZeroCopy: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			data := make([]byte, tt.dataSize)
+			data[0] = 1
+			data[tt.dataSize-1] = 2
+			chunkSize := gRPCChunkSize(tt.chunkSize)
+			mockSender := &mockZeroCopySender{}
+			w := &gRPCWriter{
+				buf:           nil, // Allocated lazily on first buffered write.
+				chunkSize:     chunkSize,
+				forceOneShot:  tt.forceOneShot,
+				writeQuantum:  maxPerMessageWriteSize,
+				preRunCtx:     context.Background(),
+				sendableUnits: 10,
+				writesChan:    make(chan gRPCWriterCommand, 1),
+				donec:         make(chan struct{}),
+				streamSender:  mockSender,
+				settings:      &settings{},
+			}
+			w.progress = func(int64) {}
+			w.setObj = func(*ObjectAttrs) {}
+			w.setSize = func(int64) {}
+
+			go func() {
+				w.writeLoop(context.Background())
+				close(w.donec)
+			}()
+
+			if _, err := w.Write(data); err != nil {
+				t.Fatalf("Write failed: %v", err)
+			}
+			if err := w.Close(); err != nil {
+				t.Fatalf("Close failed: %v", err)
+			}
+			mockSender.wg.Wait()
+
+			mockSender.mu.Lock()
+			defer mockSender.mu.Unlock()
+
+			reqs := filterDataRequests(mockSender.requests)
+			if len(reqs) == 0 {
+				t.Fatalf("Expected at least 1 data request, got 0")
+			}
+
+			// Verify memory address logic:
+			// The last byte of the last request buffer should match the last byte of the input data for zero-copy.
+			// For buffering/copying, the pointers must differ.
+			idx := len(reqs) - 1
+			bufIdx := len(reqs[idx].buf)
+			isZeroCopy := &reqs[idx].buf[bufIdx-1] == &data[tt.dataSize-1]
+			if isZeroCopy != tt.wantZeroCopy {
+				if tt.wantZeroCopy && tt.forceOneShot {
+					t.Errorf("One-shot upload bypassed zero-copy path; data was unexpectedly copied")
+				} else if !tt.wantZeroCopy && !tt.forceOneShot {
+					t.Errorf("Resumable upload bypassed buffering path; data was unexpectedly zero-copied")
+				} else if tt.wantZeroCopy && !tt.forceOneShot {
+					t.Errorf("Resumable upload bypassed zero-copy path; data was unexpectedly copied")
+				}
+			}
+		})
+	}
+}
+
+type mockZeroCopySender struct {
+	mu        sync.Mutex
+	requests  []gRPCBidiWriteRequest
+	errResult error
+	wg        sync.WaitGroup // Waits for all async operations to complete.
+}
+
+func (m *mockZeroCopySender) connect(ctx context.Context, cs gRPCBufSenderChans, opts ...gax.CallOption) {
+	m.wg.Add(1)
+	go func() {
+		defer m.wg.Done()
+
+		// Track active flush goroutines to prevent closing the channel prematurely.
+		var completionWg sync.WaitGroup
+
+		defer func() {
+			completionWg.Wait()
+			close(cs.completions)
+		}()
+
+		for req := range cs.requests {
+			m.mu.Lock()
+			m.requests = append(m.requests, req)
+			m.mu.Unlock()
+
+			if req.requestAck {
+				select {
+				case cs.requestAcks <- struct{}{}:
+				case <-ctx.Done():
+					return
+				}
+			}
+
+			if req.flush {
+				completionWg.Add(1)
+				// Send completions asynchronously to avoid blocking the request loop.
+				go func(offset int64) {
+					defer completionWg.Done()
+					select {
+					case cs.completions <- gRPCBidiWriteCompletion{
+						flushOffset: offset,
+					}:
+					case <-ctx.Done():
+					}
+				}(req.offset + int64(len(req.buf)))
+			}
+		}
+	}()
+}
+
+func (m *mockZeroCopySender) err() error { return m.errResult }
+
+// filterDataRequests returns only requests containing data, ignoring protocol overhead.
+func filterDataRequests(reqs []gRPCBidiWriteRequest) []gRPCBidiWriteRequest {
+	var dataReqs []gRPCBidiWriteRequest
+	for _, r := range reqs {
+		if len(r.buf) > 0 {
+			dataReqs = append(dataReqs, r)
+		}
+	}
+	return dataReqs
+}