test/milvus
1
0
Fork 0
mirror of https://gitee.com/milvus-io/milvus.git synced 2024-12-04 04:49:08 +08:00
Code Issues Actions 14 Packages Projects Releases Wiki Activity

Fix search waiting for seconds in singlenode

Browse Source 
Signed-off-by: yukun <kun.yu@zilliz.com>
This commit is contained in:
yukun 2021-03-09 17:42:13 +08:00 committed by yefu.chen
parent 408fe6e7ad
commit 0d499a9f99
3 changed files with 45 additions and 33 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

71
internal/msgstream/rmqms/rmq_msgstream.go
View File

@ -84,7 +84,7 @@ func (ms *RmqMsgStream) Close() {
} }
for _, consumer := range ms.consumers { for _, consumer := range ms.consumers {
_ = rocksmq.Rmq.DestroyConsumerGroup(consumer.GroupName, consumer.ChannelName) _ = rocksmq.Rmq.DestroyConsumerGroup(consumer.GroupName, consumer.ChannelName)
close(consumer.MsgNum) close(consumer.MsgMutex)
} }
} }
@ -112,12 +112,13 @@ func (ms *RmqMsgStream) AsConsumer(channels []string, groupName string) {
for _, channelName := range channels { for _, channelName := range channels {
consumer, err := rocksmq.Rmq.CreateConsumerGroup(groupName, channelName) consumer, err := rocksmq.Rmq.CreateConsumerGroup(groupName, channelName)
if err == nil { if err == nil {
consumer.MsgNum = make(chan int, ms.rmqBufSize) consumer.MsgMutex = make(chan struct{}, ms.rmqBufSize)
//consumer.MsgMutex <- struct{}{}
ms.consumers = append(ms.consumers, *consumer) ms.consumers = append(ms.consumers, *consumer)
ms.consumerChannels = append(ms.consumerChannels, channelName) ms.consumerChannels = append(ms.consumerChannels, channelName)
ms.consumerReflects = append(ms.consumerReflects, reflect.SelectCase{ ms.consumerReflects = append(ms.consumerReflects, reflect.SelectCase{
Dir: reflect.SelectRecv, Dir: reflect.SelectRecv,
Chan: reflect.ValueOf(consumer.MsgNum), Chan: reflect.ValueOf(consumer.MsgMutex),
}) })
ms.wait.Add(1) ms.wait.Add(1)
go ms.receiveMsg(*consumer) go ms.receiveMsg(*consumer)
@ -244,30 +245,35 @@ func (ms *RmqMsgStream) receiveMsg(consumer rocksmq.Consumer) {
select { select {
case <-ms.ctx.Done(): case <-ms.ctx.Done():
return return
case msgNum, ok := <-consumer.MsgNum: case _, ok := <-consumer.MsgMutex:
if !ok { if !ok {
return return
} }
rmqMsg, err := rocksmq.Rmq.Consume(consumer.GroupName, consumer.ChannelName, msgNum)
if err != nil {
log.Printf("Failed to consume message in rocksmq, error = %v", err)
continue
}
tsMsgList := make([]msgstream.TsMsg, 0) tsMsgList := make([]msgstream.TsMsg, 0)
for j := 0; j < len(rmqMsg); j++ { for {
headerMsg := commonpb.MsgHeader{} rmqMsgs, err := rocksmq.Rmq.Consume(consumer.GroupName, consumer.ChannelName, 1)
err := proto.Unmarshal(rmqMsg[j].Payload, &headerMsg)
if err != nil { if err != nil {
log.Printf("Failed to unmarshal message header, error = %v", err) log.Printf("Failed to consume message in rocksmq, error = %v", err)
continue continue
} }
tsMsg, err := ms.unmarshal.Unmarshal(rmqMsg[j].Payload, headerMsg.Base.MsgType) if len(rmqMsgs) == 0 {
break
}
rmqMsg := rmqMsgs[0]
headerMsg := commonpb.MsgHeader{}
err = proto.Unmarshal(rmqMsg.Payload, &headerMsg)
if err != nil {
log.Printf("Failed to unmar`shal message header, error = %v", err)
continue
}
tsMsg, err := ms.unmarshal.Unmarshal(rmqMsg.Payload, headerMsg.Base.MsgType)
if err != nil { if err != nil {
log.Printf("Failed to unmarshal tsMsg, error = %v", err) log.Printf("Failed to unmarshal tsMsg, error = %v", err)
continue continue
} }
tsMsgList = append(tsMsgList, tsMsg) tsMsgList = append(tsMsgList, tsMsg)
} }
if len(tsMsgList) > 0 { if len(tsMsgList) > 0 {
msgPack := util.MsgPack{Msgs: tsMsgList} msgPack := util.MsgPack{Msgs: tsMsgList}
ms.receiveBuf <- &msgPack ms.receiveBuf <- &msgPack
@ -326,12 +332,13 @@ func (ms *RmqTtMsgStream) AsConsumer(channels []string,
if err != nil { if err != nil {
panic(err.Error()) panic(err.Error())
} }
consumer.MsgNum = make(chan int, ms.rmqBufSize) consumer.MsgMutex = make(chan struct{}, ms.rmqBufSize)
//consumer.MsgMutex <- struct{}{}
ms.consumers = append(ms.consumers, *consumer) ms.consumers = append(ms.consumers, *consumer)
ms.consumerChannels = append(ms.consumerChannels, consumer.ChannelName) ms.consumerChannels = append(ms.consumerChannels, consumer.ChannelName)
ms.consumerReflects = append(ms.consumerReflects, reflect.SelectCase{ ms.consumerReflects = append(ms.consumerReflects, reflect.SelectCase{
Dir: reflect.SelectRecv, Dir: reflect.SelectRecv,
Chan: reflect.ValueOf(consumer.MsgNum), Chan: reflect.ValueOf(consumer.MsgMutex),
}) })
} }
} }
@ -432,25 +439,28 @@ func (ms *RmqTtMsgStream) findTimeTick(consumer rocksmq.Consumer,
select { select {
case <-ms.ctx.Done(): case <-ms.ctx.Done():
return return
case num, ok := <-consumer.MsgNum: case _, ok := <-consumer.MsgMutex:
if !ok { if !ok {
log.Printf("consumer closed!") log.Printf("consumer closed!")
return return
} }
rmqMsg, err := rocksmq.Rmq.Consume(consumer.GroupName, consumer.ChannelName, num) for {
if err != nil { rmqMsgs, err := rocksmq.Rmq.Consume(consumer.GroupName, consumer.ChannelName, 1)
log.Printf("Failed to consume message in rocksmq, error = %v", err) if err != nil {
continue log.Printf("Failed to consume message in rocksmq, error = %v", err)
} continue
}
for j := 0; j < len(rmqMsg); j++ { if len(rmqMsgs) == 0 {
return
}
rmqMsg := rmqMsgs[0]
headerMsg := commonpb.MsgHeader{} headerMsg := commonpb.MsgHeader{}
err := proto.Unmarshal(rmqMsg[j].Payload, &headerMsg) err = proto.Unmarshal(rmqMsg.Payload, &headerMsg)
if err != nil { if err != nil {
log.Printf("Failed to unmarshal message header, error = %v", err) log.Printf("Failed to unmarshal message header, error = %v", err)
continue continue
} }
tsMsg, err := ms.unmarshal.Unmarshal(rmqMsg[j].Payload, headerMsg.Base.MsgType) tsMsg, err := ms.unmarshal.Unmarshal(rmqMsg.Payload, headerMsg.Base.MsgType)
if err != nil { if err != nil {
log.Printf("Failed to unmarshal tsMsg, error = %v", err) log.Printf("Failed to unmarshal tsMsg, error = %v", err)
continue continue
@ -458,7 +468,7 @@ func (ms *RmqTtMsgStream) findTimeTick(consumer rocksmq.Consumer,
tsMsg.SetPosition(&msgstream.MsgPosition{ tsMsg.SetPosition(&msgstream.MsgPosition{
ChannelName: filepath.Base(consumer.ChannelName), ChannelName: filepath.Base(consumer.ChannelName),
MsgID: strconv.Itoa(int(rmqMsg[j].MsgID)), MsgID: strconv.Itoa(int(rmqMsg.MsgID)),
}) })
ms.unsolvedMutex.Lock() ms.unsolvedMutex.Lock()
@ -469,7 +479,8 @@ func (ms *RmqTtMsgStream) findTimeTick(consumer rocksmq.Consumer,
findMapMutex.Lock() findMapMutex.Lock()
eofMsgMap[consumer] = tsMsg.(*TimeTickMsg).Base.Timestamp eofMsgMap[consumer] = tsMsg.(*TimeTickMsg).Base.Timestamp
findMapMutex.Unlock() findMapMutex.Unlock()
return //consumer.MsgMutex <- struct{}{}
//return
} }
} }
} }
@ -504,8 +515,8 @@ func (ms *RmqTtMsgStream) Seek(mp *msgstream.MsgPosition) error {
ms.unsolvedMutex.Lock() ms.unsolvedMutex.Lock()
ms.unsolvedBuf[consumer] = make([]TsMsg, 0) ms.unsolvedBuf[consumer] = make([]TsMsg, 0)
// When rmq seek is called, msgNum can't be used before current msgs all consumed, because // When rmq seek is called, msgMutex can't be used before current msgs all consumed, because
// new msgNum is not generated. So just try to consume msgs // new msgMutex is not generated. So just try to consume msgs
for { for {
rmqMsg, err := rocksmq.Rmq.Consume(consumer.GroupName, consumer.ChannelName, 1) rmqMsg, err := rocksmq.Rmq.Consume(consumer.GroupName, consumer.ChannelName, 1)
if err != nil { if err != nil {

2
internal/util/rocksmq/global_rmq.go
View File

@ -15,7 +15,7 @@ var once sync.Once
type Consumer struct { type Consumer struct {
GroupName string GroupName string
ChannelName string ChannelName string
MsgNum chan int MsgMutex chan struct{}
} }
func InitRmq(rocksdbName string, idAllocator allocator.GIDAllocator) error { func InitRmq(rocksdbName string, idAllocator allocator.GIDAllocator) error {

5
internal/util/rocksmq/rocksmq.go
View File

@ -247,8 +247,8 @@ func (rmq *RocksMQ) Produce(channelName string, messages []ProducerMessage) erro
} }
for _, consumer := range rmq.notify[channelName] { for _, consumer := range rmq.notify[channelName] {
if consumer.MsgNum != nil { if consumer.MsgMutex != nil {
consumer.MsgNum <- msgLen consumer.MsgMutex <- struct{}{}
} }
} }
return nil return nil
@ -308,6 +308,7 @@ func (rmq *RocksMQ) Consume(groupName string, channelName string, n int) ([]Cons
return nil, err return nil, err
} }
// When already consume to last mes, an empty slice will be returned
if len(consumerMessage) == 0 { if len(consumerMessage) == 0 {
return consumerMessage, nil return consumerMessage, nil
} }