refine index node (#18545)

Signed-off-by: Zach41 <[email protected]>
milvus-io · Aug 9, 2022 · d3c478f · d3c478f
1 parent 4edc8d3
commit d3c478f
Show file tree

Hide file tree

Showing 24 changed files with 3,219 additions and 2,225 deletions.
diff --git a/internal/core/src/pb/common.pb.cc b/internal/core/src/pb/common.pb.cc
diff --git a/internal/core/src/pb/common.pb.h b/internal/core/src/pb/common.pb.h
diff --git a/internal/distributed/indexnode/service.go b/internal/distributed/indexnode/service.go
@@ -35,6 +35,7 @@ import (
 	"github.com/milvus-io/milvus/internal/indexnode"
 	"github.com/milvus-io/milvus/internal/log"
 	"github.com/milvus-io/milvus/internal/proto/commonpb"
+	"github.com/milvus-io/milvus/internal/proto/indexnodepb"
 	"github.com/milvus-io/milvus/internal/proto/indexpb"
 	"github.com/milvus-io/milvus/internal/proto/internalpb"
 	"github.com/milvus-io/milvus/internal/proto/milvuspb"
@@ -232,13 +233,24 @@ func (s *Server) GetStatisticsChannel(ctx context.Context, req *internalpb.GetSt
 	return s.indexnode.GetStatisticsChannel(ctx)
 }
 
-// CreateIndex sends the create index request to IndexNode.
-func (s *Server) CreateIndex(ctx context.Context, req *indexpb.CreateIndexRequest) (*commonpb.Status, error) {
-	return s.indexnode.CreateIndex(ctx, req)
+// CreateJob sends the create index request to IndexNode.
+func (s *Server) CreateJob(ctx context.Context, req *indexnodepb.CreateJobRequest) (*commonpb.Status, error) {
+	return s.indexnode.CreateJob(ctx, req)
 }
 
-func (s *Server) GetTaskSlots(ctx context.Context, req *indexpb.GetTaskSlotsRequest) (*indexpb.GetTaskSlotsResponse, error) {
-	return s.indexnode.GetTaskSlots(ctx, req)
+// QueryJobs querys index jobs statues
+func (s *Server) QueryJobs(ctx context.Context, req *indexnodepb.QueryJobsRequest) (*indexnodepb.QueryJobsRespond, error) {
+	return s.indexnode.QueryJobs(ctx, req)
+}
+
+// DropJobs drops index build jobs
+func (s *Server) DropJobs(ctx context.Context, req *indexnodepb.DropJobsRequest) (*commonpb.Status, error) {
+	return s.indexnode.DropJobs(ctx, req)
+}
+
+// GetJobNum gets indexnode's job statisctics
+func (s *Server) GetJobStats(ctx context.Context, req *indexnodepb.GetJobStatsRequest) (*indexnodepb.GetJobStatsRespond, error) {
+	return s.indexnode.GetJobStats(ctx, req)
 }
 
 // GetMetrics gets the metrics info of IndexNode.

diff --git a/internal/indexnode/chunk_mgr_factory.go b/internal/indexnode/chunk_mgr_factory.go
@@ -0,0 +1,39 @@
+package indexnode
+
+import (
+	"context"
+	"fmt"
+	"sync"
+
+	"github.com/milvus-io/milvus/internal/storage"
+)
+
+type StorageFactory interface {
+	NewChunkManager(ctx context.Context, bucket, storageAccessKey string) (storage.ChunkManager, error)
+}
+
+type chunkMgr struct {
+	cached sync.Map
+}
+
+func (m *chunkMgr) NewChunkManager(ctx context.Context, bucket, storageAccessKey string) (storage.ChunkManager, error) {
+	key := m.cacheKey(bucket, storageAccessKey)
+	if v, ok := m.cached.Load(key); ok {
+		return v.(storage.ChunkManager), nil
+	}
+	opts := []storage.Option{
+		storage.AccessKeyID(storageAccessKey),
+		storage.BucketName(bucket),
+	}
+	factory := storage.NewChunkManagerFactory("local", "minio", opts...)
+	mgr, err := factory.NewVectorStorageChunkManager(ctx)
+	if err != nil {
+		return nil, err
+	}
+	v, _ := m.cached.LoadOrStore(key, mgr)
+	return v.(storage.ChunkManager), nil
+}
+
+func (m *chunkMgr) cacheKey(bucket, storageAccessKey string) string {
+	return fmt.Sprintf("%s/%s", bucket, storageAccessKey)
+}