milvus/internal/querynode/stats_service.go
neza2017 aeca8f85a5 Run cwrapper unittest in ci
Signed-off-by: neza2017 <yefu.chen@zilliz.com>
2020-12-05 17:39:58 +08:00

90 lines
2.1 KiB
Go

package querynode
import (
"context"
"fmt"
"log"
"strconv"
"time"
"github.com/zilliztech/milvus-distributed/internal/msgstream"
"github.com/zilliztech/milvus-distributed/internal/proto/internalpb"
)
type statsService struct {
ctx context.Context
statsStream *msgstream.MsgStream
replica *collectionReplica
}
func newStatsService(ctx context.Context, replica *collectionReplica) *statsService {
return &statsService{
ctx: ctx,
statsStream: nil,
replica: replica,
}
}
func (sService *statsService) start() {
sleepTimeInterval := Params.statsPublishInterval()
receiveBufSize := Params.statsReceiveBufSize()
// start pulsar
msgStreamURL, err := Params.pulsarAddress()
if err != nil {
log.Fatal(err)
}
producerChannels := []string{Params.statsChannelName()}
statsStream := msgstream.NewPulsarMsgStream(sService.ctx, receiveBufSize)
statsStream.SetPulsarClient(msgStreamURL)
statsStream.CreatePulsarProducers(producerChannels)
var statsMsgStream msgstream.MsgStream = statsStream
sService.statsStream = &statsMsgStream
(*sService.statsStream).Start()
// start service
fmt.Println("do segments statistic in ", strconv.Itoa(sleepTimeInterval), "ms")
for {
select {
case <-sService.ctx.Done():
return
case <-time.After(time.Duration(sleepTimeInterval) * time.Millisecond):
sService.sendSegmentStatistic()
}
}
}
func (sService *statsService) close() {
(*sService.statsStream).Close()
}
func (sService *statsService) sendSegmentStatistic() {
statisticData := (*sService.replica).getSegmentStatistics()
// fmt.Println("Publish segment statistic")
// fmt.Println(statisticData)
sService.publicStatistic(statisticData)
}
func (sService *statsService) publicStatistic(statistic *internalpb.QueryNodeSegStats) {
var msg msgstream.TsMsg = &msgstream.QueryNodeSegStatsMsg{
BaseMsg: msgstream.BaseMsg{
HashValues: []uint32{0},
},
QueryNodeSegStats: *statistic,
}
var msgPack = msgstream.MsgPack{
Msgs: []msgstream.TsMsg{msg},
}
err := (*sService.statsStream).Produce(&msgPack)
if err != nil {
log.Println(err)
}
}