CreateCollection API执行流程(addCollectionMetaStep)源码解析

milvus版本:v2.3.2

CreateCollection这个API流程较长，也是milvus的核心API之一，涉及的内容比较复杂。这里介绍和channel相关的流程。

整体架构:

CreateCollection(addCollectionMetaStep)的数据流向:

1.客户端sdk发出CreateCollection API请求。

from pymilvus import (
    connections,
    FieldSchema, CollectionSchema, DataType,
    Collection,
)

num_entities, dim = 3000, 1024

print("start connecting to Milvus")
connections.connect("default", host="192.168.230.71", port="19530")

fields = [
    FieldSchema(name="pk", dtype=DataType.VARCHAR, is_primary=True, auto_id=False, max_length=100),
    FieldSchema(name="random", dtype=DataType.DOUBLE),
    FieldSchema(name="embeddings", dtype=DataType.FLOAT_VECTOR, dim=dim)
]

schema = CollectionSchema(fields, "hello_milvus is the simplest demo to introduce the APIs")

print("Create collection `hello_milvus`")
hello_milvus = Collection("hello_milvus", schema, consistency_level="Strong",shards_num=2)

客户端SDK向proxy发送一个CreateCollection API请求，创建一个名为hello_milvus的collection。

2.客户端接受API请求，将request封装为createCollectionTask，并压入ddQueue队列。

代码路径:internal\proxy\impl.go

func (node *Proxy) CreateCollection(ctx context.Context, request *milvuspb.CreateCollectionRequest) (*commonpb.Status, error) {
	......
    // request封装为task
	cct := &createCollectionTask{
		ctx:                     ctx,
		Condition:               NewTaskCondition(ctx),
		CreateCollectionRequest: request,
		rootCoord:               node.rootCoord,
	}

	......
    // 将task压入ddQueue队列
	if err := node.sched.ddQueue.Enqueue(cct); err != nil {
		......
	}

	......
    // 等待cct执行完
	if err := cct.WaitToFinish(); err != nil {
		......
	}

	......
}

3.执行createCollectionTask的3个方法PreExecute、Execute、PostExecute。

PreExecute()一般为参数校验等工作。

Execute()一般为真正执行逻辑。

PostExecute()执行完后的逻辑，什么都不做，返回nil。

代码路径:internal\proxy\task.go

func (t *createCollectionTask) Execute(ctx context.Context) error {
	var err error
	t.result, err = t.rootCoord.CreateCollection(ctx, t.CreateCollectionRequest)
	return err
}

从代码可以看出调用了rootCoord的CreateCollection接口。

4.进入rootCoord的CreateCollection接口。

代码路径:internal\rootcoord\root_coord.go

继续将请求封装为rootcoord里的createCollectionTask

func (c *Core) CreateCollection(ctx context.Context, in *milvuspb.CreateCollectionRequest) (*commonpb.Status, error) {
	......
    // 封装为createCollectionTask
	t := &createCollectionTask{
		baseTask: newBaseTask(ctx, c),
		Req:      in,
	}
    // 加入调度
	if err := c.scheduler.AddTask(t); err != nil {
		......
	}
    // 等待task完成
	if err := t.WaitToFinish(); err != nil {
		......
	}

	......
}

5.执行createCollectionTask的Prepare、Execute、NotifyDone方法。

Execute()为核心方法。

代码路径:internal\rootcoord\create_collection_task.go

func (t *createCollectionTask) Execute(ctx context.Context) error {
	// collID为collectionID,在Prepare()里分配
	// partIDs为partitionID,在Prepare()里分配
	collID := t.collID
	partIDs := t.partIDs
	// 产生时间戳
	ts, err := t.getCreateTs()
	if err != nil {
		return err
	}
	// vchanNames为虚拟channel,在Prepare()里分配
	// chanNames为物理channel,在Prepare()里分配
	vchanNames := t.channels.virtualChannels
	chanNames := t.channels.physicalChannels

	startPositions, err := t.addChannelsAndGetStartPositions(ctx, ts)
	if err != nil {
		t.core.chanTimeTick.removeDmlChannels(t.channels.physicalChannels...)
		return err
	}
	// 填充partition,创建collection的时候,默认只有一个名为"Default partition"的partition。
	partitions := make([]*model.Partition, len(partIDs))
	for i, partID := range partIDs {
		partitions[i] = &model.Partition{
			PartitionID:               partID,
			PartitionName:             t.partitionNames[i],
			PartitionCreatedTimestamp: ts,
			CollectionID:              collID,
			State:                     pb.PartitionState_PartitionCreated,
		}
	}
	// 填充collection
	// 可以看出collection由collID、dbid、schemaName、fields、vchanName、chanName、partition、shardNum等组成
	collInfo := model.Collection{
		CollectionID:         collID,
		DBID:                 t.dbID,
		Name:                 t.schema.Name,
		Description:          t.schema.Description,
		AutoID:               t.schema.AutoID,
		Fields:               model.UnmarshalFieldModels(t.schema.Fields),
		VirtualChannelNames:  vchanNames,
		PhysicalChannelNames: chanNames,
		ShardsNum:            t.Req.ShardsNum,
		ConsistencyLevel:     t.Req.ConsistencyLevel,
		StartPositions:       toKeyDataPairs(startPositions),
		CreateTime:           ts,
		State:                pb.CollectionState_CollectionCreating,
		Partitions:           partitions,
		Properties:           t.Req.Properties,
		EnableDynamicField:   t.schema.EnableDynamicField,
	}

	clone := collInfo.Clone()

	existedCollInfo, err := t.core.meta.GetCollectionByName(ctx, t.Req.GetDbName(), t.Req.GetCollectionName(), typeutil.MaxTimestamp)
	if err == nil {
		equal := existedCollInfo.Equal(*clone)
		if !equal {
			return fmt.Errorf("create duplicate collection with different parameters, collection: %s", t.Req.GetCollectionName())
		}

		log.Warn("add duplicate collection", zap.String("collection", t.Req.GetCollectionName()), zap.Uint64("ts", ts))
		return nil
	}
	// 分为多个step执行，每一个undoTask由todoStep和undoStep构成
	// 执行todoStep,报错则执行undoStep
	undoTask := newBaseUndoTask(t.core.stepExecutor)
	undoTask.AddStep(&expireCacheStep{
		baseStep:        baseStep{core: t.core},
		dbName:          t.Req.GetDbName(),
		collectionNames: []string{t.Req.GetCollectionName()},
		collectionID:    InvalidCollectionID,
		ts:              ts,
	}, &nullStep{})
	undoTask.AddStep(&nullStep{}, &removeDmlChannelsStep{
		baseStep:  baseStep{core: t.core},
		pChannels: chanNames,
	}) 
	undoTask.AddStep(&addCollectionMetaStep{
		baseStep: baseStep{core: t.core},
		coll:     &collInfo,
	}, &deleteCollectionMetaStep{
		baseStep:     baseStep{core: t.core},
		collectionID: collID,
		ts: ts,
	})

	undoTask.AddStep(&nullStep{}, &unwatchChannelsStep{
		baseStep:     baseStep{core: t.core},
		collectionID: collID,
		channels:     t.channels,
		isSkip:       !Params.CommonCfg.TTMsgEnabled.GetAsBool(),
	})
	undoTask.AddStep(&watchChannelsStep{
		baseStep: baseStep{core: t.core},
		info: &watchInfo{
			ts:             ts,
			collectionID:   collID,
			vChannels:      t.channels.virtualChannels,
			startPositions: toKeyDataPairs(startPositions),
			schema: &schemapb.CollectionSchema{
				Name:        collInfo.Name,
				Description: collInfo.Description,
				AutoID:      collInfo.AutoID,
				Fields:      model.MarshalFieldModels(collInfo.Fields),
			},
		},
	}, &nullStep{})
	undoTask.AddStep(&changeCollectionStateStep{
		baseStep:     baseStep{core: t.core},
		collectionID: collID,
		state:        pb.CollectionState_CollectionCreated,
		ts:           ts,
	}, &nullStep{})

	return undoTask.Execute(ctx)
}

创建collection涉及多个步骤，可以看出这里依次分为expireCacheStep、addCollectionMetaStep、watchChannelsStep、changeCollectionStateStep这几个步骤，addCollectionMetaStep是关于etcd元数据的step，已在另一篇文章对其进行详细解析。本篇幅对watchChannelsStep进行解析。

6.进入watchChannelsStep，执行其Execute()方法。

代码路径:internal\rootcoord\step.go

func (s *watchChannelsStep) Execute(ctx context.Context) ([]nestedStep, error) {
	err := s.core.broker.WatchChannels(ctx, s.info)
	return nil, err
}

在这里重点研究s.core.broker.WatchChannels()这个方法做了什么事情。

调用栈如下:

s.core.broker.WatchChannels()
  |--WatchChannels()(internal\rootcoord\broker.go)
    |--b.s.dataCoord.WatchChannels()
      |--WatchChannels()(internal\datacoord\services.go)
        |--s.channelManager.Watch()
          |--c.updateWithTimer()(internal\datacoord\channel_manager.go)
            |--c.store.Update()
              |--c.update()(internal\datacoord\channel_store.go)
                |--c.txn()(同上)
                  |--c.store.MultiSaveAndRemove()(同上)
                    |--MultiSaveAndRemove()(internal\kv\etcd\etcd_kv.go)
        |--s.meta.catalog.MarkChannelAdded()

WatchChannels这个操作最终是在etcd写入kv。那么我们研究写入的kv是什么。

根据堆栈顺序来进行分析。

1.WatchChannels()方法

代码路径:internal\datacoord\services.go

// WatchChannels notifies DataCoord to watch vchannels of a collection.
func (s *Server) WatchChannels(ctx context.Context, req *datapb.WatchChannelsRequest) (*datapb.WatchChannelsResponse, error) {
	log := log.Ctx(ctx).With(
		zap.Int64("collectionID", req.GetCollectionID()),
		zap.Strings("channels", req.GetChannelNames()),
	)
	log.Info("receive watch channels request")
	resp := &datapb.WatchChannelsResponse{
		Status: merr.Success(),
	}

	if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
		return &datapb.WatchChannelsResponse{
			Status: merr.Status(err),
		}, nil
	}
    // req.GetChannelNames()得到的值为:
    // by-dev-rootcoord-dml_2_445674962009727985v0
    // by-dev-rootcoord-dml_3_445674962009727985v1
	for _, channelName := range req.GetChannelNames() {
		ch := &channel{
			Name:            channelName,
			CollectionID:    req.GetCollectionID(),
			StartPositions:  req.GetStartPositions(),
			Schema:          req.GetSchema(),
			CreateTimestamp: req.GetCreateTimestamp(),
		}
        // 循环执行watch()
		err := s.channelManager.Watch(ctx, ch)
		if err != nil {
			log.Warn("fail to watch channelName", zap.Error(err))
			resp.Status = merr.Status(err)
			return resp, nil
		}
        // 向etcd写入另外一个kv
		if err := s.meta.catalog.MarkChannelAdded(ctx, ch.Name); err != nil {
			// TODO: add background task to periodically cleanup the orphaned channel add marks.
			log.Error("failed to mark channel added", zap.Error(err))
			resp.Status = merr.Status(err)
			return resp, nil
		}
	}

	return resp, nil
}

函数入参req的值如下:

在这里有2个channelName，是虚拟channel，为什么是2个channel?因为客户端SDK创建collection传入了shards_num=2。一个shard对应一个虚拟channel。

channel名称by-dev-rootcoord-dml_2_445674962009727985v0中的445674962009727985是collectionID。

2.进入到s.channelManager.Watch()

代码路径:internal\datacoord\channel_manager.go

// Watch tries to add the channel to cluster. Watch is a no op if the channel already exists.
func (c *ChannelManager) Watch(ctx context.Context, ch *channel) error {
	log := log.Ctx(ctx)
	c.mu.Lock()
	defer c.mu.Unlock()
    // 使用分配策略:datacoord.AverageAssignPolicy
	updates := c.assignPolicy(c.store, []*channel{ch})
	if len(updates) == 0 {
		return nil
	}
	log.Info("try to update channel watch info with ToWatch state",
		zap.String("channel", ch.String()),
		zap.Array("updates", updates))
    // 操作etcd
	err := c.updateWithTimer(updates, datapb.ChannelWatchState_ToWatch)
	if err != nil {
		log.Warn("fail to update channel watch info with ToWatch state",
			zap.String("channel", ch.String()), zap.Array("updates", updates), zap.Error(err))
	}
	return err
}

updates的值为:

updates变量是一个ChannelOpSet类型。这时候ChannelWatchInfos为空。

type ChannelOpSet []*ChannelOp

type ChannelOp struct {
	Type              ChannelOpType
	NodeID            int64
	Channels          []*channel
	ChannelWatchInfos []*datapb.ChannelWatchInfo
}

3.进入c.updateWithTimer()

代码路径:internal\datacoord\channel_manager.go

func (c *ChannelManager) updateWithTimer(updates ChannelOpSet, state datapb.ChannelWatchState) error {
	channelsWithTimer := []string{}
	// updates此时数组长度为1
	for _, op := range updates {
		if op.Type == Add {
			// 填充ChannelWatchInfos
			channelsWithTimer = append(channelsWithTimer, c.fillChannelWatchInfoWithState(op, state)...)
		}
	}
    // 操作etcd
	err := c.store.Update(updates)
	if err != nil {
		log.Warn("fail to update", zap.Array("updates", updates), zap.Error(err))
		c.stateTimer.removeTimers(channelsWithTimer)
	}
	c.lastActiveTimestamp = time.Now()
	return err
}

4.进入c.store.Update()

代码路径:internal\datacoord\channel_store.go

// Update applies the channel operations in opSet.
func (c *ChannelStore) Update(opSet ChannelOpSet) error {
	totalChannelNum := 0
	for _, op := range opSet {
		totalChannelNum += len(op.Channels)
	}
	// totalChannelNum = 1
	// maxOperationsPerTxn = 64
	if totalChannelNum <= maxOperationsPerTxn {
		// 走这条路径
		return c.update(opSet)
	}
	// 如果超过则分批执行
	......
}

5.进入c.update(opSet)

代码路径:internal\datacoord\channel_store.go

// update applies the ADD/DELETE operations to the current channel store.
func (c *ChannelStore) update(opSet ChannelOpSet) error {
	// Update ChannelStore's kv store.
    // 操作etcd
	if err := c.txn(opSet); err != nil {
		return err
	}

	// Update node id -> channel mapping.
	for _, op := range opSet {
		switch op.Type {
		case Add:
			for _, ch := range op.Channels {
				if c.checkIfExist(op.NodeID, ch) {
					continue // prevent adding duplicated channel info
				}
				// Append target channels to channel store.
				c.channelsInfo[op.NodeID].Channels = append(c.channelsInfo[op.NodeID].Channels, ch)
			}
		case Delete:
			// Remove target channels from channel store.
			del := make(map[string]struct{})
			for _, ch := range op.Channels {
				del[ch.Name] = struct{}{}
			}
			prev := c.channelsInfo[op.NodeID].Channels
			curr := make([]*channel, 0, len(prev))
			for _, ch := range prev {
				if _, ok := del[ch.Name]; !ok {
					curr = append(curr, ch)
				}
			}
			c.channelsInfo[op.NodeID].Channels = curr
		default:
			return errUnknownOpType
		}
		metrics.DataCoordDmlChannelNum.WithLabelValues(strconv.FormatInt(op.NodeID, 10)).Set(float64(len(c.channelsInfo[op.NodeID].Channels)))
	}
	return nil
}

6.进入c.txn(opSet)

代码路径:internal\datacoord\channel_store.go

// txn updates the channelStore's kv store with the given channel ops.
func (c *ChannelStore) txn(opSet ChannelOpSet) error {
	saves := make(map[string]string)
	var removals []string
	for _, op := range opSet {
		for i, ch := range op.Channels {
			// 构建key的规则
			k := buildNodeChannelKey(op.NodeID, ch.Name)
			switch op.Type {
			case Add:
				// 构建value,ChannelWatchInfo
				info, err := proto.Marshal(op.ChannelWatchInfos[i])
				if err != nil {
					return err
				}
				saves[k] = string(info)
			case Delete:
				removals = append(removals, k)
			default:
				return errUnknownOpType
			}
		}
	}
	return c.store.MultiSaveAndRemove(saves, removals)
}

因为op.Type是Add，所以removals是nil。

key的值:

channelwatch/1/by-dev-rootcoord-dml_2_445674962009727985v0

规则为:channelwatch/{nodeID}/{chName}

saves变量的值:

后面已经不用再跟踪下去。

使用etcd-manager查看etcd。

7.进入s.meta.catalog.MarkChannelAdded()

代码路径:internal\metastore\kv\datacoord\kv_catalog.go

func (kc *Catalog) MarkChannelAdded(ctx context.Context, channel string) error {
	// 构建key的规则:datacoord-meta/channel-removal/{channelName}
	key := buildChannelRemovePath(channel)
	// 构建value:NonRemoveFlagTomestone = "non-removed"
	err := kc.MetaKv.Save(key, NonRemoveFlagTomestone)
	if err != nil {
		log.Error("failed to mark channel added", zap.String("channel", channel), zap.Error(err))
		return err
	}
	log.Info("NON remove flag tombstone added", zap.String("channel", channel))
	return nil
}

构建key的规则:

datacoord-meta/channel-removal/{channelName}