fix(storage): fix unfinalized write size (#12016)

tritone · web-flow · commit 6217f8fd3cd8 · 2025-04-21T10:04:06.000-04:00
* fix(storage): fix unfinalized write size

This makes Writer.Attrs().Size true to the last persisted
size value received from GCS.

* fix hang on takeover
diff --git a/storage/grpc_client.go b/storage/grpc_client.go
@@ -1782,7 +1782,7 @@ func (c *grpcStorageClient) OpenWriter(params *openWriterParams, opts ...storage
 		params.setTakeoverOffset(wbs.takeoverOffset)
 		offset = wbs.takeoverOffset
 		gw.streamSender = wbs
-		o = wbs.takeoverObj
+		o = wbs.objResource
 	}
 
 	// This function reads the data sent to the pipe and sends sets of messages
diff --git a/storage/grpc_writer.go b/storage/grpc_writer.go
@@ -46,7 +46,7 @@ type gRPCAppendBidiWriteBufferSender struct {
 	progress          func(int64)
 	flushOffset       int64
 	takeoverOffset    int64
-	takeoverObj       *storagepb.Object // Object returned by takeover stream reopening.
+	objResource       *storagepb.Object // Captures received obj to set w.Attrs.
 
 	// Fields used to report responses from the receive side of the stream
 	// recvs is closed when the current recv goroutine is complete. recvErr is set
@@ -96,14 +96,14 @@ func (w *gRPCWriter) newGRPCAppendTakeoverWriteBufferSender(ctx context.Context)
 	if err := s.connect(ctx); err != nil {
 		return nil, fmt.Errorf("storage: opening appendable write stream: %w", err)
 	}
-	_, err := s.sendOnConnectedStream(nil, 0, true, false, true)
+	_, err := s.sendOnConnectedStream(nil, 0, false, false, true)
 	if err != nil {
 		return nil, err
 	}
 	firstResp := <-s.recvs
 	// Object resource is returned in the first response on takeover, so capture
 	// this now.
-	s.takeoverObj = firstResp.GetResource()
+	s.objResource = firstResp.GetResource()
 	s.takeoverOffset = firstResp.GetResource().GetSize()
 	return s, nil
 }
@@ -298,8 +298,8 @@ func (s *gRPCAppendBidiWriteBufferSender) sendOnConnectedStream(buf []byte, offs
 			// When closing the stream, update the object resource to reflect
 			// the persisted size. We get a new object from the stream if
 			// the object was finalized, but not if it's unfinalized.
-			if s.takeoverObj != nil && resp.GetPersistedSize() > 0 {
-				s.takeoverObj.Size = resp.GetPersistedSize()
+			if s.objResource != nil && resp.GetPersistedSize() > 0 {
+				s.objResource.Size = resp.GetPersistedSize()
 			}
 		}
 		if s.recvErr != io.EOF {
@@ -358,6 +358,9 @@ func (s *gRPCAppendBidiWriteBufferSender) sendBuffer(ctx context.Context, buf []
 		}
 
 		obj, err = s.sendOnConnectedStream(buf, offset, flush, finishWrite, sendFirstMessage)
+		if obj != nil {
+			s.objResource = obj
+		}
 		if err == nil {
 			return
 		}
diff --git a/storage/integration_test.go b/storage/integration_test.go
@@ -3280,6 +3280,10 @@ func TestIntegration_WriterAppend(t *testing.T) {
 
 				// Write remaining data.
 				h.mustWrite(w, content)
+				// Check that local Writer.Attrs() is populated as expected.
+				if w.Attrs() == nil || w.Attrs().Size != int64(len(tc.content)) {
+					t.Errorf("Writer.Attrs(): got %+v, expected size = %v", w.Attrs().Size, int64(len(tc.content)))
+				}
 
 				// Download content again and validate.
 				// Disabled due to b/395944605; unskip after this is resolved.

Original file line number	Diff line number	Diff line change
`@@ -1782,7 +1782,7 @@ func (c grpcStorageClient) OpenWriter(params openWriterParams, opts ...storage`
`1782`	`1782`	`params.setTakeoverOffset(wbs.takeoverOffset)`
`1783`	`1783`	`offset = wbs.takeoverOffset`
`1784`	`1784`	`gw.streamSender = wbs`
`1785`		`- o = wbs.takeoverObj`
	`1785`	`+ o = wbs.objResource`
`1786`	`1786`	`}`
`1787`	`1787`
`1788`	`1788`	`// This function reads the data sent to the pipe and sends sets of messages`