milvus/internal/msgstream/mq_msgstream.go
zhenshan.cao bdd39c0623
Fix bug: check message payload before unmarshaling (#12315)
Signed-off-by: zhenshan.cao <zhenshan.cao@zilliz.com>
2021-11-29 14:31:18 +08:00

1019 lines
28 KiB
Go

// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package msgstream
import (
"context"
"errors"
"fmt"
"path/filepath"
"sync"
"time"
"github.com/golang/protobuf/proto"
"go.uber.org/zap"
"github.com/milvus-io/milvus/internal/log"
"github.com/milvus-io/milvus/internal/proto/commonpb"
"github.com/milvus-io/milvus/internal/proto/internalpb"
"github.com/milvus-io/milvus/internal/util/mqclient"
"github.com/milvus-io/milvus/internal/util/retry"
"github.com/milvus-io/milvus/internal/util/trace"
"github.com/opentracing/opentracing-go"
)
var _ MsgStream = (*mqMsgStream)(nil)
type mqMsgStream struct {
ctx context.Context
client mqclient.Client
producers map[string]mqclient.Producer
producerChannels []string
consumers map[string]mqclient.Consumer
consumerChannels []string
readers map[string]mqclient.Reader
readerChannels []string
repackFunc RepackFunc
unmarshal UnmarshalDispatcher
receiveBuf chan *MsgPack
wait *sync.WaitGroup
streamCancel func()
bufSize int64
producerLock *sync.Mutex
consumerLock *sync.Mutex
readerLock *sync.Mutex
}
// NewMqMsgStream is used to generate a new mqMsgStream object
func NewMqMsgStream(ctx context.Context,
receiveBufSize int64,
bufSize int64,
client mqclient.Client,
unmarshal UnmarshalDispatcher) (*mqMsgStream, error) {
streamCtx, streamCancel := context.WithCancel(ctx)
producers := make(map[string]mqclient.Producer)
consumers := make(map[string]mqclient.Consumer)
readers := make(map[string]mqclient.Reader)
producerChannels := make([]string, 0)
consumerChannels := make([]string, 0)
readerChannels := make([]string, 0)
receiveBuf := make(chan *MsgPack, receiveBufSize)
stream := &mqMsgStream{
ctx: streamCtx,
client: client,
producers: producers,
producerChannels: producerChannels,
consumers: consumers,
consumerChannels: consumerChannels,
readers: readers,
readerChannels: readerChannels,
unmarshal: unmarshal,
bufSize: bufSize,
receiveBuf: receiveBuf,
streamCancel: streamCancel,
producerLock: &sync.Mutex{},
consumerLock: &sync.Mutex{},
readerLock: &sync.Mutex{},
wait: &sync.WaitGroup{},
}
return stream, nil
}
// AsProducer create producer to send message to channels
func (ms *mqMsgStream) AsProducer(channels []string) {
for _, channel := range channels {
if len(channel) == 0 {
log.Error("MsgStream asProducer's channel is a empty string")
break
}
fn := func() error {
pp, err := ms.client.CreateProducer(mqclient.ProducerOptions{Topic: channel})
if err != nil {
return err
}
if pp == nil {
return errors.New("Producer is nil")
}
ms.producerLock.Lock()
defer ms.producerLock.Unlock()
ms.producers[channel] = pp
ms.producerChannels = append(ms.producerChannels, channel)
return nil
}
err := retry.Do(context.TODO(), fn, retry.Attempts(20), retry.Sleep(time.Millisecond*200))
if err != nil {
errMsg := "Failed to create producer " + channel + ", error = " + err.Error()
panic(errMsg)
}
}
}
// Create consumer to receive message from channels
func (ms *mqMsgStream) AsConsumer(channels []string, subName string) {
ms.AsConsumerWithPosition(channels, subName, mqclient.SubscriptionPositionEarliest)
}
// Create consumer to receive message from channels, with initial position
// if initial position is set to latest, last message in the channel is exclusive
func (ms *mqMsgStream) AsConsumerWithPosition(channels []string, subName string, position mqclient.SubscriptionInitialPosition) {
for _, channel := range channels {
if _, ok := ms.consumers[channel]; ok {
continue
}
fn := func() error {
pc, err := ms.client.Subscribe(mqclient.ConsumerOptions{
Topic: channel,
SubscriptionName: subName,
Type: mqclient.Exclusive,
SubscriptionInitialPosition: position,
})
if err != nil {
return err
}
if pc == nil {
return errors.New("Consumer is nil")
}
ms.consumerLock.Lock()
defer ms.consumerLock.Unlock()
ms.consumers[channel] = pc
ms.consumerChannels = append(ms.consumerChannels, channel)
return nil
}
err := retry.Do(context.TODO(), fn, retry.Attempts(20), retry.Sleep(time.Millisecond*200))
if err != nil {
errMsg := "Failed to create consumer " + channel + ", error = " + err.Error()
panic(errMsg)
}
}
}
// AsProducer create producer to send message to channels
func (ms *mqMsgStream) AsReader(channels []string, subName string) {
for _, channel := range channels {
if len(channel) == 0 {
log.Error("MsgStream asProducer's channel is a empty string")
break
}
fn := func() error {
r, err := ms.client.CreateReader(mqclient.ReaderOptions{
Topic: channel,
StartMessageID: ms.client.EarliestMessageID(),
SubscriptionRolePrefix: subName,
})
if err != nil {
return err
}
if r == nil {
return errors.New("reader is nil")
}
ms.readerLock.Lock()
defer ms.readerLock.Unlock()
ms.readers[channel] = r
ms.readerChannels = append(ms.readerChannels, channel)
return nil
}
err := retry.Do(context.TODO(), fn, retry.Attempts(20), retry.Sleep(time.Millisecond*200))
if err != nil {
errMsg := "Failed to create producer " + channel + ", error = " + err.Error()
panic(errMsg)
}
}
}
func (ms *mqMsgStream) SetRepackFunc(repackFunc RepackFunc) {
ms.repackFunc = repackFunc
}
func (ms *mqMsgStream) Start() {
for _, c := range ms.consumers {
ms.wait.Add(1)
go ms.receiveMsg(c)
}
}
func (ms *mqMsgStream) Close() {
ms.streamCancel()
ms.wait.Wait()
for _, producer := range ms.producers {
if producer != nil {
producer.Close()
}
}
for _, consumer := range ms.consumers {
if consumer != nil {
consumer.Close()
}
}
}
func (ms *mqMsgStream) ComputeProduceChannelIndexes(tsMsgs []TsMsg) [][]int32 {
if len(tsMsgs) <= 0 {
return nil
}
reBucketValues := make([][]int32, len(tsMsgs))
channelNum := uint32(len(ms.producerChannels))
if channelNum == 0 {
return nil
}
for idx, tsMsg := range tsMsgs {
hashValues := tsMsg.HashKeys()
bucketValues := make([]int32, len(hashValues))
for index, hashValue := range hashValues {
bucketValues[index] = int32(hashValue % channelNum)
}
reBucketValues[idx] = bucketValues
}
return reBucketValues
}
func (ms *mqMsgStream) GetProduceChannels() []string {
return ms.producerChannels
}
func (ms *mqMsgStream) Produce(msgPack *MsgPack) error {
if msgPack == nil || len(msgPack.Msgs) <= 0 {
log.Debug("Warning: Receive empty msgPack")
return nil
}
if len(ms.producers) <= 0 {
return errors.New("nil producer in msg stream")
}
tsMsgs := msgPack.Msgs
reBucketValues := ms.ComputeProduceChannelIndexes(msgPack.Msgs)
var result map[int32]*MsgPack
var err error
if ms.repackFunc != nil {
result, err = ms.repackFunc(tsMsgs, reBucketValues)
} else {
msgType := (tsMsgs[0]).Type()
switch msgType {
case commonpb.MsgType_Insert:
result, err = InsertRepackFunc(tsMsgs, reBucketValues)
case commonpb.MsgType_Delete:
result, err = DeleteRepackFunc(tsMsgs, reBucketValues)
default:
result, err = DefaultRepackFunc(tsMsgs, reBucketValues)
}
}
if err != nil {
return err
}
for k, v := range result {
channel := ms.producerChannels[k]
for i := 0; i < len(v.Msgs); i++ {
sp, spanCtx := MsgSpanFromCtx(v.Msgs[i].TraceCtx(), v.Msgs[i])
mb, err := v.Msgs[i].Marshal(v.Msgs[i])
if err != nil {
return err
}
m, err := convertToByteArray(mb)
if err != nil {
return err
}
msg := &mqclient.ProducerMessage{Payload: m, Properties: map[string]string{}}
trace.InjectContextToPulsarMsgProperties(sp.Context(), msg.Properties)
ms.producerLock.Lock()
if _, err := ms.producers[channel].Send(
spanCtx,
msg,
); err != nil {
ms.producerLock.Unlock()
trace.LogError(sp, err)
sp.Finish()
return err
}
sp.Finish()
ms.producerLock.Unlock()
}
}
return nil
}
// ProduceMark send msg pack to all producers and returns corresponding msg id
// the returned message id serves as marking
func (ms *mqMsgStream) ProduceMark(msgPack *MsgPack) (map[string][]MessageID, error) {
ids := make(map[string][]MessageID)
if msgPack == nil || len(msgPack.Msgs) <= 0 {
return ids, errors.New("empty msgs")
}
if len(ms.producers) <= 0 {
return ids, errors.New("nil producer in msg stream")
}
tsMsgs := msgPack.Msgs
reBucketValues := ms.ComputeProduceChannelIndexes(msgPack.Msgs)
var result map[int32]*MsgPack
var err error
if ms.repackFunc != nil {
result, err = ms.repackFunc(tsMsgs, reBucketValues)
} else {
msgType := (tsMsgs[0]).Type()
switch msgType {
case commonpb.MsgType_Insert:
result, err = InsertRepackFunc(tsMsgs, reBucketValues)
case commonpb.MsgType_Delete:
result, err = DeleteRepackFunc(tsMsgs, reBucketValues)
default:
result, err = DefaultRepackFunc(tsMsgs, reBucketValues)
}
}
if err != nil {
return ids, err
}
for k, v := range result {
channel := ms.producerChannels[k]
for i, tsMsg := range v.Msgs {
sp, spanCtx := MsgSpanFromCtx(v.Msgs[i].TraceCtx(), tsMsg)
mb, err := tsMsg.Marshal(tsMsg)
if err != nil {
return ids, err
}
m, err := convertToByteArray(mb)
if err != nil {
return ids, err
}
msg := &mqclient.ProducerMessage{Payload: m, Properties: map[string]string{}}
trace.InjectContextToPulsarMsgProperties(sp.Context(), msg.Properties)
ms.producerLock.Lock()
id, err := ms.producers[channel].Send(
spanCtx,
msg,
)
if err != nil {
ms.producerLock.Unlock()
trace.LogError(sp, err)
sp.Finish()
return ids, err
}
ids[channel] = append(ids[channel], id)
sp.Finish()
ms.producerLock.Unlock()
}
}
return ids, nil
}
// Broadcast put msgPack to all producer in current msgstream
// which ignores repackFunc logic
func (ms *mqMsgStream) Broadcast(msgPack *MsgPack) error {
if msgPack == nil || len(msgPack.Msgs) <= 0 {
log.Debug("Warning: Receive empty msgPack")
return nil
}
for _, v := range msgPack.Msgs {
sp, spanCtx := MsgSpanFromCtx(v.TraceCtx(), v)
mb, err := v.Marshal(v)
if err != nil {
return err
}
m, err := convertToByteArray(mb)
if err != nil {
return err
}
msg := &mqclient.ProducerMessage{Payload: m, Properties: map[string]string{}}
trace.InjectContextToPulsarMsgProperties(sp.Context(), msg.Properties)
ms.producerLock.Lock()
for _, producer := range ms.producers {
if _, err := producer.Send(
spanCtx,
msg,
); err != nil {
ms.producerLock.Unlock()
trace.LogError(sp, err)
sp.Finish()
return err
}
}
ms.producerLock.Unlock()
sp.Finish()
}
return nil
}
// BroadcastMark broadcast msg pack to all producers and returns corresponding msg id
// the returned message id serves as marking
func (ms *mqMsgStream) BroadcastMark(msgPack *MsgPack) (map[string][]MessageID, error) {
ids := make(map[string][]MessageID)
if msgPack == nil || len(msgPack.Msgs) <= 0 {
return ids, errors.New("empty msgs")
}
for _, v := range msgPack.Msgs {
sp, spanCtx := MsgSpanFromCtx(v.TraceCtx(), v)
mb, err := v.Marshal(v)
if err != nil {
return ids, err
}
m, err := convertToByteArray(mb)
if err != nil {
return ids, err
}
msg := &mqclient.ProducerMessage{Payload: m, Properties: map[string]string{}}
trace.InjectContextToPulsarMsgProperties(sp.Context(), msg.Properties)
ms.producerLock.Lock()
for channel, producer := range ms.producers {
id, err := producer.Send(spanCtx, msg)
if err != nil {
ms.producerLock.Unlock()
trace.LogError(sp, err)
sp.Finish()
return ids, err
}
ids[channel] = append(ids[channel], id)
}
ms.producerLock.Unlock()
sp.Finish()
}
return ids, nil
}
func (ms *mqMsgStream) Consume() *MsgPack {
for {
select {
case <-ms.ctx.Done():
//log.Debug("context closed")
return nil
case cm, ok := <-ms.receiveBuf:
if !ok {
log.Debug("buf chan closed")
return nil
}
return cm
}
}
}
func (ms *mqMsgStream) getTsMsgFromConsumerMsg(msg mqclient.Message) (TsMsg, error) {
header := commonpb.MsgHeader{}
if msg.Payload() == nil {
return nil, fmt.Errorf("failed to unmarshal message header, payload is empty")
}
err := proto.Unmarshal(msg.Payload(), &header)
if err != nil {
return nil, fmt.Errorf("failed to unmarshal message header, err %s", err.Error())
}
if header.Base == nil {
return nil, fmt.Errorf("failed to unmarshal message, header is uncomplete")
}
tsMsg, err := ms.unmarshal.Unmarshal(msg.Payload(), header.Base.MsgType)
if err != nil {
return nil, fmt.Errorf("failed to unmarshal tsMsg, err %s", err.Error())
}
// set msg info to tsMsg
tsMsg.SetPosition(&MsgPosition{
ChannelName: filepath.Base(msg.Topic()),
MsgID: msg.ID().Serialize(),
})
return tsMsg, nil
}
func (ms *mqMsgStream) receiveMsg(consumer mqclient.Consumer) {
defer ms.wait.Done()
for {
select {
case <-ms.ctx.Done():
return
case msg, ok := <-consumer.Chan():
if !ok {
return
}
consumer.Ack(msg)
if msg.Payload() == nil {
log.Warn("MqMsgStream get msg whose payload is nil")
continue
}
tsMsg, err := ms.getTsMsgFromConsumerMsg(msg)
if err != nil {
log.Error("Failed to getTsMsgFromConsumerMsg", zap.Error(err))
continue
}
pos := tsMsg.Position()
tsMsg.SetPosition(&MsgPosition{
ChannelName: pos.ChannelName,
MsgID: pos.MsgID,
MsgGroup: consumer.Subscription(),
Timestamp: tsMsg.BeginTs(),
})
sp, ok := ExtractFromPulsarMsgProperties(tsMsg, msg.Properties())
if ok {
tsMsg.SetTraceCtx(opentracing.ContextWithSpan(context.Background(), sp))
}
msgPack := MsgPack{
Msgs: []TsMsg{tsMsg},
StartPositions: []*internalpb.MsgPosition{tsMsg.Position()},
EndPositions: []*internalpb.MsgPosition{tsMsg.Position()},
}
ms.receiveBuf <- &msgPack
sp.Finish()
}
}
}
func (ms *mqMsgStream) Chan() <-chan *MsgPack {
return ms.receiveBuf
}
func (ms *mqMsgStream) SeekReaders(msgPositions []*internalpb.MsgPosition) error {
for _, mp := range msgPositions {
reader, ok := ms.readers[mp.ChannelName]
if !ok {
return fmt.Errorf("channel %s not subscribed", mp.ChannelName)
}
messageID, err := ms.client.BytesToMsgID(mp.MsgID)
if err != nil {
return err
}
log.Debug("MsgStream reader begin to seek", zap.Any("MessageID", mp.MsgID))
err = reader.Seek(messageID)
if err != nil {
log.Debug("Failed to seek", zap.Error(err))
return err
}
}
return nil
}
func (ms *mqMsgStream) Next(ctx context.Context, channelName string) (TsMsg, error) {
reader, ok := ms.readers[channelName]
if !ok {
return nil, fmt.Errorf("reader for channel %s is not exist", channelName)
}
msg, err := reader.Next(ctx)
if err != nil {
return nil, err
}
if msg.Payload() == nil {
log.Warn("mqMsgStream reader Next get msg whose payload is nil")
return nil, nil
}
tsMsg, err := ms.getTsMsgFromConsumerMsg(msg)
if err != nil {
log.Error("Failed to getTsMsgFromConsumerMsg", zap.Error(err))
return nil, errors.New("Failed to getTsMsgFromConsumerMsg")
}
pos := tsMsg.Position()
tsMsg.SetPosition(&MsgPosition{
ChannelName: pos.ChannelName,
MsgID: pos.MsgID,
Timestamp: tsMsg.BeginTs(),
})
return tsMsg, nil
}
func (ms *mqMsgStream) HasNext(channelName string) bool {
reader, ok := ms.readers[channelName]
if !ok {
return false
}
return reader.HasNext()
}
// Seek reset the subscription associated with this consumer to a specific position, the seek position is exclusive
// User has to ensure mq_msgstream is not closed before seek, and the seek position is already written.
func (ms *mqMsgStream) Seek(msgPositions []*internalpb.MsgPosition) error {
for _, mp := range msgPositions {
consumer, ok := ms.consumers[mp.ChannelName]
if !ok {
return fmt.Errorf("channel %s not subscribed", mp.ChannelName)
}
messageID, err := ms.client.BytesToMsgID(mp.MsgID)
if err != nil {
return err
}
log.Debug("MsgStream begin to seek", zap.Any("MessageID", mp.MsgID))
err = consumer.Seek(messageID, false)
if err != nil {
log.Debug("Failed to seek", zap.Error(err))
return err
}
log.Debug("MsgStream seek finished", zap.Any("MessageID", messageID))
}
return nil
}
var _ MsgStream = (*MqTtMsgStream)(nil)
// MqTtMsgStream is a msgstream that contains timeticks
type MqTtMsgStream struct {
mqMsgStream
chanMsgBuf map[mqclient.Consumer][]TsMsg
chanMsgPos map[mqclient.Consumer]*internalpb.MsgPosition
chanStopChan map[mqclient.Consumer]chan bool
chanTtMsgTime map[mqclient.Consumer]Timestamp
chanMsgBufMutex *sync.Mutex
chanTtMsgTimeMutex *sync.RWMutex
chanWaitGroup *sync.WaitGroup
lastTimeStamp Timestamp
syncConsumer chan int
}
// NewMqTtMsgStream is used to generate a new MqTtMsgStream object
func NewMqTtMsgStream(ctx context.Context,
receiveBufSize int64,
bufSize int64,
client mqclient.Client,
unmarshal UnmarshalDispatcher) (*MqTtMsgStream, error) {
msgStream, err := NewMqMsgStream(ctx, receiveBufSize, bufSize, client, unmarshal)
if err != nil {
return nil, err
}
chanMsgBuf := make(map[mqclient.Consumer][]TsMsg)
chanMsgPos := make(map[mqclient.Consumer]*internalpb.MsgPosition)
chanStopChan := make(map[mqclient.Consumer]chan bool)
chanTtMsgTime := make(map[mqclient.Consumer]Timestamp)
syncConsumer := make(chan int, 1)
return &MqTtMsgStream{
mqMsgStream: *msgStream,
chanMsgBuf: chanMsgBuf,
chanMsgPos: chanMsgPos,
chanStopChan: chanStopChan,
chanTtMsgTime: chanTtMsgTime,
chanMsgBufMutex: &sync.Mutex{},
chanTtMsgTimeMutex: &sync.RWMutex{},
chanWaitGroup: &sync.WaitGroup{},
syncConsumer: syncConsumer,
}, nil
}
func (ms *MqTtMsgStream) addConsumer(consumer mqclient.Consumer, channel string) {
if len(ms.consumers) == 0 {
ms.syncConsumer <- 1
}
ms.consumers[channel] = consumer
ms.consumerChannels = append(ms.consumerChannels, channel)
ms.chanMsgBuf[consumer] = make([]TsMsg, 0)
ms.chanMsgPos[consumer] = &internalpb.MsgPosition{
ChannelName: channel,
MsgID: make([]byte, 0),
Timestamp: ms.lastTimeStamp,
}
ms.chanStopChan[consumer] = make(chan bool)
ms.chanTtMsgTime[consumer] = 0
}
// AsConsumer subscribes channels as consumer for a MsgStream
func (ms *MqTtMsgStream) AsConsumer(channels []string, subName string) {
ms.AsConsumerWithPosition(channels, subName, mqclient.SubscriptionPositionEarliest)
}
// AsConsumerWithPosition subscribes channels as consumer for a MsgStream and seeks to a certain position.
func (ms *MqTtMsgStream) AsConsumerWithPosition(channels []string, subName string, position mqclient.SubscriptionInitialPosition) {
for _, channel := range channels {
if _, ok := ms.consumers[channel]; ok {
continue
}
fn := func() error {
pc, err := ms.client.Subscribe(mqclient.ConsumerOptions{
Topic: channel,
SubscriptionName: subName,
Type: mqclient.Exclusive,
SubscriptionInitialPosition: position,
})
if err != nil {
return err
}
if pc == nil {
return errors.New("Consumer is nil")
}
ms.consumerLock.Lock()
defer ms.consumerLock.Unlock()
ms.addConsumer(pc, channel)
return nil
}
err := retry.Do(context.TODO(), fn, retry.Attempts(20), retry.Sleep(time.Millisecond*200))
if err != nil {
errMsg := "Failed to create consumer " + channel + ", error = " + err.Error()
panic(errMsg)
}
}
}
// Start will start a goroutine which keep carrying msg from pulsar/rocksmq to golang chan
func (ms *MqTtMsgStream) Start() {
if ms.consumers != nil {
ms.wait.Add(1)
go ms.bufMsgPackToChannel()
}
}
// Close will stop goroutine and free internal producers and consumers
func (ms *MqTtMsgStream) Close() {
ms.streamCancel()
close(ms.syncConsumer)
ms.wait.Wait()
for _, producer := range ms.producers {
if producer != nil {
producer.Close()
}
}
for _, consumer := range ms.consumers {
if consumer != nil {
consumer.Close()
}
}
for _, reader := range ms.readers {
if reader != nil {
reader.Close()
}
}
}
func (ms *MqTtMsgStream) bufMsgPackToChannel() {
defer ms.wait.Done()
chanTtMsgSync := make(map[mqclient.Consumer]bool)
// block here until addConsumer
if _, ok := <-ms.syncConsumer; !ok {
log.Debug("consumer closed!")
return
}
for {
select {
case <-ms.ctx.Done():
return
default:
ms.consumerLock.Lock()
// wait all channels get ttMsg
for _, consumer := range ms.consumers {
if !chanTtMsgSync[consumer] {
ms.chanWaitGroup.Add(1)
go ms.consumeToTtMsg(consumer)
}
}
ms.chanWaitGroup.Wait()
// block here until all channels reach same timetick
currTs, ok := ms.allChanReachSameTtMsg(chanTtMsgSync)
if !ok || currTs <= ms.lastTimeStamp {
//log.Printf("All timeTick's timestamps are inconsistent")
ms.consumerLock.Unlock()
continue
}
timeTickBuf := make([]TsMsg, 0)
startMsgPosition := make([]*internalpb.MsgPosition, 0)
endMsgPositions := make([]*internalpb.MsgPosition, 0)
ms.chanMsgBufMutex.Lock()
for consumer, msgs := range ms.chanMsgBuf {
if len(msgs) == 0 {
continue
}
tempBuffer := make([]TsMsg, 0)
var timeTickMsg TsMsg
for _, v := range msgs {
if v.Type() == commonpb.MsgType_TimeTick {
timeTickMsg = v
continue
}
if v.EndTs() <= currTs {
timeTickBuf = append(timeTickBuf, v)
//log.Debug("pack msg", zap.Uint64("curr", v.EndTs()), zap.Uint64("currTs", currTs))
} else {
tempBuffer = append(tempBuffer, v)
}
}
ms.chanMsgBuf[consumer] = tempBuffer
startMsgPosition = append(startMsgPosition, proto.Clone(ms.chanMsgPos[consumer]).(*internalpb.MsgPosition))
var newPos *internalpb.MsgPosition
if len(tempBuffer) > 0 {
// if tempBuffer is not empty, use tempBuffer[0] to seek
newPos = &internalpb.MsgPosition{
ChannelName: tempBuffer[0].Position().ChannelName,
MsgID: tempBuffer[0].Position().MsgID,
Timestamp: currTs,
MsgGroup: consumer.Subscription(),
}
endMsgPositions = append(endMsgPositions, newPos)
} else if timeTickMsg != nil {
// if tempBuffer is empty, use timeTickMsg to seek
newPos = &internalpb.MsgPosition{
ChannelName: timeTickMsg.Position().ChannelName,
MsgID: timeTickMsg.Position().MsgID,
Timestamp: currTs,
MsgGroup: consumer.Subscription(),
}
endMsgPositions = append(endMsgPositions, newPos)
}
ms.chanMsgPos[consumer] = newPos
}
ms.chanMsgBufMutex.Unlock()
ms.consumerLock.Unlock()
msgPack := MsgPack{
BeginTs: ms.lastTimeStamp,
EndTs: currTs,
Msgs: timeTickBuf,
StartPositions: startMsgPosition,
EndPositions: endMsgPositions,
}
//log.Debug("send msg pack", zap.Int("len", len(msgPack.Msgs)), zap.Uint64("currTs", currTs))
ms.receiveBuf <- &msgPack
ms.lastTimeStamp = currTs
}
}
}
// Save all msgs into chanMsgBuf[] till receive one ttMsg
func (ms *MqTtMsgStream) consumeToTtMsg(consumer mqclient.Consumer) {
defer ms.chanWaitGroup.Done()
for {
select {
case <-ms.ctx.Done():
return
case <-ms.chanStopChan[consumer]:
return
case msg, ok := <-consumer.Chan():
if !ok {
log.Debug("consumer closed!")
return
}
consumer.Ack(msg)
if msg.Payload() == nil {
log.Warn("MqTtMsgStream get msg whose payload is nil")
continue
}
tsMsg, err := ms.getTsMsgFromConsumerMsg(msg)
if err != nil {
log.Error("Failed to getTsMsgFromConsumerMsg", zap.Error(err))
continue
}
sp, ok := ExtractFromPulsarMsgProperties(tsMsg, msg.Properties())
if ok {
tsMsg.SetTraceCtx(opentracing.ContextWithSpan(context.Background(), sp))
}
ms.chanMsgBufMutex.Lock()
ms.chanMsgBuf[consumer] = append(ms.chanMsgBuf[consumer], tsMsg)
ms.chanMsgBufMutex.Unlock()
if tsMsg.Type() == commonpb.MsgType_TimeTick {
ms.chanTtMsgTimeMutex.Lock()
ms.chanTtMsgTime[consumer] = tsMsg.(*TimeTickMsg).Base.Timestamp
ms.chanTtMsgTimeMutex.Unlock()
sp.Finish()
return
}
sp.Finish()
}
}
}
// return true only when all channels reach same timetick
func (ms *MqTtMsgStream) allChanReachSameTtMsg(chanTtMsgSync map[mqclient.Consumer]bool) (Timestamp, bool) {
tsMap := make(map[Timestamp]int)
var maxTime Timestamp = 0
for _, t := range ms.chanTtMsgTime {
tsMap[t]++
if t > maxTime {
maxTime = t
}
}
// when all channels reach same timetick, timeMap should contain only 1 timestamp
if len(tsMap) <= 1 {
for consumer := range ms.chanTtMsgTime {
chanTtMsgSync[consumer] = false
}
return maxTime, true
}
for consumer := range ms.chanTtMsgTime {
ms.chanTtMsgTimeMutex.RLock()
chanTtMsgSync[consumer] = (ms.chanTtMsgTime[consumer] == maxTime)
ms.chanTtMsgTimeMutex.RUnlock()
}
return 0, false
}
// Seek to the specified position
func (ms *MqTtMsgStream) Seek(msgPositions []*internalpb.MsgPosition) error {
var consumer mqclient.Consumer
var mp *MsgPosition
var err error
fn := func() error {
var ok bool
consumer, ok = ms.consumers[mp.ChannelName]
if !ok {
return fmt.Errorf("please subcribe the channel, channel name =%s", mp.ChannelName)
}
if consumer == nil {
return fmt.Errorf("consumer is nil")
}
seekMsgID, err := ms.client.BytesToMsgID(mp.MsgID)
if err != nil {
return err
}
err = consumer.Seek(seekMsgID, true)
if err != nil {
return err
}
return nil
}
ms.consumerLock.Lock()
defer ms.consumerLock.Unlock()
for idx := range msgPositions {
mp = msgPositions[idx]
if len(mp.MsgID) == 0 {
return fmt.Errorf("when msgID's length equal to 0, please use AsConsumer interface")
}
if err = retry.Do(context.TODO(), fn, retry.Attempts(20), retry.Sleep(time.Millisecond*200)); err != nil {
return fmt.Errorf("Failed to seek, error %s", err.Error())
}
ms.addConsumer(consumer, mp.ChannelName)
ms.chanMsgPos[consumer] = (proto.Clone(mp)).(*MsgPosition)
// skip all data before current tt
runLoop := true
for runLoop {
select {
case <-ms.ctx.Done():
return nil
case msg, ok := <-consumer.Chan():
if !ok {
return fmt.Errorf("consumer closed")
}
consumer.Ack(msg)
headerMsg := commonpb.MsgHeader{}
err := proto.Unmarshal(msg.Payload(), &headerMsg)
if err != nil {
return fmt.Errorf("Failed to unmarshal message header, err %s", err.Error())
}
tsMsg, err := ms.unmarshal.Unmarshal(msg.Payload(), headerMsg.Base.MsgType)
if err != nil {
return fmt.Errorf("Failed to unmarshal tsMsg, err %s", err.Error())
}
if tsMsg.Type() == commonpb.MsgType_TimeTick && tsMsg.BeginTs() >= mp.Timestamp {
runLoop = false
break
} else if tsMsg.BeginTs() > mp.Timestamp {
tsMsg.SetPosition(&MsgPosition{
ChannelName: filepath.Base(msg.Topic()),
MsgID: msg.ID().Serialize(),
})
ms.chanMsgBuf[consumer] = append(ms.chanMsgBuf[consumer], tsMsg)
}
}
}
}
return nil
}