milvus/internal/datanode/collection_replica.go
XuanYang-cn 377ee8687d Modify flush logic
Signed-off-by: XuanYang-cn <xuan.yang@zilliz.com>
2021-01-22 09:36:40 +08:00

224 lines
6.1 KiB
Go

package datanode
import (
"log"
"sync"
"github.com/zilliztech/milvus-distributed/internal/errors"
"github.com/zilliztech/milvus-distributed/internal/proto/internalpb2"
)
type collectionReplica interface {
// collection
getCollectionNum() int
addCollection(collectionID UniqueID, schemaBlob string) error
removeCollection(collectionID UniqueID) error
getCollectionByID(collectionID UniqueID) (*Collection, error)
getCollectionByName(collectionName string) (*Collection, error)
getCollectionIDByName(collectionName string) (UniqueID, error)
hasCollection(collectionID UniqueID) bool
// segment
addSegment(segmentID UniqueID, collName string, partitionName string) error
removeSegment(segmentID UniqueID) error
hasSegment(segmentID UniqueID) bool
updateSegmentRowNums(segmentID UniqueID, numRows int64) error
getSegmentStatisticsUpdates(segmentID UniqueID) (*internalpb2.SegmentStatisticsUpdates, error)
getSegmentByID(segmentID UniqueID) (*Segment, error)
}
type (
Segment struct {
segmentID UniqueID
collectionID UniqueID
partitionName string
numRows int64
memorySize int64
}
collectionReplicaImpl struct {
mu sync.RWMutex
collections []*Collection
segments []*Segment
}
)
//----------------------------------------------------------------------------------------------------- collection
func (colReplica *collectionReplicaImpl) getSegmentByID(segmentID UniqueID) (*Segment, error) {
// GOOSE TODO: read write lock
colReplica.mu.RLock()
defer colReplica.mu.RUnlock()
for _, segment := range colReplica.segments {
if segment.segmentID == segmentID {
return segment, nil
}
}
return nil, errors.Errorf("cannot find segment, id = %v", segmentID)
}
func (colReplica *collectionReplicaImpl) addSegment(segmentID UniqueID, collName string, partitionName string) error {
colReplica.mu.Lock()
defer colReplica.mu.Unlock()
log.Println("Add Segment", segmentID)
collID, err := colReplica.getCollectionIDByName(collName)
if err != nil {
return err
}
seg := &Segment{
segmentID: segmentID,
collectionID: collID,
partitionName: partitionName,
}
colReplica.segments = append(colReplica.segments, seg)
return nil
}
func (colReplica *collectionReplicaImpl) removeSegment(segmentID UniqueID) error {
colReplica.mu.Lock()
defer colReplica.mu.Unlock()
for index, ele := range colReplica.segments {
if ele.segmentID == segmentID {
log.Println("Removing segment:", segmentID)
numOfSegs := len(colReplica.segments)
colReplica.segments[index] = colReplica.segments[numOfSegs-1]
colReplica.segments = colReplica.segments[:numOfSegs-1]
return nil
}
}
return errors.Errorf("Error, there's no segment %v", segmentID)
}
func (colReplica *collectionReplicaImpl) hasSegment(segmentID UniqueID) bool {
colReplica.mu.RLock()
defer colReplica.mu.RUnlock()
for _, ele := range colReplica.segments {
if ele.segmentID == segmentID {
return true
}
}
return false
}
func (colReplica *collectionReplicaImpl) updateSegmentRowNums(segmentID UniqueID, numRows int64) error {
colReplica.mu.Lock()
defer colReplica.mu.Unlock()
for _, ele := range colReplica.segments {
if ele.segmentID == segmentID {
log.Printf("updating segment(%v) row nums: (%v)", segmentID, numRows)
ele.memorySize = 0
ele.numRows += numRows
return nil
}
}
return errors.Errorf("Error, there's no segment %v", segmentID)
}
func (colReplica *collectionReplicaImpl) getSegmentStatisticsUpdates(segmentID UniqueID) (*internalpb2.SegmentStatisticsUpdates, error) {
colReplica.mu.RLock()
defer colReplica.mu.RUnlock()
for _, ele := range colReplica.segments {
if ele.segmentID == segmentID {
updates := &internalpb2.SegmentStatisticsUpdates{
SegmentID: segmentID,
MemorySize: ele.memorySize,
NumRows: ele.numRows,
}
return updates, nil
}
}
return nil, errors.Errorf("Error, there's no segment %v", segmentID)
}
func (colReplica *collectionReplicaImpl) getCollectionNum() int {
colReplica.mu.RLock()
defer colReplica.mu.RUnlock()
return len(colReplica.collections)
}
func (colReplica *collectionReplicaImpl) addCollection(collectionID UniqueID, schemaBlob string) error {
colReplica.mu.Lock()
defer colReplica.mu.Unlock()
var newCollection = newCollection(collectionID, schemaBlob)
colReplica.collections = append(colReplica.collections, newCollection)
log.Println("Create collection: ", newCollection.Name())
return nil
}
func (colReplica *collectionReplicaImpl) getCollectionIDByName(collName string) (UniqueID, error) {
colReplica.mu.RLock()
defer colReplica.mu.RUnlock()
for _, collection := range colReplica.collections {
if collection.Name() == collName {
return collection.ID(), nil
}
}
return 0, errors.Errorf("There is no collection name=%v", collName)
}
func (colReplica *collectionReplicaImpl) removeCollection(collectionID UniqueID) error {
// GOOSE TODO: optimize
colReplica.mu.Lock()
defer colReplica.mu.Unlock()
tmpCollections := make([]*Collection, 0)
for _, col := range colReplica.collections {
if col.ID() != collectionID {
tmpCollections = append(tmpCollections, col)
} else {
log.Println("Drop collection : ", col.Name())
}
}
colReplica.collections = tmpCollections
return nil
}
func (colReplica *collectionReplicaImpl) getCollectionByID(collectionID UniqueID) (*Collection, error) {
colReplica.mu.RLock()
defer colReplica.mu.RUnlock()
for _, collection := range colReplica.collections {
if collection.ID() == collectionID {
return collection, nil
}
}
return nil, errors.Errorf("cannot find collection, id = %v", collectionID)
}
func (colReplica *collectionReplicaImpl) getCollectionByName(collectionName string) (*Collection, error) {
colReplica.mu.RLock()
defer colReplica.mu.RUnlock()
for _, collection := range colReplica.collections {
if collection.Name() == collectionName {
return collection, nil
}
}
return nil, errors.Errorf("Cannot found collection: %v", collectionName)
}
func (colReplica *collectionReplicaImpl) hasCollection(collectionID UniqueID) bool {
colReplica.mu.RLock()
defer colReplica.mu.RUnlock()
for _, col := range colReplica.collections {
if col.ID() == collectionID {
return true
}
}
return false
}