From 5afdf3d49f475e39056bd23e5165d65dfd87e225 Mon Sep 17 00:00:00 2001 From: Ten Thousand Leaves <69466447+soothing-rain@users.noreply.github.com> Date: Mon, 9 May 2022 20:47:52 +0800 Subject: [PATCH] Move locks back to meta table methods (#16848) /kind improvement Signed-off-by: Yuchen Gao --- internal/rootcoord/meta_table.go | 13 ++++++++ internal/rootcoord/meta_table_test.go | 8 +++++ internal/rootcoord/root_coord.go | 48 +++++++++++++-------------- internal/rootcoord/root_coord_test.go | 2 +- 4 files changed, 45 insertions(+), 26 deletions(-) diff --git a/internal/rootcoord/meta_table.go b/internal/rootcoord/meta_table.go index 3241c52c71..1c65ac3464 100644 --- a/internal/rootcoord/meta_table.go +++ b/internal/rootcoord/meta_table.go @@ -408,6 +408,19 @@ func (mt *MetaTable) HasCollection(collID typeutil.UniqueID, ts typeutil.Timesta return err == nil } +// GetCollectionIDByName returns the collection ID according to its name. +// Returns an error if no matching ID is found. +func (mt *MetaTable) GetCollectionIDByName(cName string) (typeutil.UniqueID, error) { + mt.ddLock.RLock() + defer mt.ddLock.RUnlock() + var cID UniqueID + var ok bool + if cID, ok = mt.collName2ID[cName]; !ok { + return 0, fmt.Errorf("collection ID not found for collection name %s", cName) + } + return cID, nil +} + // GetCollectionByID return collection meta by collection id func (mt *MetaTable) GetCollectionByID(collectionID typeutil.UniqueID, ts typeutil.Timestamp) (*pb.CollectionInfo, error) { mt.ddLock.RLock() diff --git a/internal/rootcoord/meta_table_test.go b/internal/rootcoord/meta_table_test.go index 20be0c69bc..6c9dc5942e 100644 --- a/internal/rootcoord/meta_table_test.go +++ b/internal/rootcoord/meta_table_test.go @@ -1312,6 +1312,14 @@ func TestMetaWithTimestamp(t *testing.T) { assert.NotNil(t, err) _, err = mt.GetPartitionByName(2, partName2, tsoStart) assert.NotNil(t, err) + + var cID UniqueID + cID, err = mt.GetCollectionIDByName(collName1) + assert.NoError(t, err) + assert.Equal(t, collID1, cID) + + _, err = mt.GetCollectionIDByName("badname") + assert.Error(t, err) } func TestFixIssue10540(t *testing.T) { diff --git a/internal/rootcoord/root_coord.go b/internal/rootcoord/root_coord.go index 0e60805e2e..77a8293cc5 100644 --- a/internal/rootcoord/root_coord.go +++ b/internal/rootcoord/root_coord.go @@ -2277,18 +2277,19 @@ func (c *Core) Import(ctx context.Context, req *milvuspb.ImportRequest) (*milvus } // Get collection/partition ID from collection/partition name. - var cID int64 - var ok bool - c.MetaTable.ddLock.RLock() - defer c.MetaTable.ddLock.RUnlock() - if cID, ok = c.MetaTable.collName2ID[req.GetCollectionName()]; !ok { - log.Error("failed to find collection ID for collection name", - zap.String("collection name", req.GetCollectionName())) - return nil, fmt.Errorf("collection ID not found for collection name %s", req.GetCollectionName()) - } - var pID int64 + var cID UniqueID var err error + if cID, err = c.MetaTable.GetCollectionIDByName(req.GetCollectionName()); err != nil { + log.Error("failed to find collection ID from its name", + zap.String("collection name", req.GetCollectionName()), + zap.Error(err)) + return nil, err + } + var pID UniqueID if pID, err = c.MetaTable.getPartitionByName(cID, req.GetPartitionName(), 0); err != nil { + log.Error("failed to get partition ID from its name", + zap.String("partition name", req.GetPartitionName()), + zap.Error(err)) return nil, err } log.Info("receive import request", @@ -2357,22 +2358,19 @@ func (c *Core) ReportImport(ctx context.Context, ir *rootcoordpb.ImportResult) ( }, nil } - // Reverse look up collection name on collection ID. + // Look up collection name on collection ID. var colName string - c.MetaTable.ddLock.RLock() - defer c.MetaTable.ddLock.RUnlock() - for k, v := range c.MetaTable.collName2ID { - if v == ti.GetCollectionId() { - colName = k - } - } - if colName == "" { - log.Error("Collection name not found for collection ID", zap.Int64("collection ID", ti.GetCollectionId())) + var colMeta *etcdpb.CollectionInfo + if colMeta, err = c.MetaTable.GetCollectionByID(ti.GetCollectionId(), 0); err != nil { + log.Error("failed to get collection name", + zap.Int64("collection ID", ti.GetCollectionId()), + zap.Error(err)) return &commonpb.Status{ ErrorCode: commonpb.ErrorCode_CollectionNameNotFound, - Reason: "Collection name not found for collection ID" + strconv.FormatInt(ti.GetCollectionId(), 10), + Reason: "failed to get collection name for collection ID" + strconv.FormatInt(ti.GetCollectionId(), 10), }, nil } + colName = colMeta.GetSchema().GetName() // When DataNode has done its thing, remove it from the busy node list. func() { @@ -2490,10 +2488,10 @@ func (c *Core) postImportPersistLoop(ctx context.Context, taskID int64, colID in c.wg.Add(1) c.checkSegmentLoadedLoop(ctx, taskID, colID, segIDs) // Check if collection has any indexed fields. If so, start a loop to check segments' index states. - c.MetaTable.ddLock.RLock() - defer c.MetaTable.ddLock.RUnlock() - colMeta := c.MetaTable.collID2Meta[colID] - if len(colMeta.GetFieldIndexes()) != 0 { + if colMeta, err := c.MetaTable.GetCollectionByID(colID, 0); err != nil { + log.Error("failed to find meta for collection", + zap.Int64("collection ID", colID)) + } else if len(colMeta.GetFieldIndexes()) != 0 { c.wg.Add(1) c.checkCompleteIndexLoop(ctx, taskID, colID, colName, segIDs) } diff --git a/internal/rootcoord/root_coord_test.go b/internal/rootcoord/root_coord_test.go index d7a64b403a..1b01f1a344 100644 --- a/internal/rootcoord/root_coord_test.go +++ b/internal/rootcoord/root_coord_test.go @@ -1362,7 +1362,7 @@ func TestRootCoord_Base(t *testing.T) { }) wg.Add(1) - t.Run("import w/ collection ID not found", func(t *testing.T) { + t.Run("import with collection ID not found", func(t *testing.T) { defer wg.Done() req := &milvuspb.ImportRequest{ CollectionName: "bad name",