mirror of https://github.com/milvus-io/milvus.git
84 lines
2.1 KiB
Go
84 lines
2.1 KiB
Go
package indexnode
|
|
|
|
import (
|
|
"github.com/golang/protobuf/proto"
|
|
"github.com/milvus-io/milvus/internal/proto/commonpb"
|
|
"github.com/milvus-io/milvus/internal/proto/indexnodepb"
|
|
)
|
|
|
|
func (i *IndexNode) loadOrStoreTask(clusterID, buildID UniqueID, info *taskInfo) *taskInfo {
|
|
i.stateLock.Lock()
|
|
defer i.stateLock.Unlock()
|
|
key := taskKey{ClusterID: clusterID, BuildID: buildID}
|
|
oldInfo, ok := i.tasks[key]
|
|
if ok {
|
|
return oldInfo
|
|
}
|
|
i.tasks[key] = info
|
|
return nil
|
|
}
|
|
|
|
func (i *IndexNode) loadTaskState(clusterID, buildID UniqueID) (commonpb.IndexState, bool) {
|
|
key := taskKey{ClusterID: clusterID, BuildID: buildID}
|
|
i.stateLock.Lock()
|
|
defer i.stateLock.Unlock()
|
|
task, ok := i.tasks[key]
|
|
return task.state, ok
|
|
}
|
|
|
|
func (i *IndexNode) storeTaskState(clusterID, buildID UniqueID, state commonpb.IndexState) {
|
|
key := taskKey{ClusterID: clusterID, BuildID: buildID}
|
|
i.stateLock.Lock()
|
|
defer i.stateLock.Unlock()
|
|
if task, ok := i.tasks[key]; ok {
|
|
task.state = state
|
|
}
|
|
}
|
|
|
|
func (i *IndexNode) foreachTaskInfo(fn func(clusterID, buildID UniqueID, info *taskInfo)) {
|
|
i.stateLock.Lock()
|
|
defer i.stateLock.Unlock()
|
|
for key, info := range i.tasks {
|
|
fn(key.ClusterID, key.BuildID, info)
|
|
}
|
|
}
|
|
|
|
func (i *IndexNode) storeIndexFilesAndStatistic(clusterID, buildID UniqueID, files []string, statistic *indexnodepb.JobInfo) {
|
|
key := taskKey{ClusterID: clusterID, BuildID: buildID}
|
|
i.stateLock.Lock()
|
|
defer i.stateLock.Unlock()
|
|
if info, ok := i.tasks[key]; !ok {
|
|
return
|
|
} else {
|
|
info.indexfiles = files[:]
|
|
info.statistic = proto.Clone(statistic).(*indexnodepb.JobInfo)
|
|
}
|
|
}
|
|
|
|
func (i *IndexNode) deleteTaskInfos(keys []taskKey) []*taskInfo {
|
|
i.stateLock.Lock()
|
|
defer i.stateLock.Unlock()
|
|
deleted := make([]*taskInfo, 0, len(keys))
|
|
for _, key := range keys {
|
|
info, ok := i.tasks[key]
|
|
if ok {
|
|
deleted = append(deleted, info)
|
|
delete(i.tasks, key)
|
|
}
|
|
}
|
|
return deleted
|
|
}
|
|
|
|
func (i *IndexNode) deleteAllTasks() []*taskInfo {
|
|
i.stateLock.Lock()
|
|
deletedTasks := i.tasks
|
|
i.tasks = make(map[taskKey]*taskInfo)
|
|
i.stateLock.Unlock()
|
|
|
|
deleted := make([]*taskInfo, 0, len(deletedTasks))
|
|
for _, info := range deletedTasks {
|
|
deleted = append(deleted, info)
|
|
}
|
|
return deleted
|
|
}
|