Merge pull request #1676 from ydb-platform/1660-big-topic-message

Fixed hangup when try to send batch of messages with size more, then …

Author: asmyasnikov

CHANGELOG.md

@@ -1,3 +1,5 @@
+* Fixed hangup when try to send batch of messages with size more, then grpc limits from topic writer internals
+
 ## v3.101.2
 * Added a new metric `ydb_go_sdk_ydb_info` with the current version of the SDK
 

config/config.go

@@ -20,18 +20,19 @@ import (
 type Config struct {
 	config.Common
 
-	trace          *trace.Driver
-	dialTimeout    time.Duration
-	connectionTTL  time.Duration
-	balancerConfig *balancerConfig.Config
-	secure         bool
-	endpoint       string
-	database       string
-	metaOptions    []meta.Option
-	grpcOptions    []grpc.DialOption
-	credentials    credentials.Credentials
-	tlsConfig      *tls.Config
-	meta           *meta.Meta
+	trace              *trace.Driver
+	dialTimeout        time.Duration
+	connectionTTL      time.Duration
+	balancerConfig     *balancerConfig.Config
+	secure             bool
+	endpoint           string
+	database           string
+	metaOptions        []meta.Option
+	grpcOptions        []grpc.DialOption
+	grpcMaxMessageSize int
+	credentials        credentials.Credentials
+	tlsConfig          *tls.Config
+	meta               *meta.Meta
 
 	excludeGRPCCodesForPessimization []grpcCodes.Code
 }
@@ -53,6 +54,11 @@ func (c *Config) GrpcDialOptions() []grpc.DialOption {
 	)
 }
 
+// GrpcMaxMessageSize return client settings for max grpc message size
+func (c *Config) GrpcMaxMessageSize() int {
+	return c.grpcMaxMessageSize
+}
+
 // Meta reports meta information about database connection
 func (c *Config) Meta() *meta.Meta {
 	return c.meta
@@ -245,6 +251,19 @@ func WithDialTimeout(timeout time.Duration) Option {
 	}
 }
 
+func WithGrpcMaxMessageSize(sizeBytes int) Option {
+	return func(c *Config) {
+		c.grpcMaxMessageSize = sizeBytes
+		c.grpcOptions = append(c.grpcOptions,
+			// limit size of outgoing and incoming packages
+			grpc.WithDefaultCallOptions(
+				grpc.MaxCallRecvMsgSize(c.grpcMaxMessageSize),
+				grpc.MaxCallSendMsgSize(c.grpcMaxMessageSize),
+			),
+		)
+	}
+}
+
 func WithBalancer(balancer *balancerConfig.Config) Option {
 	return func(c *Config) {
 		c.balancerConfig = balancer

config/defaults.go

@@ -107,9 +107,10 @@ func defaultConfig() (c *Config) {
 		credentials: credentials.NewAnonymousCredentials(
 			credentials.WithSourceInfo(stack.Record(0)),
 		),
-		balancerConfig: balancers.Default(),
-		tlsConfig:      defaultTLSConfig(),
-		dialTimeout:    DefaultDialTimeout,
-		trace:          &trace.Driver{},
+		balancerConfig:     balancers.Default(),
+		tlsConfig:          defaultTLSConfig(),
+		dialTimeout:        DefaultDialTimeout,
+		trace:              &trace.Driver{},
+		grpcMaxMessageSize: DefaultGRPCMsgSize,
 	}
 }

driver.go

@@ -34,6 +34,7 @@ import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/stack"
 	internalTable "github.com/ydb-platform/ydb-go-sdk/v3/internal/table"
 	tableConfig "github.com/ydb-platform/ydb-go-sdk/v3/internal/table/config"
+	internalTopic "github.com/ydb-platform/ydb-go-sdk/v3/internal/topic"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/topic/topicclientinternal"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xcontext"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
@@ -588,6 +589,7 @@ func (d *Driver) connect(ctx context.Context) (err error) {
 				[]topicoptions.TopicOption{
 					topicoptions.WithOperationTimeout(d.config.OperationTimeout()),
 					topicoptions.WithOperationCancelAfter(d.config.OperationCancelAfter()),
+					internalTopic.WithGrpcMessageSize(d.config.GrpcMaxMessageSize()),
 				},
 				d.topicOptions...,
 			)...,

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,55 @@ 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) FillCache() *WriteRequest {
+	r.Size()
+
+	return r
+}
+
+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 +237,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

@@ -1,11 +1,40 @@
 package topic
 
 import (
+	"time"
+
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/config"
 	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
 )
 
 type Config struct {
 	config.Common
-	Trace *trace.Topic
+	Trace              *trace.Topic
+	MaxGrpcMessageSize int
+}
+
+type Option func(c *Config)
+
+func PublicWithTrace(trace trace.Topic, opts ...trace.TopicComposeOption) Option { //nolint:gocritic
+	return func(c *Config) {
+		c.Trace = c.Trace.Compose(&trace, opts...)
+	}
+}
+
+func PublicWithOperationTimeout(operationTimeout time.Duration) Option {
+	return func(c *Config) {
+		config.SetOperationTimeout(&c.Common, operationTimeout)
+	}
+}
+
+func PublicWithOperationCancelAfter(operationCancelAfter time.Duration) Option {
+	return func(c *Config) {
+		config.SetOperationCancelAfter(&c.Common, operationCancelAfter)
+	}
+}
+
+func WithGrpcMessageSize(sizeBytes int) Option {
+	return func(c *Config) {
+		c.MaxGrpcMessageSize = sizeBytes
+	}
 }