mirror of https://github.com/milvus-io/milvus.git
278 lines
9.2 KiB
Go
278 lines
9.2 KiB
Go
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
|
// with the License. You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
|
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
|
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
|
|
|
package queryservice
|
|
|
|
import (
|
|
"errors"
|
|
|
|
internalPb "github.com/milvus-io/milvus/internal/proto/internalpb"
|
|
"github.com/milvus-io/milvus/internal/proto/querypb"
|
|
"github.com/milvus-io/milvus/internal/proto/schemapb"
|
|
)
|
|
|
|
type Replica interface {
|
|
getCollections(dbID UniqueID) ([]*collection, error)
|
|
getPartitions(dbID UniqueID, collectionID UniqueID) ([]*partition, error)
|
|
getSegments(dbID UniqueID, collectionID UniqueID, partitionID UniqueID) ([]*segment, error)
|
|
getCollectionByID(dbID UniqueID, collectionID UniqueID) (*collection, error)
|
|
getPartitionByID(dbID UniqueID, collectionID UniqueID, partitionID UniqueID) (*partition, error)
|
|
addCollection(dbID UniqueID, collectionID UniqueID, schema *schemapb.CollectionSchema) error
|
|
addPartition(dbID UniqueID, collectionID UniqueID, partitionID UniqueID) error
|
|
updatePartitionState(dbID UniqueID, collectionID UniqueID, partitionID UniqueID, state querypb.PartitionState) error
|
|
getPartitionStates(dbID UniqueID, collectionID UniqueID, partitionIDs []UniqueID) ([]*querypb.PartitionStates, error)
|
|
releaseCollection(dbID UniqueID, collectionID UniqueID) error
|
|
releasePartition(dbID UniqueID, collectionID UniqueID, partitionID UniqueID) error
|
|
addDmChannel(dbID UniqueID, collectionID UniqueID, channel string, watchedStartPos *internalPb.MsgPosition) error
|
|
addExcludeSegmentIDs(dbID UniqueID, collectionID UniqueID, excludeSegments []UniqueID) error
|
|
//getAssignedNodeIDByChannelName(dbID UniqueID, collectionID UniqueID, channel string) (int64, error)
|
|
}
|
|
|
|
type segment struct {
|
|
id UniqueID
|
|
}
|
|
|
|
type partition struct {
|
|
id UniqueID
|
|
segments map[UniqueID]*segment
|
|
state querypb.PartitionState
|
|
}
|
|
|
|
type collection struct {
|
|
id UniqueID
|
|
partitions map[UniqueID]*partition
|
|
dmChannels []string
|
|
dmChannels2Pos map[string]*internalPb.MsgPosition
|
|
excludeSegmentIds []UniqueID
|
|
schema *schemapb.CollectionSchema
|
|
}
|
|
|
|
type metaReplica struct {
|
|
dbID []UniqueID
|
|
db2collections map[UniqueID][]*collection
|
|
}
|
|
|
|
func newMetaReplica() Replica {
|
|
db2collections := make(map[UniqueID][]*collection)
|
|
db2collections[0] = make([]*collection, 0)
|
|
dbIDs := make([]UniqueID, 0)
|
|
dbIDs = append(dbIDs, UniqueID(0))
|
|
return &metaReplica{
|
|
dbID: dbIDs,
|
|
db2collections: db2collections,
|
|
}
|
|
}
|
|
|
|
func (mp *metaReplica) addCollection(dbID UniqueID, collectionID UniqueID, schema *schemapb.CollectionSchema) error {
|
|
//TODO:: assert dbID = 0 exist
|
|
if _, ok := mp.db2collections[dbID]; ok {
|
|
partitions := make(map[UniqueID]*partition)
|
|
channels := make([]string, 0)
|
|
startPos := make(map[string]*internalPb.MsgPosition)
|
|
excludeSegmentIDs := make([]UniqueID, 0)
|
|
newCollection := &collection{
|
|
id: collectionID,
|
|
partitions: partitions,
|
|
schema: schema,
|
|
dmChannels: channels,
|
|
dmChannels2Pos: startPos,
|
|
excludeSegmentIds: excludeSegmentIDs,
|
|
}
|
|
mp.db2collections[dbID] = append(mp.db2collections[dbID], newCollection)
|
|
return nil
|
|
}
|
|
|
|
return errors.New("addCollection: can't find dbID when add collection")
|
|
}
|
|
|
|
func (mp *metaReplica) addPartition(dbID UniqueID, collectionID UniqueID, partitionID UniqueID) error {
|
|
if collections, ok := mp.db2collections[dbID]; ok {
|
|
for _, collection := range collections {
|
|
if collection.id == collectionID {
|
|
partitions := collection.partitions
|
|
segments := make(map[UniqueID]*segment)
|
|
partition := &partition{
|
|
id: partitionID,
|
|
state: querypb.PartitionState_NotPresent,
|
|
segments: segments,
|
|
}
|
|
partitions[partitionID] = partition
|
|
return nil
|
|
}
|
|
}
|
|
}
|
|
return errors.New("addPartition: can't find collection when add partition")
|
|
}
|
|
|
|
func (mp *metaReplica) getCollections(dbID UniqueID) ([]*collection, error) {
|
|
if collections, ok := mp.db2collections[dbID]; ok {
|
|
return collections, nil
|
|
}
|
|
|
|
return nil, errors.New("getCollections: can't find collectionID")
|
|
}
|
|
|
|
func (mp *metaReplica) getPartitions(dbID UniqueID, collectionID UniqueID) ([]*partition, error) {
|
|
if collections, ok := mp.db2collections[dbID]; ok {
|
|
for _, collection := range collections {
|
|
if collectionID == collection.id {
|
|
partitions := make([]*partition, 0)
|
|
for _, partition := range collection.partitions {
|
|
partitions = append(partitions, partition)
|
|
}
|
|
return partitions, nil
|
|
}
|
|
}
|
|
}
|
|
|
|
return nil, errors.New("getPartitions: can't find partitionIDs")
|
|
}
|
|
|
|
func (mp *metaReplica) getSegments(dbID UniqueID, collectionID UniqueID, partitionID UniqueID) ([]*segment, error) {
|
|
if collections, ok := mp.db2collections[dbID]; ok {
|
|
for _, collection := range collections {
|
|
if collectionID == collection.id {
|
|
if partition, ok := collection.partitions[partitionID]; ok {
|
|
segments := make([]*segment, 0)
|
|
for _, segment := range partition.segments {
|
|
segments = append(segments, segment)
|
|
}
|
|
return segments, nil
|
|
}
|
|
}
|
|
}
|
|
}
|
|
return nil, errors.New("getSegments: can't find segmentID")
|
|
}
|
|
|
|
func (mp *metaReplica) getCollectionByID(dbID UniqueID, collectionID UniqueID) (*collection, error) {
|
|
if collections, ok := mp.db2collections[dbID]; ok {
|
|
for _, collection := range collections {
|
|
if collectionID == collection.id {
|
|
return collection, nil
|
|
}
|
|
}
|
|
}
|
|
|
|
return nil, errors.New("getCollectionByID: can't find collectionID")
|
|
}
|
|
|
|
func (mp *metaReplica) getPartitionByID(dbID UniqueID, collectionID UniqueID, partitionID UniqueID) (*partition, error) {
|
|
if collections, ok := mp.db2collections[dbID]; ok {
|
|
for _, collection := range collections {
|
|
if collectionID == collection.id {
|
|
partitions := collection.partitions
|
|
if partition, ok := partitions[partitionID]; ok {
|
|
return partition, nil
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
return nil, errors.New("getPartitionByID: can't find partitionID")
|
|
}
|
|
|
|
func (mp *metaReplica) updatePartitionState(dbID UniqueID,
|
|
collectionID UniqueID,
|
|
partitionID UniqueID,
|
|
state querypb.PartitionState) error {
|
|
for _, collection := range mp.db2collections[dbID] {
|
|
if collection.id == collectionID {
|
|
if partition, ok := collection.partitions[partitionID]; ok {
|
|
partition.state = state
|
|
return nil
|
|
}
|
|
}
|
|
}
|
|
return errors.New("updatePartitionState: update partition state fail")
|
|
}
|
|
|
|
func (mp *metaReplica) getPartitionStates(dbID UniqueID,
|
|
collectionID UniqueID,
|
|
partitionIDs []UniqueID) ([]*querypb.PartitionStates, error) {
|
|
partitionStates := make([]*querypb.PartitionStates, 0)
|
|
for _, collection := range mp.db2collections[dbID] {
|
|
if collection.id == collectionID {
|
|
for _, partitionID := range partitionIDs {
|
|
if partition, ok := collection.partitions[partitionID]; ok {
|
|
partitionStates = append(partitionStates, &querypb.PartitionStates{
|
|
PartitionID: partitionID,
|
|
State: partition.state,
|
|
})
|
|
} else {
|
|
partitionStates = append(partitionStates, &querypb.PartitionStates{
|
|
PartitionID: partitionID,
|
|
State: querypb.PartitionState_NotPresent,
|
|
})
|
|
}
|
|
}
|
|
}
|
|
}
|
|
return partitionStates, nil
|
|
}
|
|
|
|
func (mp *metaReplica) releaseCollection(dbID UniqueID, collectionID UniqueID) error {
|
|
if collections, ok := mp.db2collections[dbID]; ok {
|
|
for i, coll := range collections {
|
|
if collectionID == coll.id {
|
|
newSize := len(collections) - 1
|
|
newColls := make([]*collection, 0, newSize)
|
|
collections[i] = collections[newSize]
|
|
newColls = append(newColls, collections[:newSize]...)
|
|
mp.db2collections[dbID] = newColls
|
|
return nil
|
|
}
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (mp *metaReplica) releasePartition(dbID UniqueID, collectionID UniqueID, partitionID UniqueID) error {
|
|
if collections, ok := mp.db2collections[dbID]; ok {
|
|
for _, collection := range collections {
|
|
if collectionID == collection.id {
|
|
if _, ok := collection.partitions[partitionID]; ok {
|
|
delete(collection.partitions, partitionID)
|
|
return nil
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (mp *metaReplica) addDmChannel(dbID UniqueID, collectionID UniqueID, channel string, watchedStartPos *internalPb.MsgPosition) error {
|
|
if collections, ok := mp.db2collections[dbID]; ok {
|
|
for _, collection := range collections {
|
|
if collectionID == collection.id {
|
|
collection.dmChannels = append(collection.dmChannels, channel)
|
|
collection.dmChannels2Pos[channel] = watchedStartPos
|
|
return nil
|
|
}
|
|
}
|
|
}
|
|
return errors.New("addDmChannels: can't find dbID or collectionID")
|
|
}
|
|
|
|
func (mp *metaReplica) addExcludeSegmentIDs(dbID UniqueID, collectionID UniqueID, excludeSegments []UniqueID) error {
|
|
if collections, ok := mp.db2collections[dbID]; ok {
|
|
for _, collection := range collections {
|
|
if collectionID == collection.id {
|
|
collection.excludeSegmentIds = append(collection.excludeSegmentIds, excludeSegments...)
|
|
return nil
|
|
}
|
|
}
|
|
}
|
|
return errors.New("addExcludeSegmentIDs: can't find dbID or collectionID")
|
|
}
|