milvus/internal/rootcoord/create_partition_task.go

package rootcoord

import (
	"context"
	"fmt"

	pb "github.com/milvus-io/milvus/internal/proto/etcdpb"

	"github.com/milvus-io/milvus-proto/go-api/commonpb"

	"github.com/milvus-io/milvus/internal/log"
	"github.com/milvus-io/milvus/internal/metastore/model"
	"go.uber.org/zap"

	"github.com/milvus-io/milvus-proto/go-api/milvuspb"
)

type createPartitionTask struct {
	baseTask
	Req      *milvuspb.CreatePartitionRequest
	collMeta *model.Collection
}

func (t *createPartitionTask) Prepare(ctx context.Context) error {
	if err := CheckMsgType(t.Req.GetBase().GetMsgType(), commonpb.MsgType_CreatePartition); err != nil {
		return err
	}
	collMeta, err := t.core.meta.GetCollectionByName(ctx, t.Req.GetCollectionName(), t.GetTs())
	if err != nil {
		return err
	}
	t.collMeta = collMeta
	return nil
}

func (t *createPartitionTask) Execute(ctx context.Context) error {
	for _, partition := range t.collMeta.Partitions {
		if partition.PartitionName == t.Req.GetPartitionName() {
			log.Warn("add duplicate partition", zap.String("collection", t.Req.GetCollectionName()), zap.String("partition", t.Req.GetPartitionName()), zap.Uint64("ts", t.GetTs()))
			return nil
		}
	}

	cfgMaxPartitionNum := Params.RootCoordCfg.MaxPartitionNum.GetAsInt()
	if len(t.collMeta.Partitions) >= cfgMaxPartitionNum {
		return fmt.Errorf("partition number (%d) exceeds max configuration (%d), collection: %s",
			len(t.collMeta.Partitions), cfgMaxPartitionNum, t.collMeta.Name)
	}

	partID, err := t.core.idAllocator.AllocOne()
	if err != nil {
		return err
	}
	partition := &model.Partition{
		PartitionID:               partID,
		PartitionName:             t.Req.GetPartitionName(),
		PartitionCreatedTimestamp: t.GetTs(),
		Extra:                     nil,
		CollectionID:              t.collMeta.CollectionID,
		State:                     pb.PartitionState_PartitionCreated,
	}

	undoTask := newBaseUndoTask(t.core.stepExecutor)
	undoTask.AddStep(&expireCacheStep{
		baseStep:        baseStep{core: t.core},
		collectionNames: []string{t.collMeta.Name},
		collectionID:    t.collMeta.CollectionID,
		ts:              t.GetTs(),
	}, &nullStep{})
	undoTask.AddStep(&addPartitionMetaStep{
		baseStep:  baseStep{core: t.core},
		partition: partition,
	}, &nullStep{}) // adding partition is atomic enough.

	return undoTask.Execute(ctx)
}