github.com/ydb-platform/ydb-go-sdk/v3@v3.89.2/internal/topic/topicreadercommon/message.go (about)

     1  package topicreadercommon
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"errors"
     7  	"fmt"
     8  	"time"
     9  
    10  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/empty"
    11  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/grpcwrapper/rawtopic/rawtopiccommon"
    12  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
    13  )
    14  
    15  var errMessageWasReadEarly = xerrors.Wrap(errors.New("ydb: message was read early"))
    16  
    17  // PublicMessage is representation of topic message
    18  type PublicMessage struct {
    19  	empty.DoNotCopy
    20  
    21  	SeqNo                int64
    22  	CreatedAt            time.Time
    23  	MessageGroupID       string
    24  	WriteSessionMetadata map[string]string
    25  	Offset               int64
    26  	WrittenAt            time.Time
    27  	ProducerID           string
    28  	Metadata             map[string][]byte // Metadata, nil if no metadata
    29  
    30  	commitRange        CommitRange
    31  	data               oneTimeReader
    32  	rawDataLen         int
    33  	bufferBytesAccount int
    34  	UncompressedSize   int // as sent by sender, server/sdk doesn't check the field. It may be empty or wrong.
    35  	dataConsumed       bool
    36  }
    37  
    38  func (m *PublicMessage) Context() context.Context {
    39  	return m.commitRange.session().Context()
    40  }
    41  
    42  func (m *PublicMessage) Topic() string {
    43  	return m.commitRange.session().Topic
    44  }
    45  
    46  func (m *PublicMessage) PartitionID() int64 {
    47  	return m.commitRange.session().PartitionID
    48  }
    49  
    50  func (m *PublicMessage) getCommitRange() PublicCommitRange {
    51  	return m.commitRange.getCommitRange()
    52  }
    53  
    54  // UnmarshalTo can call most once per message, it read all data from internal reader and
    55  // call PublicMessageContentUnmarshaler.UnmarshalYDBTopicMessage with uncompressed content
    56  func (m *PublicMessage) UnmarshalTo(dst PublicMessageContentUnmarshaler) error {
    57  	if m.dataConsumed {
    58  		return xerrors.WithStackTrace(errMessageWasReadEarly)
    59  	}
    60  
    61  	m.dataConsumed = true
    62  
    63  	return callbackOnReaderContent(globalReadMessagePool, m, m.UncompressedSize, dst)
    64  }
    65  
    66  // Read implements io.Reader
    67  // Read uncompressed message content
    68  // return topicreader.UnexpectedCodec if message compressed with unknown codec
    69  //
    70  // Content of the message released from the memory after first read error
    71  // including io.EOF.
    72  func (m *PublicMessage) Read(p []byte) (n int, err error) {
    73  	m.dataConsumed = true
    74  
    75  	return m.data.Read(p)
    76  }
    77  
    78  // PublicMessageContentUnmarshaler is interface for unmarshal message content
    79  type PublicMessageContentUnmarshaler interface {
    80  	// UnmarshalYDBTopicMessage MUST NOT use data after return.
    81  	// If you need content after return from Consume - copy data content to
    82  	// own slice with copy(dst, data)
    83  	UnmarshalYDBTopicMessage(data []byte) error
    84  }
    85  
    86  func createReader(decoders DecoderMap, codec rawtopiccommon.Codec, rawBytes []byte) oneTimeReader {
    87  	reader, err := decoders.Decode(codec, bytes.NewReader(rawBytes))
    88  	if err != nil {
    89  		reader = errorReader{
    90  			err: fmt.Errorf("failed to decode message with codec '%v': %w", codec, err),
    91  		}
    92  	}
    93  
    94  	return newOneTimeReader(reader)
    95  }
    96  
    97  type errorReader struct {
    98  	err error
    99  }
   100  
   101  func (u errorReader) Read(p []byte) (n int, err error) {
   102  	return 0, u.err
   103  }
   104  
   105  type PublicMessageBuilder struct {
   106  	mess *PublicMessage
   107  }
   108  
   109  func NewPublicMessageBuilder() *PublicMessageBuilder {
   110  	res := &PublicMessageBuilder{}
   111  	res.initMessage()
   112  
   113  	return res
   114  }
   115  
   116  func (pmb *PublicMessageBuilder) initMessage() {
   117  	pmb.mess = &PublicMessage{
   118  		commitRange: CommitRange{PartitionSession: NewPartitionSession(
   119  			context.Background(),
   120  			"",
   121  			0,
   122  			0,
   123  			"",
   124  			0,
   125  			0,
   126  			0,
   127  		)},
   128  	}
   129  }
   130  
   131  // Seqno set message Seqno
   132  func (pmb *PublicMessageBuilder) Seqno(seqNo int64) *PublicMessageBuilder {
   133  	pmb.mess.SeqNo = seqNo
   134  
   135  	return pmb
   136  }
   137  
   138  // CreatedAt set message CreatedAt
   139  func (pmb *PublicMessageBuilder) CreatedAt(createdAt time.Time) *PublicMessageBuilder {
   140  	pmb.mess.CreatedAt = createdAt
   141  
   142  	return pmb
   143  }
   144  
   145  func (pmb *PublicMessageBuilder) Metadata(metadata map[string][]byte) *PublicMessageBuilder {
   146  	pmb.mess.Metadata = make(map[string][]byte, len(metadata))
   147  	for key, val := range metadata {
   148  		pmb.mess.Metadata[key] = bytes.Clone(val)
   149  	}
   150  
   151  	return pmb
   152  }
   153  
   154  // MessageGroupID set message MessageGroupID
   155  func (pmb *PublicMessageBuilder) MessageGroupID(messageGroupID string) *PublicMessageBuilder {
   156  	pmb.mess.MessageGroupID = messageGroupID
   157  
   158  	return pmb
   159  }
   160  
   161  // WriteSessionMetadata set message WriteSessionMetadata
   162  func (pmb *PublicMessageBuilder) WriteSessionMetadata(writeSessionMetadata map[string]string) *PublicMessageBuilder {
   163  	pmb.mess.WriteSessionMetadata = writeSessionMetadata
   164  
   165  	return pmb
   166  }
   167  
   168  // Offset set message Offset
   169  func (pmb *PublicMessageBuilder) Offset(offset int64) *PublicMessageBuilder {
   170  	pmb.mess.Offset = offset
   171  	pmb.mess.commitRange.CommitOffsetStart = rawtopiccommon.Offset(offset)
   172  	pmb.mess.commitRange.CommitOffsetEnd = rawtopiccommon.Offset(offset + 1)
   173  
   174  	return pmb
   175  }
   176  
   177  // WrittenAt set message WrittenAt
   178  func (pmb *PublicMessageBuilder) WrittenAt(writtenAt time.Time) *PublicMessageBuilder {
   179  	pmb.mess.WrittenAt = writtenAt
   180  
   181  	return pmb
   182  }
   183  
   184  // ProducerID set message ProducerID
   185  func (pmb *PublicMessageBuilder) ProducerID(producerID string) *PublicMessageBuilder {
   186  	pmb.mess.ProducerID = producerID
   187  
   188  	return pmb
   189  }
   190  
   191  // DataAndUncompressedSize set message uncompressed content and field UncompressedSize
   192  func (pmb *PublicMessageBuilder) DataAndUncompressedSize(data []byte) *PublicMessageBuilder {
   193  	copyData := make([]byte, len(data))
   194  	copy(copyData, data)
   195  	pmb.mess.data = oneTimeReader{reader: bytes.NewReader(data)}
   196  	pmb.mess.dataConsumed = false
   197  	pmb.mess.rawDataLen = len(copyData)
   198  	pmb.mess.UncompressedSize = len(copyData)
   199  
   200  	return pmb
   201  }
   202  
   203  func (pmb *PublicMessageBuilder) CommitRange(cr CommitRange) *PublicMessageBuilder {
   204  	pmb.mess.commitRange = cr
   205  
   206  	return pmb
   207  }
   208  
   209  // UncompressedSize set message UncompressedSize
   210  func (pmb *PublicMessageBuilder) UncompressedSize(uncompressedSize int) *PublicMessageBuilder {
   211  	pmb.mess.UncompressedSize = uncompressedSize
   212  
   213  	return pmb
   214  }
   215  
   216  // Context set message Context
   217  func (pmb *PublicMessageBuilder) Context(ctx context.Context) *PublicMessageBuilder {
   218  	pmb.mess.commitRange.PartitionSession.SetContext(ctx)
   219  
   220  	return pmb
   221  }
   222  
   223  // Topic set message Topic
   224  func (pmb *PublicMessageBuilder) Topic(topic string) *PublicMessageBuilder {
   225  	pmb.mess.commitRange.PartitionSession.Topic = topic
   226  
   227  	return pmb
   228  }
   229  
   230  // PartitionID set message PartitionID
   231  func (pmb *PublicMessageBuilder) PartitionID(partitionID int64) *PublicMessageBuilder {
   232  	pmb.mess.commitRange.PartitionSession.PartitionID = partitionID
   233  
   234  	return pmb
   235  }
   236  
   237  func (pmb *PublicMessageBuilder) PartitionSession(session *PartitionSession) *PublicMessageBuilder {
   238  	pmb.mess.commitRange.PartitionSession = session
   239  
   240  	return pmb
   241  }
   242  
   243  func (pmb *PublicMessageBuilder) RawDataLen(val int) *PublicMessageBuilder {
   244  	pmb.mess.rawDataLen = val
   245  
   246  	return pmb
   247  }
   248  
   249  // Build return builded message and reset internal state for create new message
   250  func (pmb *PublicMessageBuilder) Build() *PublicMessage {
   251  	mess := pmb.mess
   252  	pmb.initMessage()
   253  
   254  	return mess
   255  }
   256  
   257  func MessageGetBufferBytesAccount(m *PublicMessage) int {
   258  	return m.bufferBytesAccount
   259  }
   260  
   261  func MessageWithSetCommitRangeForTest(m *PublicMessage, commitRange CommitRange) *PublicMessage {
   262  	m.commitRange = commitRange
   263  
   264  	return m
   265  }
   266  
   267  func MessageSetNilDataForTest(m *PublicMessage) {
   268  	m.data = newOneTimeReader(nil)
   269  	m.bufferBytesAccount = 0
   270  	m.dataConsumed = false
   271  }