fix(storage): fix race condition during retries in gRPC writer (#14649)

krishnamd-jkp · krishnamd-jkp · commit 04b6c635c09d · 2026-06-03T06:31:11.000Z
diff --git a/storage/grpc_writer.go b/storage/grpc_writer.go
@@ -443,6 +443,31 @@ func (w *gRPCWriter) writeLoop(ctx context.Context) error {
 	defer cancel()
 	w.streamSender.connect(ctx, bscs, w.settings.gax...)
 
+	// Drain any initial completions (like QueryWriteStatus results).
+Loop:
+	for {
+		select {
+		case c, ok := <-completions:
+			if !ok {
+				return w.streamSender.err()
+			}
+			w.handleCompletion(c)
+		default:
+			break Loop
+		}
+	}
+
+	if w.bufFlushedIdx > 0 {
+		copy(w.buf, w.buf[w.bufFlushedIdx:])
+		w.buf = w.buf[:len(w.buf)-w.bufFlushedIdx]
+		w.bufBaseOffset += int64(w.bufFlushedIdx)
+		w.bufUnsentIdx -= w.bufFlushedIdx
+		if w.bufUnsentIdx < 0 {
+			w.bufUnsentIdx = 0
+		}
+		w.bufFlushedIdx = -1
+	}
+
 	// Send any full quantum in w.buf, possibly including a flush
 	if err := w.withCommandRetryDeadline(func() error {
 		sentOffset, ok := w.sendBufferToTarget(chcs, w.buf, w.bufBaseOffset, cap(w.buf),
@@ -591,6 +616,26 @@ func (c *gRPCWriterCommandWrite) handle(w *gRPCWriter, cs gRPCWriterCommandHandl
 		return nil
 	}
 
+	if !c.hasStarted {
+		c.initialOffset = w.bufBaseOffset + int64(len(w.buf))
+		c.hasStarted = true
+	} else {
+		// Retrying this command; check if server has persisted some bytes of this command's payload.
+		bytesPersisted := w.bufBaseOffset - c.initialOffset
+		if bytesPersisted > 0 {
+			if int64(len(c.p)) < bytesPersisted {
+				bytesPersisted = int64(len(c.p))
+			}
+			c.p = c.p[bytesPersisted:]
+			c.initialOffset = w.bufBaseOffset
+
+			if len(c.p) == 0 {
+				c.markDone()
+				return nil
+			}
+		}
+	}
+
 	// Zero-Copy send.
 	if w.forceOneShot {
 		err := c.zeroCopyWrite(w, cs)
@@ -650,6 +695,7 @@ func (c *gRPCWriterCommandWrite) handle(w *gRPCWriter, cs gRPCWriterCommandHandl
 	w.buf = w.buf[:wblen+toNextWriteQuantum]
 	copied := copy(w.buf[wblen:], c.p)
 	c.p = c.p[copied:]
+	c.initialOffset += int64(copied)
 	firstFullBufFromCmd := cap(w.buf) - len(w.buf)
 
 	sending := w.buf[w.bufUnsentIdx:]
@@ -674,6 +720,7 @@ func (c *gRPCWriterCommandWrite) handle(w *gRPCWriter, cs gRPCWriterCommandHandl
 			trim = len(c.p)
 		}
 		c.p = c.p[trim:]
+		c.initialOffset += int64(trim)
 		cmdBaseOffset = bufTail
 	}
 	offset := cmdBaseOffset
diff --git a/storage/internal/test/conformance/retry_tests.json b/storage/internal/test/conformance/retry_tests.json
@@ -258,6 +258,19 @@
         },
         {
           "instructions": ["return-503-after-8192K", "return-408"]
+        },
+        {
+          "instructions": [ "return-503-after-4080K"]
+        },
+        {
+          "instructions": [
+            "return-503-after-1000K",
+            "return-503-after-1024K",
+            "return-503-after-1025K",
+            "return-503-after-4000K",
+            "return-503-after-4096K",
+            "return-503-after-4097K"
+          ]
         }
       ],
       "methods": [
diff --git a/storage/retry_conformance_test.go b/storage/retry_conformance_test.go
@@ -43,7 +43,8 @@ import (
 const (
 	projectID           = "my-project-id"
 	serviceAccountEmail = "my-sevice-account@my-project-id.iam.gserviceaccount.com"
-	MiB                 = 1 << 10 << 10
+	MiB                 = 1 << 20
+	KiB                 = 1 << 10
 )
 
 var (
@@ -782,6 +783,56 @@ var methods = map[string][]retryFunc{
 			}
 			return nil
 		},
+		func(ctx context.Context, c *Client, fs *resources, preconditions bool) error {
+			obj := c.Bucket(fs.bucket.Name).Object(objectIDs.New())
+			if preconditions {
+				obj = obj.If(Conditions{DoesNotExist: true})
+			}
+			w := obj.NewWriter(ctx)
+			w.ChunkSize = 4 * MiB
+			data := generateRandomBytes(32 * MiB)
+			writeSize := 2 * MiB
+			// Write 2MiB at a time to use buffers and not user's buffer.
+			for i := 0; i < 16; i++ {
+				if _, err := w.Write(data[i*writeSize : (i+1)*writeSize]); err != nil {
+					return fmt.Errorf("writing object: %v", err)
+				}
+			}
+			if err := w.Close(); err != nil {
+				return fmt.Errorf("closing object: %v", err)
+			}
+			return nil
+		},
+		func(ctx context.Context, c *Client, fs *resources, preconditions bool) error {
+			obj := c.Bucket(fs.bucket.Name).Object(objectIDs.New())
+			if preconditions {
+				obj = obj.If(Conditions{DoesNotExist: true})
+			}
+			w := obj.NewWriter(ctx)
+			w.ChunkSize = 4 * MiB
+			data := generateRandomBytes(32 * MiB)
+			totalWriteSize := 0
+			times := 0
+			for totalWriteSize < len(data) {
+				times++
+				// Write smaller messages and larger messages to exercise different paths.
+				writeSize := 5 * KiB
+				if times%2 == 0 {
+					writeSize = 6 * MiB
+				}
+				if totalWriteSize+writeSize > len(data) {
+					writeSize = len(data) - totalWriteSize
+				}
+				if _, err := w.Write(data[totalWriteSize : totalWriteSize+writeSize]); err != nil {
+					return fmt.Errorf("writing object: %v", err)
+				}
+				totalWriteSize += writeSize
+			}
+			if err := w.Close(); err != nil {
+				return fmt.Errorf("closing object: %v", err)
+			}
+			return nil
+		},
 	},
 	"storage.appendable.upload": {
 		func(ctx context.Context, c *Client, fs *resources, preconditions bool) error {
