milvus-io
/
milvus
mirror of https://github.com/milvus-io/milvus.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Refactor meta in data coord (#6300) 

Signed-off-by: sunby <bingyi.sun@zilliz.com>
pull/6340/head
sunby 2021-07-07 14:02:01 +08:00 committed by GitHub
parent a1ba9e9dcd
commit 017484e093
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
9 changed files with 295 additions and 476 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
internal/datacoord/flush_monitor_test.go
View File

@ -42,22 +42,17 @@ func TestFlushMonitor(t *testing.T) {
// create seg0 for partition0, seg0/seg1 for partition1
segID0_0, err := mockAllocator.allocID()
assert.Nil(t, err)
segInfo0_0, err := buildSegment(collID, partID0, segID0_0, channelName)
assert.Nil(t, err)
segInfo0_0 := buildSegment(collID, partID0, segID0_0, channelName)
segID1_0, err := mockAllocator.allocID()
assert.Nil(t, err)
segInfo1_0, err := buildSegment(collID, partID1, segID1_0, channelName)
assert.Nil(t, err)
segInfo1_0 := buildSegment(collID, partID1, segID1_0, channelName)
segID1_1, err := mockAllocator.allocID()
assert.Nil(t, err)
segInfo1_1, err := buildSegment(collID, partID1, segID1_1, channelName)
assert.Nil(t, err)
segInfo1_1 := buildSegment(collID, partID1, segID1_1, channelName)
// check AddSegment
err = meta.AddSegment(segInfo0_0)
assert.Nil(t, err)
err = meta.AddSegment(segInfo0_0)
assert.NotNil(t, err)
err = meta.AddSegment(segInfo1_0)
assert.Nil(t, err)
err = meta.AddSegment(segInfo1_1)
@ -77,9 +72,8 @@ func TestFlushMonitor(t *testing.T) {
fm.segmentPolicy = estSegmentSizePolicy(1024*1024, 1024*1024*2) // row size 1Mib Limit 2 MB
segID3Rows, err := mockAllocator.allocID()
assert.Nil(t, err)
segInfo3Rows, err := buildSegment(collID, partID1, segID3Rows, channelName)
segInfo3Rows := buildSegment(collID, partID1, segID3Rows, channelName)
segInfo3Rows.NumOfRows = 3
assert.Nil(t, err)
ids := fm.CheckSegments([]*datapb.SegmentInfo{segInfo3Rows})
if assert.Equal(t, 1, len(ids)) {
@ -95,8 +89,7 @@ func TestFlushMonitor(t *testing.T) {
for i := 0; i < 100; i++ {
segID, err := mockAllocator.allocID()
assert.Nil(t, err)
seg, err := buildSegment(collID, partID0, segID, channelName2)
assert.Nil(t, err)
seg := buildSegment(collID, partID0, segID, channelName2)
seg.DmlPosition = &internalpb.MsgPosition{
Timestamp: uint64(i + 1),
}
@ -108,8 +101,7 @@ func TestFlushMonitor(t *testing.T) {
exSegID, err := mockAllocator.allocID()
assert.Nil(t, err)
seg, err := buildSegment(collID, partID0, exSegID, channelName2)
assert.Nil(t, err)
seg := buildSegment(collID, partID0, exSegID, channelName2)
seg.DmlPosition = &internalpb.MsgPosition{
Timestamp: uint64(0), // the oldest
}

47
internal/datacoord/grpc_services.go
View File

@ -95,7 +95,7 @@ func (s *Server) AssignSegmentID(ctx context.Context, req *datapb.AssignSegmentI
zap.String("channelName", r.GetChannelName()),
zap.Uint32("count", r.GetCount()))
if !s.meta.HasCollection(r.CollectionID) {
if coll := s.meta.GetCollection(r.CollectionID); coll == nil {
if err := s.loadCollectionFromRootCoord(ctx, r.CollectionID); err != nil {
errMsg := fmt.Sprintf("Can not load collection %d", r.CollectionID)
appendFailedAssignment(errMsg)
@ -163,10 +163,10 @@ func (s *Server) GetSegmentStates(ctx context.Context, req *datapb.GetSegmentSta
Status: &commonpb.Status{},
SegmentID: segmentID,
}
segmentInfo, err := s.meta.GetSegment(segmentID)
if err != nil {
segmentInfo := s.meta.GetSegment(segmentID)
if segmentInfo == nil {
state.Status.ErrorCode = commonpb.ErrorCode_UnexpectedError
state.Status.Reason = fmt.Sprintf("Failed to get segment %d, %s", segmentID, err.Error())
state.Status.Reason = fmt.Sprintf("Failed to get segment %d", segmentID)
} else {
state.Status.ErrorCode = commonpb.ErrorCode_Success
state.State = segmentInfo.GetState()
@ -227,11 +227,7 @@ func (s *Server) GetCollectionStatistics(ctx context.Context, req *datapb.GetCol
resp.Status.Reason = serverNotServingErrMsg
return resp, nil
}
nums, err := s.meta.GetNumRowsOfCollection(req.CollectionID)
if err != nil {
resp.Status.Reason = err.Error()
return resp, nil
}
nums := s.meta.GetNumRowsOfCollection(req.CollectionID)
resp.Status.ErrorCode = commonpb.ErrorCode_Success
resp.Stats = append(resp.Stats, &commonpb.KeyValuePair{Key: "row_count", Value: strconv.FormatInt(nums, 10)})
return resp, nil
@ -247,11 +243,7 @@ func (s *Server) GetPartitionStatistics(ctx context.Context, req *datapb.GetPart
resp.Status.Reason = serverNotServingErrMsg
return resp, nil
}
nums, err := s.meta.GetNumRowsOfPartition(req.CollectionID, req.PartitionID)
if err != nil {
resp.Status.Reason = err.Error()
return resp, nil
}
nums := s.meta.GetNumRowsOfPartition(req.CollectionID, req.PartitionID)
resp.Status.ErrorCode = commonpb.ErrorCode_Success
resp.Stats = append(resp.Stats, &commonpb.KeyValuePair{Key: "row_count", Value: strconv.FormatInt(nums, 10)})
return resp, nil
@ -278,9 +270,9 @@ func (s *Server) GetSegmentInfo(ctx context.Context, req *datapb.GetSegmentInfoR
}
infos := make([]*datapb.SegmentInfo, 0, len(req.SegmentIDs))
for _, id := range req.SegmentIDs {
info, err := s.meta.GetSegment(id)
if err != nil {
resp.Status.Reason = err.Error()
info := s.meta.GetSegment(id)
if info == nil {
resp.Status.Reason = fmt.Sprintf("Failed to get segment %d", id)
return resp, nil
}
infos = append(infos, info)
@ -304,10 +296,10 @@ func (s *Server) SaveBinlogPaths(ctx context.Context, req *datapb.SaveBinlogPath
zap.Any("checkpoints", req.GetCheckPoints()))
// check segment id & collection id matched
_, err := s.meta.GetCollection(req.GetCollectionID())
if err != nil {
log.Error("Failed to get collection info", zap.Int64("collectionID", req.GetCollectionID()), zap.Error(err))
resp.Reason = err.Error()
if coll := s.meta.GetCollection(req.GetCollectionID()); coll == nil {
errMsg := fmt.Sprintf("Failed to get collection info %d", req.GetCollectionID())
log.Error(errMsg)
resp.Reason = errMsg
return resp, nil
}
@ -382,10 +374,11 @@ func (s *Server) GetRecoveryInfo(ctx context.Context, req *datapb.GetRecoveryInf
segmentIDs := s.meta.GetSegmentsOfPartition(collectionID, partitionID)
segment2Binlogs := make(map[UniqueID][]*datapb.FieldBinlog)
for _, id := range segmentIDs {
segment, err := s.meta.GetSegment(id)
if err != nil {
log.Error("Get segment failed", zap.Int64("segmentID", id))
resp.Status.Reason = err.Error()
segment := s.meta.GetSegment(id)
if segment == nil {
errMsg := fmt.Sprintf("Failed to get segment %d", id)
log.Error(errMsg)
resp.Status.Reason = errMsg
return resp, nil
}
if segment.State != commonpb.SegmentState_Flushed && segment.State != commonpb.SegmentState_Flushing {
@ -472,8 +465,8 @@ func (s *Server) GetFlushedSegments(ctx context.Context, req *datapb.GetFlushedS
}
ret := make([]UniqueID, 0, len(segmentIDs))
for _, id := range segmentIDs {
s, err := s.meta.GetSegment(id)
if err != nil || s.GetState() != commonpb.SegmentState_Flushed {
s := s.meta.GetSegment(id)
if s == nil || s.GetState() != commonpb.SegmentState_Flushed {
continue
}
ret = append(ret, id)

312
internal/datacoord/meta.go
View File

@ -16,12 +16,9 @@ import (
"github.com/golang/protobuf/proto"
"github.com/milvus-io/milvus/internal/kv"
"github.com/milvus-io/milvus/internal/log"
"go.uber.org/zap"
"github.com/milvus-io/milvus/internal/proto/commonpb"
"github.com/milvus-io/milvus/internal/proto/datapb"
"github.com/milvus-io/milvus/internal/proto/internalpb"
)
const (
@ -29,39 +26,18 @@ const (
segmentPrefix = metaPrefix + "/s"
)
var (
errSegmentNotFound = func(segmentID UniqueID) error {
return fmt.Errorf("segment %d not found", segmentID)
}
errCollectionNotFound = func(collectionID UniqueID) error {
return fmt.Errorf("collection %d not found", collectionID)
}
errPartitionNotFound = func(partitionID UniqueID) error {
return fmt.Errorf("partition %d not found", partitionID)
}
errCollectionExist = func(collectionName string, collectionID UniqueID) error {
return fmt.Errorf("collection %s with id %d already exist", collectionName, collectionID)
}
errPartitionExist = func(partitionID UniqueID) error {
return fmt.Errorf("partition %d already exist", partitionID)
}
errSegmentExist = func(segmentID UniqueID) error {
return fmt.Errorf("segment %d already exist", segmentID)
}
)
type meta struct {
sync.RWMutex
client kv.TxnKV // client of a reliable kv service, i.e. etcd client
collections map[UniqueID]*datapb.CollectionInfo // collection id to collection info
segments map[UniqueID]*datapb.SegmentInfo // segment id to segment info
segments *SegmentsInfo // segment id to segment info
}
func newMeta(kv kv.TxnKV) (*meta, error) {
mt := &meta{
client: kv,
collections: make(map[UniqueID]*datapb.CollectionInfo),
segments: make(map[UniqueID]*datapb.SegmentInfo),
segments: NewSegmentsInfo(),
}
err := mt.reloadFromKV()
if err != nil {
@ -82,95 +58,57 @@ func (m *meta) reloadFromKV() error {
if err != nil {
return fmt.Errorf("DataCoord reloadFromKV UnMarshalText datapb.SegmentInfo err:%w", err)
}
m.segments[segmentInfo.ID] = segmentInfo
m.segments.SetSegment(segmentInfo.GetID(), segmentInfo)
}
return nil
}
func (m *meta) AddCollection(collection *datapb.CollectionInfo) error {
func (m *meta) AddCollection(collection *datapb.CollectionInfo) {
m.Lock()
defer m.Unlock()
if _, ok := m.collections[collection.ID]; ok {
return errCollectionExist(collection.GetSchema().GetName(), collection.GetID())
}
m.collections[collection.ID] = collection
return nil
}
func (m *meta) DropCollection(collectionID UniqueID) error {
m.Lock()
defer m.Unlock()
if _, ok := m.collections[collectionID]; !ok {
return errCollectionNotFound(collectionID)
}
key := buildCollectionPath(collectionID)
if err := m.client.RemoveWithPrefix(key); err != nil {
return err
}
delete(m.collections, collectionID)
for i, info := range m.segments {
if info.CollectionID == collectionID {
delete(m.segments, i)
}
}
return nil
}
func (m *meta) HasCollection(collID UniqueID) bool {
func (m *meta) GetCollection(collectionID UniqueID) *datapb.CollectionInfo {
m.RLock()
defer m.RUnlock()
_, ok := m.collections[collID]
return ok
}
func (m *meta) GetCollection(collectionID UniqueID) (*datapb.CollectionInfo, error) {
m.RLock()
defer m.RUnlock()
collection, ok := m.collections[collectionID]
if !ok {
return nil, errCollectionNotFound(collectionID)
return nil
}
return proto.Clone(collection).(*datapb.CollectionInfo), nil
return collection
}
func (m *meta) GetNumRowsOfCollection(collectionID UniqueID) (int64, error) {
func (m *meta) GetNumRowsOfCollection(collectionID UniqueID) int64 {
m.RLock()
defer m.RUnlock()
var ret int64 = 0
for _, segment := range m.segments {
if segment.CollectionID == collectionID {
segments := m.segments.GetSegments()
for _, segment := range segments {
if segment.GetCollectionID() == collectionID {
ret += segment.GetNumOfRows()
}
}
return ret, nil
return ret
}
func (m *meta) AddSegment(segment *datapb.SegmentInfo) error {
m.Lock()
defer m.Unlock()
if _, ok := m.segments[segment.ID]; ok {
return errSegmentExist(segment.GetID())
}
m.segments[segment.ID] = segment
m.segments.SetSegment(segment.GetID(), segment)
if err := m.saveSegmentInfo(segment); err != nil {
return err
}
return nil
}
func (m *meta) UpdateSegmentStatistic(stats *internalpb.SegmentStatisticsUpdates) error {
func (m *meta) SetRowCount(segmentID UniqueID, rowCount int64) error {
m.Lock()
defer m.Unlock()
seg, ok := m.segments[stats.SegmentID]
if !ok {
return errSegmentNotFound(stats.SegmentID)
}
seg.NumOfRows = stats.NumRows
if err := m.saveSegmentInfo(seg); err != nil {
return err
m.segments.SetRowCount(segmentID, rowCount)
if segment := m.segments.GetSegment(segmentID); segment != nil {
return m.saveSegmentInfo(segment)
}
return nil
}
@ -178,14 +116,9 @@ func (m *meta) UpdateSegmentStatistic(stats *internalpb.SegmentStatisticsUpdates
func (m *meta) SetLastExpireTime(segmentID UniqueID, expireTs Timestamp) error {
m.Lock()
defer m.Unlock()
segment, ok := m.segments[segmentID]
if !ok {
return errSegmentNotFound(segmentID)
}
segment.LastExpireTime = expireTs
if err := m.saveSegmentInfo(segment); err != nil {
return err
m.segments.SetLasteExpiraTime(segmentID, expireTs)
if segment := m.segments.GetSegment(segmentID); segment != nil {
return m.saveSegmentInfo(segment)
}
return nil
}
@ -193,41 +126,26 @@ func (m *meta) SetLastExpireTime(segmentID UniqueID, expireTs Timestamp) error {
func (m *meta) DropSegment(segmentID UniqueID) error {
m.Lock()
defer m.Unlock()
segment, ok := m.segments[segmentID]
if !ok {
return errSegmentNotFound(segmentID)
}
segment := m.segments.GetSegment(segmentID)
m.segments.DropSegment(segmentID)
if err := m.removeSegmentInfo(segment); err != nil {
return err
}
delete(m.segments, segmentID)
return nil
}
func (m *meta) GetSegment(segID UniqueID) (*datapb.SegmentInfo, error) {
func (m *meta) GetSegment(segID UniqueID) *datapb.SegmentInfo {
m.RLock()
defer m.RUnlock()
segment, ok := m.segments[segID]
if !ok {
return nil, errSegmentNotFound(segID)
}
return proto.Clone(segment).(*datapb.SegmentInfo), nil
return m.segments.GetSegment(segID)
}
func (m *meta) SealSegment(segID UniqueID) error {
func (m *meta) SetState(segmentID UniqueID, state commonpb.SegmentState) error {
m.Lock()
defer m.Unlock()
segInfo, ok := m.segments[segID]
if !ok {
return errSegmentNotFound(segID)
}
segInfo.State = commonpb.SegmentState_Sealed
if err := m.saveSegmentInfo(segInfo); err != nil {
return err
m.segments.SetState(segmentID, state)
if segInfo := m.segments.GetSegment(segmentID); segInfo != nil {
return m.saveSegmentInfo(segInfo)
}
return nil
}
@ -237,53 +155,43 @@ func (m *meta) SaveBinlogAndCheckPoints(segID UniqueID, flushed bool,
startPositions []*datapb.SegmentStartPosition) error {
m.Lock()
defer m.Unlock()
segment, ok := m.segments[segID]
if !ok {
return errSegmentNotFound(segID)
}
kv := make(map[string]string)
for k, v := range binlogs {
kv[k] = v
}
if flushed {
segment.State = commonpb.SegmentState_Flushing
m.segments.SetState(segID, commonpb.SegmentState_Flushing)
}
modifiedSegments := make(map[UniqueID]struct{})
modSegments := make([]UniqueID, 0)
for _, pos := range startPositions {
segment, ok := m.segments[pos.GetSegmentID()]
if !ok {
log.Warn("Failed to find segment", zap.Int64("id", pos.GetSegmentID()))
continue
}
if len(pos.GetStartPosition().GetMsgID()) != 0 {
continue
}
segment.StartPosition = pos.GetStartPosition()
modifiedSegments[segment.GetID()] = struct{}{}
if segment := m.segments.GetSegment(pos.GetSegmentID()); segment != nil {
m.segments.SetStartPosition(pos.GetSegmentID(), pos.GetStartPosition())
modSegments = append(modSegments, pos.GetSegmentID())
}
}
for _, cp := range checkpoints {
segment, ok := m.segments[cp.SegmentID]
if !ok {
log.Warn("Failed to find segment", zap.Int64("id", cp.SegmentID))
continue
if segment := m.segments.GetSegment(cp.GetSegmentID()); segment != nil {
if segment.DmlPosition != nil && segment.DmlPosition.Timestamp >= cp.Position.Timestamp {
// segment position in etcd is larger than checkpoint, then dont change it
continue
}
m.segments.SetDmlPositino(cp.GetSegmentID(), cp.GetPosition())
m.segments.SetRowCount(cp.GetSegmentID(), cp.GetNumOfRows())
modSegments = append(modSegments, segment.GetID())
}
if segment.DmlPosition != nil && segment.DmlPosition.Timestamp >= cp.Position.Timestamp {
// segment position in etcd is larger than checkpoint, then dont change it
continue
}
segment.DmlPosition = cp.Position
segment.NumOfRows = cp.NumOfRows
modifiedSegments[segment.GetID()] = struct{}{}
}
for id := range modifiedSegments {
segment = m.segments[id]
segBytes := proto.MarshalTextString(segment)
key := buildSegmentPath(segment.GetCollectionID(), segment.GetPartitionID(), segment.GetID())
kv[key] = segBytes
for _, id := range modSegments {
if segment := m.segments.GetSegment(id); segment != nil {
segBytes := proto.MarshalTextString(segment)
key := buildSegmentPath(segment.GetCollectionID(), segment.GetPartitionID(), segment.GetID())
kv[key] = segBytes
}
}
if err := m.saveKvTxn(kv); err != nil {
@ -296,36 +204,22 @@ func (m *meta) GetSegmentsByChannel(dmlCh string) []*datapb.SegmentInfo {
m.RLock()
defer m.RUnlock()
infos := make([]*datapb.SegmentInfo, 0)
for _, segment := range m.segments {
segments := m.segments.GetSegments()
for _, segment := range segments {
if segment.InsertChannel != dmlCh {
continue
}
cInfo := proto.Clone(segment).(*datapb.SegmentInfo)
infos = append(infos, cInfo)
infos = append(infos, segment)
}
return infos
}
func (m *meta) FlushSegment(segID UniqueID) error {
m.Lock()
defer m.Unlock()
segInfo, ok := m.segments[segID]
if !ok {
return errSegmentNotFound(segID)
}
segInfo.State = commonpb.SegmentState_Flushed
if err := m.saveSegmentInfo(segInfo); err != nil {
return err
}
return nil
}
func (m *meta) GetSegmentsOfCollection(collectionID UniqueID) []UniqueID {
m.RLock()
defer m.RUnlock()
ret := make([]UniqueID, 0)
for _, info := range m.segments {
segments := m.segments.GetSegments()
for _, info := range segments {
if info.CollectionID == collectionID {
ret = append(ret, info.ID)
}
@ -336,9 +230,9 @@ func (m *meta) GetSegmentsOfCollection(collectionID UniqueID) []UniqueID {
func (m *meta) GetSegmentsOfPartition(collectionID, partitionID UniqueID) []UniqueID {
m.RLock()
defer m.RUnlock()
ret := make([]UniqueID, 0)
for _, info := range m.segments {
segments := m.segments.GetSegments()
for _, info := range segments {
if info.CollectionID == collectionID && info.PartitionID == partitionID {
ret = append(ret, info.ID)
}
@ -346,107 +240,43 @@ func (m *meta) GetSegmentsOfPartition(collectionID, partitionID UniqueID) []Uniq
return ret
}
func (m *meta) AddPartition(collectionID UniqueID, partitionID UniqueID) error {
m.Lock()
defer m.Unlock()
coll, ok := m.collections[collectionID]
if !ok {
return errCollectionNotFound(collectionID)
}
for _, t := range coll.Partitions {
if t == partitionID {
return errPartitionExist(partitionID)
}
}
coll.Partitions = append(coll.Partitions, partitionID)
return nil
}
func (m *meta) DropPartition(collectionID UniqueID, partitionID UniqueID) error {
m.Lock()
defer m.Unlock()
collection, ok := m.collections[collectionID]
if !ok {
return errCollectionNotFound(collectionID)
}
idx := -1
for i, id := range collection.Partitions {
if partitionID == id {
idx = i
break
}
}
if idx == -1 {
return errPartitionNotFound(partitionID)
}
prefix := buildPartitionPath(collectionID, partitionID)
if err := m.client.RemoveWithPrefix(prefix); err != nil {
return err
}
collection.Partitions = append(collection.Partitions[:idx], collection.Partitions[idx+1:]...)
for i, info := range m.segments {
if info.PartitionID == partitionID {
delete(m.segments, i)
}
}
return nil
}
func (m *meta) HasPartition(collID UniqueID, partID UniqueID) bool {
m.RLock()
defer m.RUnlock()
coll, ok := m.collections[collID]
if !ok {
return false
}
for _, id := range coll.Partitions {
if partID == id {
return true
}
}
return false
}
func (m *meta) GetNumRowsOfPartition(collectionID UniqueID, partitionID UniqueID) (int64, error) {
func (m *meta) GetNumRowsOfPartition(collectionID UniqueID, partitionID UniqueID) int64 {
m.RLock()
defer m.RUnlock()
var ret int64 = 0
for _, info := range m.segments {
segments := m.segments.GetSegments()
for _, info := range segments {
if info.CollectionID == collectionID && info.PartitionID == partitionID {
ret += info.NumOfRows
}
}
return ret, nil
return ret
}
func (m *meta) GetUnFlushedSegments() []*datapb.SegmentInfo {
m.RLock()
defer m.RUnlock()
segments := make([]*datapb.SegmentInfo, 0)
for _, info := range m.segments {
ret := make([]*datapb.SegmentInfo, 0)
segments := m.segments.GetSegments()
for _, info := range segments {
if info.State != commonpb.SegmentState_Flushing && info.State != commonpb.SegmentState_Flushed {
cInfo := proto.Clone(info).(*datapb.SegmentInfo)
segments = append(segments, cInfo)
ret = append(ret, info)
}
}
return segments
return ret
}
func (m *meta) GetFlushingSegments() []*datapb.SegmentInfo {
m.RLock()
defer m.RUnlock()
segments := make([]*datapb.SegmentInfo, 0)
for _, info := range m.segments {
ret := make([]*datapb.SegmentInfo, 0)
segments := m.segments.GetSegments()
for _, info := range segments {
if info.State == commonpb.SegmentState_Flushing {
cInfo := proto.Clone(info).(*datapb.SegmentInfo)
segments = append(segments, cInfo)
ret = append(ret, info)
}
}
return segments
return ret
}
func (m *meta) saveSegmentInfo(segment *datapb.SegmentInfo) error {
@ -477,7 +307,7 @@ func buildPartitionPath(collectionID UniqueID, partitionID UniqueID) string {
return fmt.Sprintf("%s/%d/%d/", segmentPrefix, collectionID, partitionID)
}
func buildSegment(collectionID UniqueID, partitionID UniqueID, segmentID UniqueID, channelName string) (*datapb.SegmentInfo, error) {
func buildSegment(collectionID UniqueID, partitionID UniqueID, segmentID UniqueID, channelName string) *datapb.SegmentInfo {
return &datapb.SegmentInfo{
ID: segmentID,
CollectionID: collectionID,
@ -485,5 +315,5 @@ func buildSegment(collectionID UniqueID, partitionID UniqueID, segmentID UniqueI
InsertChannel: channelName,
NumOfRows: 0,
State: commonpb.SegmentState_Growing,
}, nil
}
}

159
internal/datacoord/meta_test.go
View File

@ -42,113 +42,46 @@ func TestMeta_Basic(t *testing.T) {
}
t.Run("Test Collection", func(t *testing.T) {
// check add collection
err = meta.AddCollection(collInfo)
assert.Nil(t, err)
// check add existed collection
err = meta.AddCollection(collInfo)
assert.NotNil(t, err)
meta.AddCollection(collInfo)
// check has collection
has := meta.HasCollection(collID)
assert.True(t, has)
collInfo := meta.GetCollection(collID)
assert.NotNil(t, collInfo)
// check partition info
collInfo, err = meta.GetCollection(collID)
assert.Nil(t, err)
assert.EqualValues(t, collID, collInfo.ID)
assert.EqualValues(t, testSchema, collInfo.Schema)
assert.EqualValues(t, 2, len(collInfo.Partitions))
assert.EqualValues(t, partID0, collInfo.Partitions[0])
assert.EqualValues(t, partID1, collInfo.Partitions[1])
// check drop collection
err = meta.DropCollection(collID)
assert.Nil(t, err)
has = meta.HasCollection(collID)
assert.False(t, has)
_, err = meta.GetCollection(collID)
assert.NotNil(t, err)
})
t.Run("Test Partition", func(t *testing.T) {
err = meta.AddCollection(collInfoWoPartition)
assert.Nil(t, err)
// check add partition
err = meta.AddPartition(collID, partID0)
assert.Nil(t, err)
err = meta.AddPartition(collID, partID1)
assert.Nil(t, err)
exist0 := meta.HasPartition(collID, partID0)
assert.True(t, exist0)
exist1 := meta.HasPartition(collID, partID1)
assert.True(t, exist1)
// check add existed partition
err = meta.AddPartition(collID, partID0)
assert.NotNil(t, err)
// check GetCollection
collInfo, err = meta.GetCollection(collID)
assert.Nil(t, err)
assert.EqualValues(t, 2, len(collInfo.Partitions))
assert.Contains(t, collInfo.Partitions, partID0)
assert.Contains(t, collInfo.Partitions, partID1)
// check DropPartition
err = meta.DropPartition(collID, partID0)
assert.Nil(t, err)
exist0 = meta.HasPartition(collID, partID0)
assert.False(t, exist0)
exist1 = meta.HasPartition(collID, partID1)
assert.True(t, exist1)
// check DropPartition twice
err = meta.DropPartition(collID, partID0)
assert.NotNil(t, err)
err = meta.DropCollection(collID)
assert.Nil(t, err)
})
t.Run("Test Segment", func(t *testing.T) {
err = meta.AddCollection(collInfoWoPartition)
assert.Nil(t, err)
err = meta.AddPartition(collID, partID0)
assert.Nil(t, err)
meta.AddCollection(collInfoWoPartition)
// create seg0 for partition0, seg0/seg1 for partition1
segID0_0, err := mockAllocator.allocID()
assert.Nil(t, err)
segInfo0_0, err := buildSegment(collID, partID0, segID0_0, channelName)
assert.Nil(t, err)
segInfo0_0 := buildSegment(collID, partID0, segID0_0, channelName)
segID1_0, err := mockAllocator.allocID()
assert.Nil(t, err)
segInfo1_0, err := buildSegment(collID, partID1, segID1_0, channelName)
assert.Nil(t, err)
segInfo1_0 := buildSegment(collID, partID1, segID1_0, channelName)
segID1_1, err := mockAllocator.allocID()
assert.Nil(t, err)
segInfo1_1, err := buildSegment(collID, partID1, segID1_1, channelName)
assert.Nil(t, err)
segInfo1_1 := buildSegment(collID, partID1, segID1_1, channelName)
// check AddSegment
err = meta.AddSegment(segInfo0_0)
assert.Nil(t, err)
err = meta.AddSegment(segInfo0_0)
assert.NotNil(t, err)
err = meta.AddSegment(segInfo1_0)
assert.Nil(t, err)
err = meta.AddSegment(segInfo1_1)
assert.Nil(t, err)
// check GetSegment
info0_0, err := meta.GetSegment(segID0_0)
assert.Nil(t, err)
info0_0 := meta.GetSegment(segID0_0)
assert.NotNil(t, info0_0)
assert.True(t, proto.Equal(info0_0, segInfo0_0))
info1_0, err := meta.GetSegment(segID1_0)
assert.Nil(t, err)
info1_0 := meta.GetSegment(segID1_0)
assert.NotNil(t, info1_0)
assert.True(t, proto.Equal(info1_0, segInfo1_0))
// check GetSegmentsOfCollection
@ -174,19 +107,14 @@ func TestMeta_Basic(t *testing.T) {
assert.EqualValues(t, 1, len(segIDs))
assert.Contains(t, segIDs, segID1_1)
err = meta.SealSegment(segID0_0)
err = meta.SetState(segID0_0, commonpb.SegmentState_Sealed)
assert.Nil(t, err)
err = meta.FlushSegment(segID0_0)
err = meta.SetState(segID0_0, commonpb.SegmentState_Flushed)
assert.Nil(t, err)
info0_0, err = meta.GetSegment(segID0_0)
assert.Nil(t, err)
info0_0 = meta.GetSegment(segID0_0)
assert.NotNil(t, info0_0)
assert.EqualValues(t, commonpb.SegmentState_Flushed, info0_0.State)
err = meta.DropPartition(collID, partID0)
assert.Nil(t, err)
err = meta.DropCollection(collID)
assert.Nil(t, err)
})
t.Run("Test GetCount", func(t *testing.T) {
@ -195,15 +123,13 @@ func TestMeta_Basic(t *testing.T) {
const dim = 1024
// no segment
nums, err := meta.GetNumRowsOfCollection(collID)
assert.Nil(t, err)
nums := meta.GetNumRowsOfCollection(collID)
assert.EqualValues(t, 0, nums)
// add seg1 with 100 rows
segID0, err := mockAllocator.allocID()
assert.Nil(t, err)
segInfo0, err := buildSegment(collID, partID0, segID0, channelName)
assert.Nil(t, err)
segInfo0 := buildSegment(collID, partID0, segID0, channelName)
segInfo0.NumOfRows = rowCount0
err = meta.AddSegment(segInfo0)
assert.Nil(t, err)
@ -211,62 +137,17 @@ func TestMeta_Basic(t *testing.T) {
// add seg2 with 300 rows
segID1, err := mockAllocator.allocID()
assert.Nil(t, err)
segInfo1, err := buildSegment(collID, partID0, segID1, channelName)
assert.Nil(t, err)
segInfo1 := buildSegment(collID, partID0, segID1, channelName)
segInfo1.NumOfRows = rowCount1
err = meta.AddSegment(segInfo1)
assert.Nil(t, err)
// check partition/collection statistics
nums, err = meta.GetNumRowsOfPartition(collID, partID0)
assert.Nil(t, err)
nums = meta.GetNumRowsOfPartition(collID, partID0)
assert.EqualValues(t, (rowCount0 + rowCount1), nums)
nums, err = meta.GetNumRowsOfCollection(collID)
assert.Nil(t, err)
nums = meta.GetNumRowsOfCollection(collID)
assert.EqualValues(t, (rowCount0 + rowCount1), nums)
})
t.Run("Test Invalid", func(t *testing.T) {
collIDInvalid := UniqueID(10000)
partIDInvalid := UniqueID(10001)
segIDInvalid := UniqueID(10002)
// check drop non-exist collection
err = meta.DropCollection(collID)
assert.NotNil(t, err)
// add partition wo collection
err = meta.AddPartition(collID, partID0)
assert.NotNil(t, err)
// has partition wo collection
exist := meta.HasPartition(collID, partID0)
assert.False(t, exist)
err = meta.AddCollection(collInfo)
assert.Nil(t, err)
// check drop non-exist partition
err = meta.DropPartition(collIDInvalid, partID0)
assert.NotNil(t, err)
err = meta.DropPartition(collID, partIDInvalid)
assert.NotNil(t, err)
// check drop non-exist segment
err = meta.DropSegment(segIDInvalid)
assert.NotNil(t, err)
// check seal non-exist segment
err = meta.SealSegment(segIDInvalid)
assert.NotNil(t, err)
// check flush non-exist segment
err = meta.FlushSegment(segIDInvalid)
assert.NotNil(t, err)
err = meta.DropCollection(collID)
assert.Nil(t, err)
})
}
func TestGetUnFlushedSegments(t *testing.T) {

143
internal/datacoord/segment_info.go Normal file
View File

@ -0,0 +1,143 @@
package datacoord
import (
"github.com/gogo/protobuf/proto"
"github.com/milvus-io/milvus/internal/proto/commonpb"
"github.com/milvus-io/milvus/internal/proto/datapb"
"github.com/milvus-io/milvus/internal/proto/internalpb"
)
type SegmentsInfo struct {
segments map[UniqueID]*datapb.SegmentInfo
}
func NewSegmentsInfo() *SegmentsInfo {
return &SegmentsInfo{segments: make(map[UniqueID]*datapb.SegmentInfo)}
}
func (s *SegmentsInfo) GetSegment(segmentID UniqueID) *datapb.SegmentInfo {
segment, ok := s.segments[segmentID]
if !ok {
return nil
}
return segment
}
func (s *SegmentsInfo) GetSegments() []*datapb.SegmentInfo {
segments := make([]*datapb.SegmentInfo, 0, len(s.segments))
for _, segment := range s.segments {
segments = append(segments, segment)
}
return segments
}
func (s *SegmentsInfo) DropSegment(segmentID UniqueID) {
delete(s.segments, segmentID)
}
func (s *SegmentsInfo) SetSegment(segmentID UniqueID, segment *datapb.SegmentInfo) {
s.segments[segmentID] = segment
}
func (s *SegmentsInfo) SetRowCount(segmentID UniqueID, rowCount int64) {
if segment, ok := s.segments[segmentID]; ok {
s.segments[segmentID] = s.ShadowClone(segment, SetRowCount(rowCount))
}
}
func (s *SegmentsInfo) SetLasteExpiraTime(segmentID UniqueID, expireTs Timestamp) {
if segment, ok := s.segments[segmentID]; ok {
s.segments[segmentID] = s.ShadowClone(segment, SetExpireTime(expireTs))
}
}
func (s *SegmentsInfo) SetState(segmentID UniqueID, state commonpb.SegmentState) {
if segment, ok := s.segments[segmentID]; ok {
s.segments[segmentID] = s.ShadowClone(segment, SetState(state))
}
}
func (s *SegmentsInfo) SetDmlPositino(segmentID UniqueID, pos *internalpb.MsgPosition) {
if segment, ok := s.segments[segmentID]; ok {
s.segments[segmentID] = s.Clone(segment, SetDmlPositino(pos))
}
}
func (s *SegmentsInfo) SetStartPosition(segmentID UniqueID, pos *internalpb.MsgPosition) {
if segment, ok := s.segments[segmentID]; ok {
s.segments[segmentID] = s.Clone(segment, SetStartPosition(pos))
}
}
func (s *SegmentsInfo) Clone(segment *datapb.SegmentInfo, opts ...SegmentInfoOption) *datapb.SegmentInfo {
dmlPos := proto.Clone(segment.DmlPosition).(*internalpb.MsgPosition)
startPos := proto.Clone(segment.StartPosition).(*internalpb.MsgPosition)
cloned := &datapb.SegmentInfo{
ID: segment.ID,
CollectionID: segment.CollectionID,
PartitionID: segment.PartitionID,
InsertChannel: segment.InsertChannel,
NumOfRows: segment.NumOfRows,
State: segment.State,
DmlPosition: dmlPos,
MaxRowNum: segment.MaxRowNum,
LastExpireTime: segment.LastExpireTime,
StartPosition: startPos,
}
for _, opt := range opts {
opt(cloned)
}
return cloned
}
func (s *SegmentsInfo) ShadowClone(segment *datapb.SegmentInfo, opts ...SegmentInfoOption) *datapb.SegmentInfo {
cloned := &datapb.SegmentInfo{
ID: segment.ID,
CollectionID: segment.CollectionID,
PartitionID: segment.PartitionID,
InsertChannel: segment.InsertChannel,
NumOfRows: segment.NumOfRows,
State: segment.State,
DmlPosition: segment.DmlPosition,
MaxRowNum: segment.MaxRowNum,
LastExpireTime: segment.LastExpireTime,
StartPosition: segment.StartPosition,
}
for _, opt := range opts {
opt(cloned)
}
return cloned
}
type SegmentInfoOption func(segment *datapb.SegmentInfo)
func SetRowCount(rowCount int64) SegmentInfoOption {
return func(segment *datapb.SegmentInfo) {
segment.NumOfRows = rowCount
}
}
func SetExpireTime(expireTs Timestamp) SegmentInfoOption {
return func(segment *datapb.SegmentInfo) {
segment.LastExpireTime = expireTs
}
}
func SetState(state commonpb.SegmentState) SegmentInfoOption {
return func(segment *datapb.SegmentInfo) {
segment.State = state
}
}
func SetDmlPositino(pos *internalpb.MsgPosition) SegmentInfoOption {
return func(segment *datapb.SegmentInfo) {
segment.DmlPosition = pos
}
}
func SetStartPosition(pos *internalpb.MsgPosition) SegmentInfoOption {
return func(segment *datapb.SegmentInfo) {
segment.StartPosition = pos
}
}

34
internal/datacoord/segment_manager.go
View File

@ -236,8 +236,8 @@ func (s *SegmentManager) AllocSegment(ctx context.Context, collectionID UniqueID
var status *segmentStatus
var info *datapb.SegmentInfo
for _, segStatus := range s.stats {
info, err = s.meta.GetSegment(segStatus.id)
if err != nil {
info = s.meta.GetSegment(segStatus.id)
if info == nil {
log.Warn("Failed to get seginfo from meta", zap.Int64("id", segStatus.id), zap.Error(err))
continue
}
@ -260,8 +260,8 @@ func (s *SegmentManager) AllocSegment(ctx context.Context, collectionID UniqueID
if err != nil {
return
}
info, err = s.meta.GetSegment(status.id)
if err != nil {
info = s.meta.GetSegment(status.id)
if info == nil {
log.Warn("Failed to get seg into from meta", zap.Int64("id", status.id), zap.Error(err))
return
}
@ -368,9 +368,9 @@ func (s *SegmentManager) openNewSegment(ctx context.Context, collectionID Unique
}
func (s *SegmentManager) estimateMaxNumOfRows(collectionID UniqueID) (int, error) {
collMeta, err := s.meta.GetCollection(collectionID)
if err != nil {
return -1, err
collMeta := s.meta.GetCollection(collectionID)
if collMeta == nil {
return -1, fmt.Errorf("Failed to get collection %d", collectionID)
}
return s.estimatePolicy.apply(collMeta.Schema)
}
@ -396,9 +396,9 @@ func (s *SegmentManager) SealAllSegments(ctx context.Context, collectionID Uniqu
defer s.mu.Unlock()
ret := make([]UniqueID, 0)
for _, status := range s.stats {
info, err := s.meta.GetSegment(status.id)
if err != nil {
log.Warn("Failed to get seg info from meta", zap.Int64("id", status.id), zap.Error(err))
info := s.meta.GetSegment(status.id)
if info == nil {
log.Warn("Failed to get seg info from meta", zap.Int64("id", status.id))
continue
}
if info.CollectionID != collectionID {
@ -408,7 +408,7 @@ func (s *SegmentManager) SealAllSegments(ctx context.Context, collectionID Uniqu
ret = append(ret, status.id)
continue
}
if err := s.meta.SealSegment(status.id); err != nil {
if err := s.meta.SetState(status.id, commonpb.SegmentState_Sealed); err != nil {
return nil, err
}
ret = append(ret, status.id)
@ -486,9 +486,9 @@ func (s *SegmentManager) tryToSealSegment(ts Timestamp) error {
channelInfo := make(map[string][]*datapb.SegmentInfo)
mIDSegment := make(map[UniqueID]*datapb.SegmentInfo)
for _, status := range s.stats {
info, err := s.meta.GetSegment(status.id)
if err != nil {
log.Warn("Failed to get seg info from meta", zap.Int64("id", status.id), zap.Error(err))
info := s.meta.GetSegment(status.id)
if info == nil {
log.Warn("Failed to get seg info from meta", zap.Int64("id", status.id))
continue
}
mIDSegment[status.id] = info
@ -499,7 +499,7 @@ func (s *SegmentManager) tryToSealSegment(ts Timestamp) error {
// change shouldSeal to segment seal policy logic
for _, policy := range s.segmentSealPolicies {
if policy(status, info, ts) {
if err := s.meta.SealSegment(status.id); err != nil {
if err := s.meta.SetState(status.id, commonpb.SegmentState_Sealed); err != nil {
return err
}
break
@ -513,7 +513,7 @@ func (s *SegmentManager) tryToSealSegment(ts Timestamp) error {
if info.State == commonpb.SegmentState_Sealed {
continue
}
if err := s.meta.SealSegment(info.ID); err != nil {
if err := s.meta.SetState(info.GetID(), commonpb.SegmentState_Sealed); err != nil {
return err
}
}
@ -528,7 +528,7 @@ func (s *SegmentManager) SealSegment(ctx context.Context, segmentID UniqueID) er
defer sp.Finish()
s.mu.Lock()
defer s.mu.Unlock()
if err := s.meta.SealSegment(segmentID); err != nil {
if err := s.meta.SetState(segmentID, commonpb.SegmentState_Sealed); err != nil {
return err
}
return nil

24
internal/datacoord/segment_manager_test.go
View File

@ -31,11 +31,7 @@ func TestAllocSegment(t *testing.T) {
schema := newTestSchema()
collID, err := mockAllocator.allocID()
assert.Nil(t, err)
err = meta.AddCollection(&datapb.CollectionInfo{
ID: collID,
Schema: schema,
})
assert.Nil(t, err)
meta.AddCollection(&datapb.CollectionInfo{ID: collID, Schema: schema})
cases := []struct {
collectionID UniqueID
partitionID UniqueID
@ -68,11 +64,7 @@ func TestLoadSegmentsFromMeta(t *testing.T) {
schema := newTestSchema()
collID, err := mockAllocator.allocID()
assert.Nil(t, err)
err = meta.AddCollection(&datapb.CollectionInfo{
ID: collID,
Schema: schema,
})
assert.Nil(t, err)
meta.AddCollection(&datapb.CollectionInfo{ID: collID, Schema: schema})
sealedSegment := &datapb.SegmentInfo{
ID: 1,
@ -124,12 +116,7 @@ func TestSaveSegmentsToMeta(t *testing.T) {
schema := newTestSchema()
collID, err := mockAllocator.allocID()
assert.Nil(t, err)
err = meta.AddCollection(&datapb.CollectionInfo{
ID: collID,
Schema: schema,
})
assert.Nil(t, err)
meta.AddCollection(&datapb.CollectionInfo{ID: collID, Schema: schema})
segmentManager := newSegmentManager(meta, mockAllocator)
segID, _, expireTs, err := segmentManager.AllocSegment(context.Background(), collID, 0, "c1", 1000)
assert.Nil(t, err)
@ -137,9 +124,8 @@ func TestSaveSegmentsToMeta(t *testing.T) {
assert.NotNil(t, segStatus)
_, err = segmentManager.SealAllSegments(context.Background(), collID)
assert.Nil(t, err)
segment, err := meta.GetSegment(segID)
assert.Nil(t, err)
segment := meta.GetSegment(segID)
assert.NotNil(t, segment)
assert.EqualValues(t, segment.LastExpireTime, expireTs)
assert.EqualValues(t, commonpb.SegmentState_Sealed, segment.State)
}

13
internal/datacoord/server.go
View File

@ -340,8 +340,8 @@ func (s *Server) startDataNodeTtLoop(ctx context.Context) {
log.Debug("Flush segments", zap.Int64s("segmentIDs", segments))
segmentInfos := make([]*datapb.SegmentInfo, 0, len(segments))
for _, id := range segments {
sInfo, err := s.meta.GetSegment(id)
if err != nil {
sInfo := s.meta.GetSegment(id)
if sInfo == nil {
log.Error("get segment from meta error", zap.Int64("id", id),
zap.Error(err))
continue
@ -424,8 +424,8 @@ func (s *Server) startFlushLoop(ctx context.Context) {
log.Debug("flush loop shutdown")
return
case segmentID := <-s.flushCh:
segment, err := s.meta.GetSegment(segmentID)
if err != nil {
segment := s.meta.GetSegment(segmentID)
if segment == nil {
log.Warn("failed to get flused segment", zap.Int64("id", segmentID))
continue
}
@ -441,7 +441,7 @@ func (s *Server) startFlushLoop(ctx context.Context) {
continue
}
// set segment to SegmentState_Flushed
if err = s.meta.FlushSegment(segmentID); err != nil {
if err = s.meta.SetState(segmentID, commonpb.SegmentState_Flushed); err != nil {
log.Error("flush segment complete failed", zap.Error(err))
continue
}
@ -546,7 +546,8 @@ func (s *Server) loadCollectionFromRootCoord(ctx context.Context, collectionID i
Schema: resp.Schema,
Partitions: presp.PartitionIDs,
}
return s.meta.AddCollection(collInfo)
s.meta.AddCollection(collInfo)
return nil
}
func (s *Server) prepareBinlog(req *datapb.SaveBinlogPathsRequest) (map[string]string, error) {

19
internal/datacoord/server_test.go
View File

@ -109,12 +109,7 @@ func TestFlush(t *testing.T) {
svr := newTestServer(t, nil)
defer closeTestServer(t, svr)
schema := newTestSchema()
err := svr.meta.AddCollection(&datapb.CollectionInfo{
ID: 0,
Schema: schema,
Partitions: []int64{},
})
assert.Nil(t, err)
svr.meta.AddCollection(&datapb.CollectionInfo{ID: 0, Schema: schema, Partitions: []int64{}})
segID, _, expireTs, err := svr.segmentManager.AllocSegment(context.TODO(), 0, 1, "channel-1", 1)
assert.Nil(t, err)
resp, err := svr.Flush(context.TODO(), &datapb.FlushRequest{
@ -369,12 +364,11 @@ func TestSaveBinlogPaths(t *testing.T) {
}
for _, collection := range collections {
err := svr.meta.AddCollection(&datapb.CollectionInfo{
svr.meta.AddCollection(&datapb.CollectionInfo{
ID: collection.ID,
Schema: nil,
Partitions: collection.Partitions,
})
assert.Nil(t, err)
}
segments := []struct {
@ -447,8 +441,8 @@ func TestSaveBinlogPaths(t *testing.T) {
assert.EqualValues(t, "/by-dev/test/0/1/2/1/Allo2", metas[1].BinlogPath)
}
segmentInfo, err := svr.meta.GetSegment(0)
assert.Nil(t, err)
segmentInfo := svr.meta.GetSegment(0)
assert.NotNil(t, segmentInfo)
assert.EqualValues(t, segmentInfo.DmlPosition.ChannelName, "ch1")
assert.EqualValues(t, segmentInfo.DmlPosition.MsgID, []byte{1, 2, 3})
assert.EqualValues(t, segmentInfo.NumOfRows, 10)
@ -565,12 +559,11 @@ func TestGetVChannelPos(t *testing.T) {
svr := newTestServer(t, nil)
defer closeTestServer(t, svr)
schema := newTestSchema()
err := svr.meta.AddCollection(&datapb.CollectionInfo{
svr.meta.AddCollection(&datapb.CollectionInfo{
ID: 0,
Schema: schema,
})
assert.Nil(t, err)
err = svr.meta.AddSegment(&datapb.SegmentInfo{
err := svr.meta.AddSegment(&datapb.SegmentInfo{
ID: 1,
CollectionID: 0,
PartitionID: 0,