// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package task

import (
	"context"
	"fmt"
	"math"
	"sync"
	"time"

	"github.com/blang/semver/v4"
	"github.com/cockroachdb/errors"
	"github.com/samber/lo"
	"go.uber.org/atomic"

	"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
	"github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
	"github.com/milvus-io/milvus/internal/querycoordv2/meta"
	. "github.com/milvus-io/milvus/internal/querycoordv2/params"
	"github.com/milvus-io/milvus/internal/querycoordv2/session"
	"github.com/milvus-io/milvus/internal/querycoordv2/utils"
	"github.com/milvus-io/milvus/pkg/v3/common"
	"github.com/milvus-io/milvus/pkg/v3/mlog"
	"github.com/milvus-io/milvus/pkg/v3/proto/datapb"
	"github.com/milvus-io/milvus/pkg/v3/proto/indexpb"
	"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
	"github.com/milvus-io/milvus/pkg/v3/util/commonpbutil"
	"github.com/milvus-io/milvus/pkg/v3/util/conc"
	"github.com/milvus-io/milvus/pkg/v3/util/funcutil"
	"github.com/milvus-io/milvus/pkg/v3/util/indexparams"
	"github.com/milvus-io/milvus/pkg/v3/util/merr"
	"github.com/milvus-io/milvus/pkg/v3/util/tsoutil"
	"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
)

// segmentsVersion is used for the flushed segments should not be included in the watch dm channel request
var segmentsVersion = semver.Version{
	Major: 2,
	Minor: 3,
	Patch: 4,
}

type Executor struct {
	nodeID    int64
	doneCh    chan struct{}
	wg        sync.WaitGroup
	meta      *meta.Meta
	dist      *meta.DistributionManager
	broker    meta.Broker
	targetMgr meta.TargetManagerInterface
	cluster   session.Cluster
	nodeMgr   *session.NodeManager

	executingTasks    *typeutil.ConcurrentSet[string] // task index
	channelTaskNum    atomic.Int32                    // channel task pool counter
	nonChannelTaskNum atomic.Int32                    // non-channel task pool counter

	collectionInfoSingleflight conc.Singleflight[*milvuspb.DescribeCollectionResponse]
}

func NewExecutor(nodeID int64,
	meta *meta.Meta,
	dist *meta.DistributionManager,
	broker meta.Broker,
	targetMgr meta.TargetManagerInterface,
	cluster session.Cluster,
	nodeMgr *session.NodeManager,
) *Executor {
	return &Executor{
		nodeID:    nodeID,
		doneCh:    make(chan struct{}),
		meta:      meta,
		dist:      dist,
		broker:    broker,
		targetMgr: targetMgr,
		cluster:   cluster,
		nodeMgr:   nodeMgr,

		executingTasks: typeutil.NewConcurrentSet[string](),
	}
}

func (ex *Executor) Start(ctx context.Context) {
}

func (ex *Executor) Stop() {
	ex.wg.Wait()
}

func (ex *Executor) GetTotalTaskExecutionCap() int32 {
	nodeInfo := ex.nodeMgr.Get(ex.nodeID)
	if nodeInfo == nil || nodeInfo.CPUNum() == 0 {
		return Params.QueryCoordCfg.TaskExecutionCap.GetAsInt32()
	}

	ret := int32(math.Ceil(float64(nodeInfo.CPUNum()) * Params.QueryCoordCfg.QueryNodeTaskParallelismFactor.GetAsFloat()))

	return ret
}

// GetChannelTaskCap returns the capacity reserved for channel tasks.
func (ex *Executor) GetChannelTaskCap() int32 {
	total := ex.GetTotalTaskExecutionCap()
	fraction := Params.QueryCoordCfg.ChannelTaskCapFraction.GetAsFloat()
	if fraction < 0 {
		fraction = 0
	}
	if fraction > 1 {
		fraction = 1
	}
	cap := int32(math.Ceil(float64(total) * fraction))
	if cap < 1 {
		cap = 1
	}
	return cap
}

// GetNonChannelTaskCap returns the capacity for segment/leader/other tasks.
// NOTE: when channelTaskCapFraction is 1.0, both pools get min-cap=1,
// so the sum of channel + non-channel caps may exceed total. This is
// intentional to guarantee liveness for both task types.
func (ex *Executor) GetNonChannelTaskCap() int32 {
	total := ex.GetTotalTaskExecutionCap()
	nonChannelCap := total - ex.GetChannelTaskCap()
	if nonChannelCap < 1 {
		nonChannelCap = 1
	}
	return nonChannelCap
}

// Execute executes the given action,
// does nothing and returns false if the action is already committed,
// returns true otherwise.
func (ex *Executor) Execute(task Task, step int) bool {
	exist := !ex.executingTasks.Insert(task.Index())
	if exist {
		return false
	}

	_, isChannel := task.Actions()[step].(*ChannelAction)
	if isChannel {
		cur := ex.channelTaskNum.Inc()
		if cur > ex.GetChannelTaskCap() {
			ex.channelTaskNum.Dec()
			ex.executingTasks.Remove(task.Index())
			mlog.Debug(context.TODO(), "channel task rejected: pool full",
				mlog.Int32("current", cur),
				mlog.Int32("cap", ex.GetChannelTaskCap()))
			return false
		}
	} else {
		cur := ex.nonChannelTaskNum.Inc()
		if cur > ex.GetNonChannelTaskCap() {
			ex.nonChannelTaskNum.Dec()
			ex.executingTasks.Remove(task.Index())
			mlog.Debug(context.TODO(), "non-channel task rejected: pool full",
				mlog.Int32("current", cur),
				mlog.Int32("cap", ex.GetNonChannelTaskCap()))
			return false
		}
	}

	go func() {
		mlog.Info(context.TODO(), "execute the action of task")
		switch task.Actions()[step].(type) {
		case *SegmentAction:
			ex.executeSegmentAction(task.(*SegmentTask), step)

		case *ChannelAction:
			ex.executeDmChannelAction(task.(*ChannelTask), step)

		case *LeaderAction:
			ex.executeLeaderAction(task.(*LeaderTask), step)

		case *DropIndexAction:
			ex.executeDropIndexAction(task.(*DropIndexTask), step)
		}
	}()

	return true
}

func (ex *Executor) removeTask(task Task, step int) {
	if task.Err() != nil {
		mlog.Info(context.TODO(), "execute action done, remove it",
			mlog.Int64("taskID", task.ID()),
			mlog.Int("step", step),
			mlog.Err(task.Err()))
	}

	ex.executingTasks.Remove(task.Index())
	if _, isChannel := task.Actions()[step].(*ChannelAction); isChannel {
		ex.channelTaskNum.Dec()
	} else {
		ex.nonChannelTaskNum.Dec()
	}
}

func (ex *Executor) executeSegmentAction(task *SegmentTask, step int) {
	switch task.Actions()[step].Type() {
	case ActionTypeGrow, ActionTypeUpdate, ActionTypeStatsUpdate, ActionTypeReopen:
		ex.loadSegment(task, step)

	case ActionTypeReduce:
		ex.releaseSegment(task, step)
	}
}

// loadSegment commits the request to merger,
// not really executes the request
func (ex *Executor) loadSegment(task *SegmentTask, step int) error {
	action := task.Actions()[step].(*SegmentAction)
	defer action.rpcReturned.Store(true)
	ctx := task.Context()

	var err error
	defer func() {
		if err != nil {
			task.Fail(err)
		}
		ex.removeTask(task, step)
	}()

	collectionInfo, loadMeta, channel, err := ex.getMetaInfo(ctx, task)
	if err != nil {
		return err
	}

	loadInfo, indexInfos, err := ex.getLoadInfo(ctx, task.CollectionID(), action.SegmentID, channel, task.LoadPriority())
	if err != nil {
		return err
	}
	req := packLoadSegmentRequest(
		task,
		action,
		collectionInfo.GetSchema(),
		collectionInfo.GetProperties(),
		loadMeta,
		loadInfo,
		indexInfos,
	)

	// get segment's replica first, then get shard leader by replica
	replica := ex.meta.Get(ctx, task.ReplicaID())
	if replica == nil {
		msg := "node doesn't belong to any replica"
		err := merr.WrapErrNodeNotAvailable(action.Node())
		mlog.Warn(ctx, msg, mlog.Err(err))
		return err
	}
	view := ex.dist.ChannelDistManager.GetShardLeader(task.Shard(), replica)
	if view == nil {
		msg := "no shard leader for the segment to execute loading"
		err = merr.WrapErrChannelNotFound(task.Shard(), "shard delegator not found")
		mlog.Warn(ctx, msg, mlog.Err(err))
		return err
	}

	// NOTE: for balance segment task, expected load and release execution on the same shard leader
	if GetTaskType(task) == TaskTypeMove {
		task.SetShardLeaderID(view.Node)
	}

	startTs := time.Now()
	mlog.Info(context.TODO(), "load segments...")
	status, err := ex.cluster.LoadSegments(task.Context(), view.Node, req)
	err = merr.CheckRPCCall(status, err)
	if err != nil {
		mlog.Warn(context.TODO(), "failed to load segment", mlog.Err(err))
		return err
	}

	elapsed := time.Since(startTs)
	mlog.Info(context.TODO(), "load segments done", mlog.Duration("elapsed", elapsed))

	return nil
}

// If we enable following checking when loading segments,
// 1. all segment should always be loaded by streamingnode but not 2.5 querynode, make some search and query failure when upgrading.
// Otherwise, some search and query result will be wrong when upgrading.
// We choose to disable this checking for now to promise available search and query when upgrading.
//
// Because the L0 management at 2.6 and 2.5 is different, so when upgrading mixcoord,
// the new mixcoord will make a wrong plan when balancing a segment from one query node to another by 2.5 delegator.
// We need to balance the 2.5 delegator to 2.6 delegator before balancing any segment by 2.6 mixcoord.
// func (ex *Executor) checkIfShardLeaderIsStreamingNode(view *meta.DmChannel) error {
// 	if !streamingutil.IsStreamingServiceEnabled() {
// 		return nil
// 	}
//
// 	node := ex.nodeMgr.Get(view.Node)
// 	if node == nil {
// 		return merr.WrapErrServiceInternalMsg("node %d is not found", view.Node)
// 	}
// 	nodes := snmanager.StaticStreamingNodeManager.GetStreamingQueryNodeIDs()
// 	if !nodes.Contain(view.Node) {
// 		return merr.WrapErrServiceInternalMsg("channel %s at node %d is not working at streamingnode, skip load segment", view.GetChannelName(), view.Node)
// 	}
// 	return nil
// }

func (ex *Executor) releaseSegment(task *SegmentTask, step int) {
	defer ex.removeTask(task, step)
	startTs := time.Now()
	action := task.Actions()[step].(*SegmentAction)
	defer action.rpcReturned.Store(true)

	ctx := task.Context()
	var err error
	defer func() {
		if err != nil {
			task.Fail(err)
		}
	}()

	dstNode := action.Node()

	req := packReleaseSegmentRequest(task, action)
	channel := ex.targetMgr.GetDmChannel(ctx, task.CollectionID(), task.Shard(), meta.CurrentTarget)
	if channel != nil {
		// if channel exists in current target, set cp to ReleaseSegmentRequest, need to use it as growing segment's exclude ts
		req.Checkpoint = channel.GetSeekPosition()
	}

	if action.Scope == querypb.DataScope_Streaming {
		// Any modification to the segment distribution have to set NeedTransfer true,
		// to protect the version, which serves search/query
		req.NeedTransfer = true
	} else {
		req.Shard = task.shard

		if ex.meta.Exist(ctx, task.CollectionID()) {
			// get segment's replica first, then get shard leader by replica
			replica := ex.meta.Get(ctx, task.ReplicaID())
			if replica == nil {
				msg := "node doesn't belong to any replica, try to send release to worker"
				err := merr.WrapErrNodeNotAvailable(action.Node())
				mlog.Warn(ctx, msg, mlog.Err(err))
				dstNode = action.Node()
				req.NeedTransfer = false
			} else {
				view := ex.dist.ChannelDistManager.GetShardLeader(task.Shard(), replica)
				if view == nil {
					msg := "no shard leader for the segment to execute releasing"
					err = merr.WrapErrChannelNotFound(task.Shard(), "shard delegator not found")
					mlog.Warn(ctx, msg, mlog.Err(err))
					return
				}
				// NOTE: for balance segment task, expected load and release execution on the same shard leader
				if GetTaskType(task) == TaskTypeMove && task.ShardLeaderID() != view.Node {
					msg := "shard leader changed, skip release"
					err = merr.WrapErrServiceInternalMsg("shard leader changed from %d to %d", task.ShardLeaderID(), view.Node)
					mlog.Warn(ctx, msg, mlog.Err(err))
					return
				}
				dstNode = view.Node
				req.NeedTransfer = true
			}
		}
	}

	mlog.Info(context.TODO(), "release segment...")
	status, err := ex.cluster.ReleaseSegments(ctx, dstNode, req)
	err = merr.CheckRPCCall(status, err)
	if err != nil {
		mlog.Warn(context.TODO(), "failed to release segment", mlog.Err(err))
		return
	}
	elapsed := time.Since(startTs)
	mlog.Info(context.TODO(), "release segment done", mlog.Int64("taskID", task.ID()), mlog.Duration("time taken", elapsed))
}

func (ex *Executor) executeDmChannelAction(task *ChannelTask, step int) {
	switch task.Actions()[step].Type() {
	case ActionTypeGrow:
		ex.subscribeChannel(task, step)

	case ActionTypeReduce:
		ex.unsubscribeChannel(task, step)
	}
}

func (ex *Executor) subscribeChannel(task *ChannelTask, step int) error {
	defer ex.removeTask(task, step)
	startTs := time.Now()
	action := task.Actions()[step].(*ChannelAction)

	var err error
	defer func() {
		if err != nil {
			task.Fail(err)
		}
	}()

	ctx := task.Context()

	collectionInfo, err := ex.broker.DescribeCollection(ctx, task.CollectionID())
	if err != nil {
		mlog.Warn(context.TODO(), "failed to get collection info", mlog.Err(err))
		return err
	}
	loadFields := ex.meta.GetLoadFields(ctx, task.CollectionID())
	partitions, err := utils.GetPartitions(ctx, ex.targetMgr, task.CollectionID())
	if err != nil {
		mlog.Warn(context.TODO(), "failed to get partitions of collection", mlog.Err(err))
		return err
	}
	indexInfo, err := ex.broker.ListIndexes(ctx, task.CollectionID())
	if err != nil {
		mlog.Warn(context.TODO(), "fail to get index meta of collection", mlog.Err(err))
		return err
	}
	dbResp, err := ex.broker.DescribeDatabase(ctx, collectionInfo.GetDbName())
	if err != nil {
		mlog.Warn(context.TODO(), "failed to get database info", mlog.Err(err))
		return err
	}
	loadMeta := packLoadMeta(
		ex.meta.GetLoadType(ctx, task.CollectionID()),
		collectionInfo,
		task.ResourceGroup(),
		loadFields,
		partitions...,
	)
	loadMeta.DbProperties = dbResp.GetProperties()

	dmChannel := ex.targetMgr.GetDmChannel(ctx, task.CollectionID(), action.ChannelName(), meta.NextTarget)
	if dmChannel == nil {
		msg := "channel does not exist in next target, skip it"
		mlog.Warn(ctx, msg, mlog.String("channelName", action.ChannelName()))
		return merr.WrapErrChannelReduplicate(action.ChannelName())
	}

	partitions, err = utils.GetPartitions(ctx, ex.targetMgr, task.collectionID)
	if err != nil {
		mlog.Warn(context.TODO(), "failed to get partitions", mlog.Err(err))
		return merr.Wrapf(err, "failed to get partitions for collection=%d", task.CollectionID())
	}

	version := ex.targetMgr.GetCollectionTargetVersion(ctx, task.CollectionID(), meta.NextTargetFirst)

	req := packSubChannelRequest(
		task,
		action,
		collectionInfo.GetSchema(),
		collectionInfo.GetProperties(),
		loadMeta,
		dmChannel,
		indexInfo,
		partitions,
		version,
	)
	err = fillSubChannelRequest(ctx, req, ex.broker, ex.shouldIncludeFlushedSegmentInfo(action.Node()))
	if err != nil {
		mlog.Warn(context.TODO(), "failed to subscribe channel, failed to fill the request with segments",
			mlog.Err(err))
		return err
	}

	sealedSegments := ex.targetMgr.GetSealedSegmentsByChannel(ctx, dmChannel.CollectionID, dmChannel.ChannelName, meta.NextTarget)
	sealedSegmentRowCount := lo.MapValues(sealedSegments, func(segment *datapb.SegmentInfo, _ int64) int64 {
		return segment.GetNumOfRows()
	})
	req.SealedSegmentRowCount = sealedSegmentRowCount

	ts := dmChannel.GetSeekPosition().GetTimestamp()
	mlog.Info(context.TODO(), "subscribe channel...",
		mlog.Uint64("checkpoint", ts),
		mlog.Duration("sinceCheckpoint", time.Since(tsoutil.PhysicalTime(ts))),
	)
	status, err := ex.cluster.WatchDmChannels(ctx, action.Node(), req)
	if err != nil {
		mlog.Warn(context.TODO(), "failed to subscribe channel, it may be a false failure", mlog.Err(err))
		return err
	}
	if !merr.Ok(status) {
		err = merr.Error(status)
		mlog.Warn(context.TODO(), "failed to subscribe channel", mlog.Err(err))
		return err
	}
	elapsed := time.Since(startTs)
	mlog.Info(context.TODO(), "subscribe channel done", mlog.Int64("taskID", task.ID()), mlog.Duration("time taken", elapsed))
	return nil
}

func (ex *Executor) shouldIncludeFlushedSegmentInfo(nodeID int64) bool {
	node := ex.nodeMgr.Get(nodeID)
	if node == nil {
		return false
	}
	return node.Version().LT(segmentsVersion)
}

func (ex *Executor) unsubscribeChannel(task *ChannelTask, step int) error {
	defer ex.removeTask(task, step)
	startTs := time.Now()
	action := task.Actions()[step].(*ChannelAction)

	var err error
	defer func() {
		if err != nil {
			task.Fail(err)
		}
	}()

	ctx := task.Context()

	req := packUnsubDmChannelRequest(task, action)
	mlog.Info(context.TODO(), "unsubscribe channel...")
	status, err := ex.cluster.UnsubDmChannel(ctx, action.Node(), req)
	if err != nil {
		mlog.Warn(context.TODO(), "failed to unsubscribe channel, it may be a false failure", mlog.Err(err))
		return err
	}
	if !merr.Ok(status) {
		err = merr.Error(status)
		mlog.Warn(context.TODO(), "failed to unsubscribe channel", mlog.Err(err))
		return err
	}

	elapsed := time.Since(startTs)
	mlog.Info(context.TODO(), "unsubscribe channel done", mlog.Int64("taskID", task.ID()), mlog.Duration("time taken", elapsed))
	return nil
}

func (ex *Executor) executeLeaderAction(task *LeaderTask, step int) {
	switch task.Actions()[step].Type() {
	case ActionTypeGrow:
		ex.setDistribution(task, step)

	case ActionTypeReduce:
		ex.removeDistribution(task, step)

	case ActionTypeUpdate:
		ex.updatePartStatsVersions(task, step)

	case ActionTypeStatsUpdate:
		ex.updatePartStatsVersions(task, step)
	}
}

func (ex *Executor) executeDropIndexAction(task *DropIndexTask, step int) {
	action := task.Actions()[step].(*DropIndexAction)
	defer action.rpcReturned.Store(true)
	ctx := task.Context()

	var err error
	defer func() {
		if err != nil {
			task.Fail(err)
		}
		ex.removeTask(task, step)
	}()

	replica := ex.meta.Get(ctx, task.ReplicaID())
	if replica == nil {
		err = merr.WrapErrNodeNotAvailable(action.Node())
		mlog.Warn(context.TODO(), "node doesn't belong to any replica", mlog.Err(err))
		return
	}
	view := ex.dist.ChannelDistManager.GetShardLeader(task.Shard(), replica)
	if view == nil {
		err = merr.WrapErrChannelNotFound(task.Shard(), "shard delegator not found")
		mlog.Warn(context.TODO(), "failed to get shard leader", mlog.Err(err))
		return
	}

	req := &querypb.DropIndexRequest{
		Base: commonpbutil.NewMsgBase(
			commonpbutil.WithMsgType(commonpb.MsgType_DropIndex),
			commonpbutil.WithMsgID(task.ID()),
		),
		SegmentID:    task.SegmentID(),
		IndexIDs:     action.indexIDs,
		Channel:      task.Shard(),
		NeedTransfer: true,
	}

	startTs := time.Now()
	mlog.Info(context.TODO(), "drop index...")
	status, err := ex.cluster.DropIndex(task.Context(), view.Node, req)
	if err != nil {
		mlog.Warn(context.TODO(), "failed to drop index", mlog.Err(err))
		return
	}
	if !merr.Ok(status) {
		err = merr.Error(status)
		mlog.Warn(context.TODO(), "failed to drop index", mlog.Err(err))
		return
	}

	elapsed := time.Since(startTs)
	mlog.Info(context.TODO(), "drop index done", mlog.Duration("elapsed", elapsed))
}

func (ex *Executor) updatePartStatsVersions(task *LeaderTask, step int) error {
	action := task.Actions()[step].(*LeaderAction)
	defer action.rpcReturned.Store(true)
	var err error
	defer func() {
		if err != nil {
			task.Fail(err)
		}
		ex.removeTask(task, step)
	}()

	req := &querypb.SyncDistributionRequest{
		Base: commonpbutil.NewMsgBase(
			commonpbutil.WithMsgType(commonpb.MsgType_SyncDistribution),
			commonpbutil.WithMsgID(task.ID()),
		),
		CollectionID: task.collectionID,
		Channel:      task.Shard(),
		ReplicaID:    task.ReplicaID(),
		Actions: []*querypb.SyncAction{
			{
				Type:                   querypb.SyncType_UpdatePartitionStats,
				SegmentID:              action.SegmentID(),
				NodeID:                 action.Node(),
				Version:                action.Version(),
				PartitionStatsVersions: action.partStatsVersions,
			},
		},
	}
	startTs := time.Now()
	mlog.Debug(context.TODO(), "Update partition stats versions...")
	status, err := ex.cluster.SyncDistribution(task.Context(), task.leaderID, req)
	err = merr.CheckRPCCall(status, err)
	if err != nil {
		mlog.Warn(context.TODO(), "failed to update partition stats versions", mlog.Err(err))
		return err
	}

	elapsed := time.Since(startTs)
	mlog.Debug(context.TODO(), "update partition stats done", mlog.Duration("elapsed", elapsed))

	return nil
}

func (ex *Executor) setDistribution(task *LeaderTask, step int) error {
	action := task.Actions()[step].(*LeaderAction)
	defer action.rpcReturned.Store(true)
	ctx := task.Context()

	var err error
	defer func() {
		if err != nil {
			task.Fail(err)
		}
		ex.removeTask(task, step)
	}()

	collectionInfo, loadMeta, channel, err := ex.getMetaInfo(ctx, task)
	if err != nil {
		return err
	}

	loadInfo, indexInfo, err := ex.getLoadInfo(ctx, task.CollectionID(), action.SegmentID(), channel, commonpb.LoadPriority_LOW)
	if err != nil {
		return err
	}

	req := &querypb.SyncDistributionRequest{
		Base: commonpbutil.NewMsgBase(
			commonpbutil.WithMsgType(commonpb.MsgType_LoadSegments),
			commonpbutil.WithMsgID(task.ID()),
		),
		CollectionID: task.collectionID,
		Channel:      task.Shard(),
		Schema:       collectionInfo.GetSchema(),
		LoadMeta:     loadMeta,
		ReplicaID:    task.ReplicaID(),
		Actions: []*querypb.SyncAction{
			{
				Type:        querypb.SyncType_Set,
				PartitionID: loadInfo.GetPartitionID(),
				SegmentID:   action.SegmentID(),
				NodeID:      action.Node(),
				Info:        loadInfo,
				Version:     action.Version(),
			},
		},
		IndexInfoList: indexInfo,
	}

	startTs := time.Now()
	mlog.Info(context.TODO(), "Sync Distribution...")
	status, err := ex.cluster.SyncDistribution(task.Context(), task.leaderID, req)
	err = merr.CheckRPCCall(status, err)
	if err != nil {
		mlog.Warn(context.TODO(), "failed to sync distribution", mlog.Err(err))
		return err
	}

	elapsed := time.Since(startTs)
	mlog.Info(context.TODO(), "sync distribution done", mlog.Duration("elapsed", elapsed))

	return nil
}

func (ex *Executor) removeDistribution(task *LeaderTask, step int) error {
	action := task.Actions()[step].(*LeaderAction)
	defer action.rpcReturned.Store(true)

	var err error
	defer func() {
		if err != nil {
			task.Fail(err)
		}
		ex.removeTask(task, step)
	}()

	req := &querypb.SyncDistributionRequest{
		Base: commonpbutil.NewMsgBase(
			commonpbutil.WithMsgType(commonpb.MsgType_SyncDistribution),
			commonpbutil.WithMsgID(task.ID()),
		),
		CollectionID: task.collectionID,
		Channel:      task.Shard(),
		ReplicaID:    task.ReplicaID(),
		Actions: []*querypb.SyncAction{
			{
				Type:      querypb.SyncType_Remove,
				SegmentID: action.SegmentID(),
				NodeID:    action.Node(),
			},
		},
	}

	startTs := time.Now()
	mlog.Info(context.TODO(), "Remove Distribution...")
	status, err := ex.cluster.SyncDistribution(task.Context(), task.leaderID, req)
	err = merr.CheckRPCCall(status, err)
	if err != nil {
		mlog.Warn(context.TODO(), "failed to remove distribution", mlog.Err(err))
		return err
	}

	elapsed := time.Since(startTs)
	mlog.Info(context.TODO(), "remove distribution done", mlog.Duration("elapsed", elapsed))

	return nil
}

func (ex *Executor) getMetaInfo(ctx context.Context, task Task) (*milvuspb.DescribeCollectionResponse, *querypb.LoadMetaInfo, *meta.DmChannel, error) {
	collectionID := task.CollectionID()
	shard := task.Shard()
	collectionInfo, err := ex.getCollectionInfo(ctx, collectionID)
	if err != nil {
		mlog.Warn(context.TODO(), "failed to get collection info", mlog.Err(err))
		return nil, nil, nil, err
	}
	loadFields := ex.meta.GetLoadFields(ctx, task.CollectionID())

	loadMeta := packLoadMeta(
		ex.meta.GetLoadType(ctx, task.CollectionID()),
		collectionInfo,
		task.ResourceGroup(),
		loadFields,
	)

	// get channel first, in case of target updated after segment info fetched
	channel := ex.targetMgr.GetDmChannel(ctx, collectionID, shard, meta.NextTargetFirst)
	if channel == nil {
		return nil, nil, nil, merr.WrapErrChannelNotAvailable(shard)
	}

	return collectionInfo, loadMeta, channel, nil
}

func (ex *Executor) getCollectionInfo(ctx context.Context, collectionID int64) (*milvuspb.DescribeCollectionResponse, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}

	ch := ex.collectionInfoSingleflight.DoChan(fmt.Sprint(collectionID), func() (*milvuspb.DescribeCollectionResponse, error) {
		return ex.broker.DescribeCollection(context.WithoutCancel(ctx), collectionID)
	})
	select {
	case result := <-ch:
		return result.Val, result.Err
	case <-ctx.Done():
		return nil, ctx.Err()
	}
}

func (ex *Executor) getLoadInfo(ctx context.Context, collectionID, segmentID int64, channel *meta.DmChannel, priority commonpb.LoadPriority) (*querypb.SegmentLoadInfo, []*indexpb.IndexInfo, error) {
	segmentInfos, err := ex.broker.GetSegmentInfo(ctx, segmentID)
	if err != nil || len(segmentInfos) == 0 {
		mlog.Warn(context.TODO(), "failed to get segment info from DataCoord", mlog.Err(err))
		return nil, nil, err
	}
	segment := segmentInfos[0]

	indexes, err := ex.broker.GetIndexInfo(ctx, collectionID, segment.GetID())
	if err != nil {
		if !errors.Is(err, merr.ErrIndexNotFound) {
			mlog.Warn(context.TODO(), "failed to get index of segment", mlog.Err(err))
			return nil, nil, err
		}
		indexes = nil
	}

	// Get collection index info
	indexInfos, err := ex.broker.ListIndexes(ctx, collectionID)
	if err != nil {
		mlog.Warn(context.TODO(), "fail to get index meta of collection", mlog.Err(err))
		return nil, nil, err
	}
	// update the field index params
	for _, segmentIndex := range indexes[segment.GetID()] {
		index, found := lo.Find(indexInfos, func(indexInfo *indexpb.IndexInfo) bool {
			return indexInfo.IndexID == segmentIndex.IndexID
		})
		if !found {
			mlog.Warn(context.TODO(), "no collection index info for the given segment index", mlog.String("indexName", segmentIndex.GetIndexName()))
		}

		params := funcutil.KeyValuePair2Map(segmentIndex.GetIndexParams())
		for _, kv := range index.GetUserIndexParams() {
			if indexparams.IsConfigableIndexParam(kv.GetKey()) {
				params[kv.GetKey()] = kv.GetValue()
			}
		}
		segmentIndex.IndexParams = funcutil.Map2KeyValuePair(params)
		segmentIndex.IndexParams = append(segmentIndex.IndexParams,
			&commonpb.KeyValuePair{Key: common.LoadPriorityKey, Value: priority.String()})
	}

	loadInfo := utils.PackSegmentLoadInfo(segment, channel.GetSeekPosition(), indexes[segment.GetID()])
	loadInfo.Priority = priority
	return loadInfo, indexInfos, nil
}