Repository: milvus-io/milvus
File: internal/proxy/task.go (Go)

Code before:

	if err := validateCollectionName(collName); err != nil {
		return err
	}

	// To compat with LoadCollcetion before Milvus@2.1
	if lct.ReplicaNumber == 0 {
		lct.ReplicaNumber = 1
	}

	return nil
}

func (lct *loadCollectionTask) Execute(ctx context.Context) (err error) {
	log.Debug("loadCollectionTask Execute", zap.String("role", typeutil.ProxyRole), zap.Int64("msgID", lct.Base.MsgID))
	collID, err := globalMetaCache.GetCollectionID(ctx, lct.CollectionName)
	if err != nil {
		return err
	}
	lct.collectionID = collID
	collSchema, err := globalMetaCache.GetCollectionSchema(ctx, lct.CollectionName)
	if err != nil {
		return err
	}

	request := &querypb.LoadCollectionRequest{
		Base: &commonpb.MsgBase{
			MsgType:   commonpb.MsgType_LoadCollection,
			MsgID:     lct.Base.MsgID,
			Timestamp: lct.Base.Timestamp,
			SourceID:  lct.Base.SourceID,
		},
		DbID:          0,
		CollectionID:  collID,
		Schema:        collSchema,
		ReplicaNumber: lct.ReplicaNumber,
	}
	log.Debug("send LoadCollectionRequest to query coordinator", zap.String("role", typeutil.ProxyRole),
		zap.Int64("msgID", request.Base.MsgID), zap.Int64("collectionID", request.CollectionID),
		zap.Any("schema", request.Schema))
	lct.result, err = lct.queryCoord.LoadCollection(ctx, request)
	if err != nil {
		return fmt.Errorf("call query coordinator LoadCollection: %s", err)
	}
	return nil
}

func (lct *loadCollectionTask) PostExecute(ctx context.Context) error {
	log.Debug("loadCollectionTask PostExecute", zap.String("role", typeutil.ProxyRole),
		zap.Int64("msgID", lct.Base.MsgID))
	return nil

Review comment (suggestion):

how about provide a method called CheckCollectionLoaded so you don't have to worry about thread safety

Code after:

	if err := validateCollectionName(collName); err != nil {
		return err
	}

	// To compat with LoadCollcetion before Milvus@2.1
	if lct.ReplicaNumber == 0 {
		lct.ReplicaNumber = 1
	}

	return nil
}

func (lct *loadCollectionTask) Execute(ctx context.Context) (err error) {
	log.Debug("loadCollectionTask Execute", zap.String("role", typeutil.ProxyRole), zap.Int64("msgID", lct.Base.MsgID))
	collID, err := globalMetaCache.GetCollectionID(ctx, lct.CollectionName)
	if err != nil {
		return err
	}

	// check if loaded
	info, err := globalMetaCache.GetCollectionInfo(ctx, lct.CollectionName)
	if err != nil {
		return fmt.Errorf("GetCollectionInfo failed, collection = %s, err = %s", lct.CollectionName, err)
	}
	if info.isLoaded { // TODO: may not be thread safe
		log.Debug("collection has been loaded",
			zap.String("role", typeutil.ProxyRole),
			zap.String("collectionName", lct.CollectionName),
			zap.Int64("collectionID", collID),
			zap.Int64("msgID", lct.GetBase().GetMsgID()))
		lct.result = &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_Success,
		}
		return nil
	}

	lct.collectionID = collID
	collSchema, err := globalMetaCache.GetCollectionSchema(ctx, lct.CollectionName)
	if err != nil {
		return err
	}
	request := &querypb.LoadCollectionRequest{
		Base: &commonpb.MsgBase{
			MsgType:   commonpb.MsgType_LoadCollection,
			MsgID:     lct.Base.MsgID,
			Timestamp: lct.Base.Timestamp,
			SourceID:  lct.Base.SourceID,
		},
		DbID:          0,
		CollectionID:  collID,