fix(storage): Fix takeover response handling. (#13239)

cjc25 · web-flow · commit 26d75bc08e24 · 2025-10-28T03:36:26.000Z
A takeover response can return a handle. In that case, a retry should present that handle. Otherwise it will attempt a takeover when it retries. After a new takeover, the client can't tell if a partial write after the last known flush offset was performed by itself, or by some other writer. Therefore the client can't tell if it's safe to retry from the last known flush offset, and the server will return an error if it repeats offsets.

The write handle allows the client to present enough information to the server to correctly identify replays on the same write session, avoiding an error when the client repeats bytes that are already durable.
diff --git a/storage/grpc_writer.go b/storage/grpc_writer.go
@@ -1169,8 +1169,8 @@ func (s *gRPCAppendBidiWriteBufferSender) handleStream(stream storagepb.Storage_
 
 type gRPCAppendTakeoverBidiWriteBufferSender struct {
 	gRPCAppendBidiWriteBufferSender
-	takeoverReported  bool
-	setTakeoverOffset func(int64)
+	takeoverReported         bool
+	handleTakeoverCompletion func(gRPCBidiWriteCompletion)
 }
 
 func writeObjectSpecAsAppendObjectSpec(s *storagepb.WriteObjectSpec, gen int64) *storagepb.AppendObjectSpec {
@@ -1197,8 +1197,11 @@ func (w *gRPCWriter) newGRPCAppendTakeoverWriteBufferSender() *gRPCAppendTakeove
 			objectChecksums: toProtoChecksums(w.sendCRC32C, w.attrs),
 			finalizeOnClose: w.finalizeOnClose,
 		},
-		takeoverReported:  false,
-		setTakeoverOffset: w.setTakeoverOffset,
+		takeoverReported: false,
+		handleTakeoverCompletion: func(c gRPCBidiWriteCompletion) {
+			w.handleCompletion(c)
+			w.setTakeoverOffset(c.flushOffset)
+		},
 	}
 }
 
@@ -1238,9 +1241,9 @@ func (s *gRPCAppendTakeoverBidiWriteBufferSender) connect(ctx context.Context, c
 			return
 		}
 
-		s.setTakeoverOffset(c.flushOffset)
+		s.maybeUpdateFirstMessage(resp)
 		s.takeoverReported = true
-		cs.completions <- *c
+		s.handleTakeoverCompletion(*c)
 	}
 
 	go s.handleStream(stream, cs, firstSend)
diff --git a/storage/retry_conformance_test.go b/storage/retry_conformance_test.go
@@ -731,6 +731,79 @@ var methods = map[string][]retryFunc{
 				return fmt.Errorf("Reader.Read: %w", err)
 			}
 
+			gotMd5 := md5.Sum(content)
+			expectedMd5 := md5.Sum(toWrite)
+			if d := cmp.Diff(gotMd5, expectedMd5); d != "" {
+				return fmt.Errorf("content mismatch, got %v bytes (md5: %v), want %v bytes (md5: %v)",
+					len(content), gotMd5, len(toWrite), expectedMd5)
+			}
+			return nil
+		},
+		// Appendable upload using a takeover.
+		func(ctx context.Context, c *Client, fs *resources, preconditions bool) error {
+			bucketName := fmt.Sprintf("%s-appendable", bucketIDs.New())
+			b := c.Bucket(bucketName)
+			if err := b.Create(ctx, projectID, nil); err != nil {
+				return err
+			}
+			defer b.Delete(ctx)
+
+			obj := b.Object(objectIDs.New())
+			if preconditions {
+				obj = obj.If(Conditions{DoesNotExist: true})
+			}
+
+			// Force multiple messages per chunk, and multiple chunks in the object.
+			chunkSize := 2 * maxPerMessageWriteSize
+			toWrite := generateRandomBytes(chunkSize * 3)
+
+			objW := obj.NewWriter(ctx)
+			objW.Append = true
+			objW.ChunkSize = chunkSize
+			if _, err := objW.Write(toWrite[0:maxPerMessageWriteSize]); err != nil {
+				return fmt.Errorf("Writer.Write: %w", err)
+			}
+			// Close this writer, which will create the appendable unfinalized object
+			// (there was not enough in Write to trigger a send).
+			if err := objW.Close(); err != nil {
+				return fmt.Errorf("Creation Writer.Close: %v", err)
+			}
+
+			generation := int64(0)
+			if preconditions {
+				generation = objW.Attrs().Generation
+			}
+			objT := b.Object(obj.ObjectName()).Generation(generation)
+			w, l, err := objT.NewWriterFromAppendableObject(ctx, &AppendableWriterOpts{ChunkSize: chunkSize})
+			if err != nil {
+				return fmt.Errorf("NewWriterFromAppendableObject: %v", err)
+			}
+			if l != int64(maxPerMessageWriteSize) {
+				return fmt.Errorf("NewWriterFromAppendableObject unexpected len: got %v, want %v", l, maxPerMessageWriteSize)
+			}
+
+			if _, err := w.Write(toWrite[maxPerMessageWriteSize:]); err != nil {
+				return fmt.Errorf("Writer.Write: %v", err)
+			}
+			if err := w.Close(); err != nil {
+				return fmt.Errorf("Writer.Close: %v", err)
+			}
+
+			if w.Attrs() == nil {
+				return fmt.Errorf("Writer.Attrs: expected attrs for written object, got nil")
+			}
+
+			// Don't reuse obj, in case preconditions were set on the write request.
+			r, err := b.Object(obj.ObjectName()).NewReader(ctx)
+			defer r.Close()
+			if err != nil {
+				return fmt.Errorf("obj.NewReader: %v", err)
+			}
+			content, err := io.ReadAll(r)
+			if err != nil {
+				return fmt.Errorf("Reader.Read: %v", err)
+			}
+
 			gotMd5 := md5.Sum(content)
 			expectedMd5 := md5.Sum(toWrite)
 			if d := cmp.Diff(gotMd5, expectedMd5); d != "" {
