Update WriteStateBytesStream to not use Seq2, instead iterate over 'messages' that contain both chunk data and any raised grpc errors

Author: SarahFrench

tfprotov6/state_store.go

@@ -57,7 +57,16 @@ type ReadStateBytesStream struct {
 }
 
 type WriteStateBytesStream struct {
-	Chunks iter.Seq2[WriteStateByteChunk, error]
+	Messages iter.Seq[WriteStateBytesStreamMsg]
+}
+
+// WriteStateBytesStreamMsg contains paired data:
+//  1. A chunk of state data, received from Terraform core to be persisted.
+//  2. Any gRPC-related errors the provider server encountered when
+//     receiving data from Terraform core.
+type WriteStateBytesStreamMsg struct {
+	Chunk WriteStateByteChunk
+	Err   error
 }
 
 type WriteStateBytesResponse struct {

tfprotov6/tf6server/server.go

@@ -1644,8 +1644,7 @@ func (s *server) WriteStateBytes(srv grpc.ClientStreamingServer[tfplugin6.WriteS
 		return err
 	}
 
-	// Trying an aproach using iter.Seq2
-	iterator := func(yield func(tfprotov6.WriteStateByteChunk, error) bool) {
+	iterator := func(yield func(tfprotov6.WriteStateBytesStreamMsg) bool) {
 		for {
 			chunk, err := srv.Recv()
 			if err == io.EOF {
@@ -1658,22 +1657,26 @@ func (s *server) WriteStateBytes(srv grpc.ClientStreamingServer[tfplugin6.WriteS
 				))
 			}
 
-			ok := yield(tfprotov6.WriteStateByteChunk{
-				Bytes:       chunk.Bytes,
-				TotalLength: chunk.TotalLength,
-				Range: tfprotov6.StateByteRange{
-					Start: chunk.Range.Start,
-					End:   chunk.Range.End,
-				},
-			}, err)
+			ok := yield(
+				tfprotov6.WriteStateBytesStreamMsg{
+					Chunk: tfprotov6.WriteStateByteChunk{
+						Bytes:       chunk.Bytes,
+						TotalLength: chunk.TotalLength,
+						Range: tfprotov6.StateByteRange{
+							Start: chunk.Range.Start,
+							End:   chunk.Range.End,
+						},
+					},
+					Err: err, // If this isn't nil, it'll be a gRPC error from Recv
+				})
 			if !ok {
 				return
 			}
 		}
 	}
 
 	resp, err := server.WriteStateBytes(ctx, &tfprotov6.WriteStateBytesStream{
-		Chunks: iterator,
+		Messages: iterator,
 	})
 	if err != nil {
 		return err