mirror of
https://gitee.com/milvus-io/milvus.git
synced 2025-12-06 17:18:35 +08:00
133 lines
3.1 KiB
Go
133 lines
3.1 KiB
Go
package writerclient
|
|
|
|
import (
|
|
"strconv"
|
|
|
|
"github.com/golang/protobuf/proto"
|
|
"go.etcd.io/etcd/clientv3"
|
|
|
|
"github.com/zilliztech/milvus-distributed/internal/errors"
|
|
"github.com/zilliztech/milvus-distributed/internal/kv"
|
|
etcdkv "github.com/zilliztech/milvus-distributed/internal/kv/etcd"
|
|
"github.com/zilliztech/milvus-distributed/internal/msgstream"
|
|
internalPb "github.com/zilliztech/milvus-distributed/internal/proto/internalpb"
|
|
pb "github.com/zilliztech/milvus-distributed/internal/proto/writerpb"
|
|
"github.com/zilliztech/milvus-distributed/internal/util/typeutil"
|
|
)
|
|
|
|
type UniqueID = typeutil.UniqueID
|
|
|
|
type Timestamp = typeutil.Timestamp
|
|
|
|
type Client struct {
|
|
kvClient kv.TxnBase // client of a reliable kv service, i.e. etcd client
|
|
kvPrefix string
|
|
|
|
flushStream msgstream.MsgStream
|
|
}
|
|
|
|
func NewWriterClient(etcdAddress string, kvRootPath string, writeNodeSegKvSubPath string, flushStream msgstream.MsgStream) (*Client, error) {
|
|
// init kv client
|
|
etcdClient, err := clientv3.New(clientv3.Config{Endpoints: []string{etcdAddress}})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
kvClient := etcdkv.NewEtcdKV(etcdClient, kvRootPath)
|
|
|
|
return &Client{
|
|
kvClient: kvClient,
|
|
kvPrefix: writeNodeSegKvSubPath,
|
|
flushStream: flushStream,
|
|
}, nil
|
|
}
|
|
|
|
type SegmentDescription struct {
|
|
SegmentID UniqueID
|
|
IsClosed bool
|
|
OpenTime Timestamp
|
|
CloseTime Timestamp
|
|
}
|
|
|
|
func (c *Client) FlushSegment(segmentID UniqueID) error {
|
|
baseMsg := msgstream.BaseMsg{
|
|
BeginTimestamp: 0,
|
|
EndTimestamp: 0,
|
|
HashValues: []uint32{0},
|
|
}
|
|
|
|
flushMsg := internalPb.FlushMsg{
|
|
MsgType: internalPb.MsgType_kFlush,
|
|
SegmentID: segmentID,
|
|
Timestamp: Timestamp(0),
|
|
}
|
|
|
|
fMsg := &msgstream.FlushMsg{
|
|
BaseMsg: baseMsg,
|
|
FlushMsg: flushMsg,
|
|
}
|
|
msgPack := msgstream.MsgPack{}
|
|
msgPack.Msgs = append(msgPack.Msgs, fMsg)
|
|
|
|
err := c.flushStream.Produce(&msgPack)
|
|
return err
|
|
}
|
|
|
|
func (c *Client) DescribeSegment(segmentID UniqueID) (*SegmentDescription, error) {
|
|
// query etcd
|
|
ret := &SegmentDescription{
|
|
SegmentID: segmentID,
|
|
IsClosed: false,
|
|
}
|
|
|
|
key := c.kvPrefix + strconv.FormatInt(segmentID, 10)
|
|
|
|
etcdKV, ok := c.kvClient.(*etcdkv.EtcdKV)
|
|
if !ok {
|
|
return nil, errors.New("type assertion failed for etcd kv")
|
|
}
|
|
count, err := etcdKV.GetCount(key)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
if count <= 0 {
|
|
ret.IsClosed = false
|
|
return ret, nil
|
|
}
|
|
|
|
value, err := c.kvClient.Load(key)
|
|
if err != nil {
|
|
return ret, err
|
|
}
|
|
|
|
flushMeta := pb.SegmentFlushMeta{}
|
|
err = proto.UnmarshalText(value, &flushMeta)
|
|
if err != nil {
|
|
return ret, err
|
|
}
|
|
ret.IsClosed = flushMeta.IsClosed
|
|
ret.OpenTime = flushMeta.OpenTime
|
|
ret.CloseTime = flushMeta.CloseTime
|
|
return ret, nil
|
|
}
|
|
|
|
func (c *Client) GetInsertBinlogPaths(segmentID UniqueID) (map[int64][]string, error) {
|
|
key := c.kvPrefix + strconv.FormatInt(segmentID, 10)
|
|
|
|
value, err := c.kvClient.Load(key)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
flushMeta := pb.SegmentFlushMeta{}
|
|
err = proto.UnmarshalText(value, &flushMeta)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
ret := make(map[int64][]string)
|
|
for _, field := range flushMeta.Fields {
|
|
ret[field.FieldID] = field.BinlogPaths
|
|
}
|
|
return ret, nil
|
|
}
|