Fixed hangup when try to send batch of messages with size more, then grpc limits from topic writer internals

Author: rekby

CHANGELOG.md

@@ -1,5 +1,4 @@
-## v3.101.0
-* Added `table.Client.ReadRows` method with internal retries
+* Fixed hangup when try to send batch of messages with size more, then grpc limits from topic writer internals
 
 ## v3.100.3
 * Fixed bug with concurrent rewrites source slice of `grpc.DialOption` on dial step

internal/config/config.go

@@ -7,6 +7,14 @@ import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
 )
 
+// InitDefaultGRPCMessageSize used for init value only
+// don't use it in any place, other than initial get values for addr and config
+var InitDefaultGRPCMessageSize = 64 * 1024 * 1024 // 64MB
+
+// DefaultGRPCMessageSizeAddr pointer to public DefaultGRPCMessageSize value
+// the scheme needs for prevent import public package from internal (and prevent import cycle)
+var DefaultGRPCMessageSizeAddr = &InitDefaultGRPCMessageSize
+
 var defaultRetryBudget = budget.Limited(-1)
 
 type Common struct {

internal/grpcwrapper/rawtopic/rawtopicwriter/messages.go

@@ -8,6 +8,7 @@ import (
 	"time"
 
 	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Topic"
+	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/timestamppb"
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/grpcwrapper/rawtopic/rawtopiccommon"
@@ -169,16 +170,49 @@ func (r *WriteRequest) toProto() (p *Ydb_Topic.StreamWriteMessage_FromClient_Wri
 	return res, nil
 }
 
+var writeRequestClientMessageSize = proto.Size(&Ydb_Topic.StreamWriteMessage_FromClient{
+	ClientMessage: &Ydb_Topic.StreamWriteMessage_FromClient_WriteRequest{},
+})
+
+func (r *WriteRequest) Size() int {
+	if mess, err := r.toProto(); err == nil {
+		size := proto.Size(mess.WriteRequest) + writeRequestClientMessageSize
+
+		return size
+	}
+
+	return 0
+}
+
+func (r *WriteRequest) Cut(count int) (head *WriteRequest, rest *WriteRequest) {
+	if count >= len(r.Messages) {
+		return r, nil
+	}
+
+	rest = &WriteRequest{}
+	*rest = *r
+	r.Messages, rest.Messages = r.Messages[:count], r.Messages[count:]
+
+	return r, rest
+}
+
 type MessageData struct {
 	SeqNo            int64
 	CreatedAt        time.Time
 	UncompressedSize int64
 	Partitioning     Partitioning
 	MetadataItems    []rawtopiccommon.MetadataItem
 	Data             []byte
+
+	size  int
+	proto *Ydb_Topic.StreamWriteMessage_WriteRequest_MessageData
 }
 
 func (d *MessageData) ToProto() (*Ydb_Topic.StreamWriteMessage_WriteRequest_MessageData, error) {
+	if d.proto != nil {
+		return d.proto, nil
+	}
+
 	res := &Ydb_Topic.StreamWriteMessage_WriteRequest_MessageData{
 		SeqNo:            d.SeqNo,
 		CreatedAt:        timestamppb.New(d.CreatedAt),
@@ -197,7 +231,19 @@ func (d *MessageData) ToProto() (*Ydb_Topic.StreamWriteMessage_WriteRequest_Mess
 		})
 	}
 
-	return res, nil
+	d.proto = res
+
+	return d.proto, nil
+}
+
+func (d *MessageData) ProtoWireSizeBytes() int {
+	if d.size == 0 {
+		if p, err := d.ToProto(); err == nil {
+			d.size = proto.Size(p)
+		}
+	}
+
+	return d.size
 }
 
 type WriteResult struct {

internal/grpcwrapper/rawtopic/rawtopicwriter/streamwriter.go

@@ -8,8 +8,6 @@ import (
 	"sync/atomic"
 
 	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Topic"
-	"google.golang.org/grpc/codes"
-	grpcStatus "google.golang.org/grpc/status"
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/grpcwrapper/rawtopic/rawtopiccommon"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/grpcwrapper/rawydb"
@@ -127,7 +125,7 @@ func (w *StreamWriter) Send(rawMsg ClientMessage) (err error) {
 			return writeErr
 		}
 
-		return sendWriteRequest(w.Stream.Send, writeReqProto)
+		return w.Stream.Send(&Ydb_Topic.StreamWriteMessage_FromClient{ClientMessage: writeReqProto})
 	case *UpdateTokenRequest:
 		protoMsg.ClientMessage = &Ydb_Topic.StreamWriteMessage_FromClient_UpdateTokenRequest{
 			UpdateTokenRequest: v.ToProto(),
@@ -149,8 +147,6 @@ func (w *StreamWriter) Send(rawMsg ClientMessage) (err error) {
 	return nil
 }
 
-type sendFunc func(req *Ydb_Topic.StreamWriteMessage_FromClient) error
-
 func (w *StreamWriter) CloseSend() error {
 	w.sendCloseMtx.Lock()
 	defer w.sendCloseMtx.Unlock()
@@ -175,40 +171,3 @@ type ServerMessage interface {
 type serverMessageImpl struct{}
 
 func (*serverMessageImpl) isServerMessage() {}
-
-func sendWriteRequest(send sendFunc, req *Ydb_Topic.StreamWriteMessage_FromClient_WriteRequest) error {
-	sendErr := send(&Ydb_Topic.StreamWriteMessage_FromClient{
-		ClientMessage: req,
-	})
-
-	if sendErr == nil {
-		return nil
-	}
-
-	grpcStatus, ok := grpcStatus.FromError(sendErr)
-	if !ok {
-		return sendErr
-	}
-
-	grpcMessages := req.WriteRequest.GetMessages()
-	if grpcStatus.Code() != codes.ResourceExhausted || len(grpcMessages) < 2 {
-		return sendErr
-	}
-
-	//nolint:gomnd
-	splitIndex := len(grpcMessages) / 2
-	firstMessages, lastMessages := grpcMessages[:splitIndex], grpcMessages[splitIndex:]
-	defer func() {
-		req.WriteRequest.Messages = grpcMessages
-	}()
-
-	req.WriteRequest.Messages = firstMessages
-	err := sendWriteRequest(send, req)
-	if err != nil {
-		return err
-	}
-
-	req.WriteRequest.Messages = lastMessages
-
-	return sendWriteRequest(send, req)
-}

internal/grpcwrapper/rawtopic/rawtopicwriter/streamwriter_test.go

@@ -17,6 +17,7 @@ type WritersCommonConfig struct {
 	writerMeta          map[string]string
 	defaultPartitioning rawtopicwriter.Partitioning
 	compressorCount     int
+	maxBytesPerMessage  int
 
 	Tracer             *trace.Topic
 	cred               credentials.Credentials

internal/topic/topicwriterinternal/writer_config.go

@@ -45,6 +45,12 @@ func WithCompressorCount(num int) PublicWriterOption {
 	}
 }
 
+func WithMaxGrpcMessageBytes(num int) PublicWriterOption {
+	return func(cfg *WriterReconnectorConfig) {
+		cfg.maxBytesPerMessage = num
+	}
+}
+
 func WithTokenUpdateInterval(interval time.Duration) PublicWriterOption {
 	return func(cfg *WriterReconnectorConfig) {
 		cfg.credUpdateInterval = interval

internal/topic/topicwriterinternal/writer_options.go

@@ -83,6 +83,7 @@ func NewWriterReconnectorConfig(options ...PublicWriterOption) WriterReconnector
 			clock:              clockwork.NewRealClock(),
 			compressorCount:    runtime.NumCPU(),
 			Tracer:             &trace.Topic{},
+			maxBytesPerMessage: *config.DefaultGRPCMessageSizeAddr,
 		},
 		AutoSetSeqNo:       true,
 		AutoSetCreatedTime: true,
@@ -207,7 +208,7 @@ func (w *WriterReconnector) fillFields(messages []messageWithDataContent) error
 
 func (w *WriterReconnector) start() {
 	name := fmt.Sprintf("writer %q", w.cfg.topic)
-	w.background.Start(name+", sendloop", w.connectionLoop)
+	w.background.Start(name+", connectionLoop", w.connectionLoop)
 }
 
 func (w *WriterReconnector) Write(ctx context.Context, messages []PublicMessage) (resErr error) {
@@ -643,27 +644,6 @@ func (w *WriterReconnector) GetSessionID() (sessionID string) {
 	return sessionID
 }
 
-func sendMessagesToStream(
-	stream RawTopicWriterStream,
-	targetCodec rawtopiccommon.Codec,
-	messages []messageWithDataContent,
-) error {
-	if len(messages) == 0 {
-		return nil
-	}
-
-	request, err := createWriteRequest(messages, targetCodec)
-	if err != nil {
-		return err
-	}
-	err = stream.Send(&request)
-	if err != nil {
-		return xerrors.WithStackTrace(fmt.Errorf("ydb: failed send write request: %w", err))
-	}
-
-	return nil
-}
-
 func allMessagesHasSameBufCodec(messages []messageWithDataContent) bool {
 	if len(messages) <= 1 {
 		return true
@@ -699,15 +679,17 @@ func splitMessagesByBufCodec(messages []messageWithDataContent) (res [][]message
 }
 
 func createWriteRequest(messages []messageWithDataContent, targetCodec rawtopiccommon.Codec) (
-	res rawtopicwriter.WriteRequest,
+	res *rawtopicwriter.WriteRequest,
 	err error,
 ) {
 	for i := 1; i < len(messages); i++ {
 		if messages[i-1].tx != messages[i].tx {
-			return res, xerrors.WithStackTrace(errDiffetentTransactions)
+			return nil, xerrors.WithStackTrace(errDiffetentTransactions)
 		}
 	}
 
+	res = &rawtopicwriter.WriteRequest{}
+
 	if len(messages) > 0 && messages[0].tx != nil {
 		res.Tx.ID = messages[0].tx.ID()
 		res.Tx.Session = messages[0].tx.SessionID()
@@ -718,7 +700,7 @@ func createWriteRequest(messages []messageWithDataContent, targetCodec rawtopicc
 	for i := range messages {
 		res.Messages[i], err = createRawMessageData(res.Codec, &messages[i])
 		if err != nil {
-			return res, err
+			return nil, err
 		}
 	}