// 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 delegator

import (
	"bytes"
	"context"
	"fmt"
	"path/filepath"
	"strconv"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/bytedance/mockey"
	"github.com/cockroachdb/errors"
	"github.com/samber/lo"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/suite"
	"google.golang.org/protobuf/proto"

	"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
	"github.com/milvus-io/milvus-proto/go-api/v3/msgpb"
	"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
	"github.com/milvus-io/milvus/internal/flushcommon/syncmgr"
	"github.com/milvus-io/milvus/internal/mocks"
	"github.com/milvus-io/milvus/internal/mocks/util/mock_segcore"
	"github.com/milvus-io/milvus/internal/querynodev2/cluster"
	"github.com/milvus-io/milvus/internal/querynodev2/delegator/deletebuffer"
	"github.com/milvus-io/milvus/internal/querynodev2/pkoracle"
	"github.com/milvus-io/milvus/internal/querynodev2/segments"
	"github.com/milvus-io/milvus/internal/storage"
	"github.com/milvus-io/milvus/internal/util/bloomfilter"
	"github.com/milvus-io/milvus/internal/util/function"
	"github.com/milvus-io/milvus/internal/util/initcore"
	"github.com/milvus-io/milvus/pkg/v3/common"
	"github.com/milvus-io/milvus/pkg/v3/mlog"
	"github.com/milvus-io/milvus/pkg/v3/mq/msgstream"
	"github.com/milvus-io/milvus/pkg/v3/proto/datapb"
	"github.com/milvus-io/milvus/pkg/v3/proto/internalpb"
	"github.com/milvus-io/milvus/pkg/v3/proto/planpb"
	"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
	"github.com/milvus-io/milvus/pkg/v3/proto/segcorepb"
	"github.com/milvus-io/milvus/pkg/v3/util/commonpbutil"
	"github.com/milvus-io/milvus/pkg/v3/util/funcutil"
	"github.com/milvus-io/milvus/pkg/v3/util/merr"
	"github.com/milvus-io/milvus/pkg/v3/util/metautil"
	"github.com/milvus-io/milvus/pkg/v3/util/metric"
	"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
	"github.com/milvus-io/milvus/pkg/v3/util/syncutil"
	"github.com/milvus-io/milvus/pkg/v3/util/tsoutil"
	"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
)

// segmentEntryCoreFields extracts core fields from SegmentEntry for comparison,
// excluding Candidate field which varies between tests.
func segmentEntryCoreFields(entries []SegmentEntry) []SegmentEntry {
	result := make([]SegmentEntry, len(entries))
	for i, e := range entries {
		result[i] = SegmentEntry{
			NodeID:        e.NodeID,
			SegmentID:     e.SegmentID,
			PartitionID:   e.PartitionID,
			Version:       e.Version,
			TargetVersion: e.TargetVersion,
			Level:         e.Level,
			Offline:       e.Offline,
			// Candidate is intentionally excluded
		}
	}
	return result
}

type DelegatorDataSuite struct {
	suite.Suite

	collectionID  int64
	replicaID     int64
	vchannelName  string
	version       int64
	workerManager *cluster.MockManager
	manager       *segments.Manager
	loader        *segments.MockLoader
	mq            *msgstream.MockMsgStream
	channel       metautil.Channel
	mapper        metautil.ChannelMapper

	delegator    *shardDelegator
	rootPath     string
	chunkManager storage.ChunkManager
}

func (s *DelegatorDataSuite) getIDFOracleForTest() *idfOracle {
	oracle, ok := s.delegator.getIDFOracle().(*idfOracle)
	s.Require().True(ok)
	return oracle
}

func (s *DelegatorDataSuite) SetupSuite() {
	paramtable.Init()
	paramtable.SetNodeID(1)
	paramtable.Get().Save(paramtable.Get().QueryNodeCfg.CleanExcludeSegInterval.Key, "1")
	localDataRootPath := filepath.Join(paramtable.Get().LocalStorageCfg.Path.GetValue(), typeutil.QueryNodeRole)
	initcore.InitLocalChunkManager(localDataRootPath)
	initcore.InitMmapManager(paramtable.Get(), 1)
	initcore.InitTieredStorage(paramtable.Get())

	s.collectionID = 1000
	s.replicaID = 65535
	s.vchannelName = "rootcoord-dml_1000v0"
	s.version = 2000
	var err error
	s.mapper = metautil.NewDynChannelMapper()
	s.channel, err = metautil.ParseChannel(s.vchannelName, s.mapper)
	s.Require().NoError(err)
}

func (s *DelegatorDataSuite) TearDownSuite() {
	paramtable.Get().Reset(paramtable.Get().QueryNodeCfg.CleanExcludeSegInterval.Key)
}

func (s *DelegatorDataSuite) genNormalCollection() {
	s.manager.Collection.PutOrRef(s.collectionID, &schemapb.CollectionSchema{
		Name: "TestCollection",
		Fields: []*schemapb.FieldSchema{
			{
				Name:         "id",
				FieldID:      100,
				IsPrimaryKey: true,
				DataType:     schemapb.DataType_Int64,
				AutoID:       true,
			},
			{
				Name:         "vector",
				FieldID:      101,
				IsPrimaryKey: false,
				DataType:     schemapb.DataType_BinaryVector,
				TypeParams: []*commonpb.KeyValuePair{
					{
						Key:   common.DimKey,
						Value: "128",
					},
				},
			},
		},
	}, &segcorepb.CollectionIndexMeta{
		MaxIndexRowCount: 100,
		IndexMetas: []*segcorepb.FieldIndexMeta{
			{
				FieldID:      101,
				CollectionID: s.collectionID,
				IndexName:    "binary_index",
				TypeParams: []*commonpb.KeyValuePair{
					{
						Key:   common.DimKey,
						Value: "128",
					},
				},
				IndexParams: []*commonpb.KeyValuePair{
					{
						Key:   common.IndexTypeKey,
						Value: "BIN_IVF_FLAT",
					},
					{
						Key:   common.MetricTypeKey,
						Value: metric.JACCARD,
					},
				},
			},
		},
	}, &querypb.LoadMetaInfo{
		LoadType:        querypb.LoadType_LoadCollection,
		PartitionIDs:    []int64{1001, 1002},
		SchemaBarrierTs: tsoutil.ComposeTSByTime(time.Now()),
	})
}

func (s *DelegatorDataSuite) genTextCollection() {
	s.manager.Collection.PutOrRef(s.collectionID, &schemapb.CollectionSchema{
		Name: "TestTextCollection",
		Fields: []*schemapb.FieldSchema{
			{
				Name:         "id",
				FieldID:      100,
				IsPrimaryKey: true,
				DataType:     schemapb.DataType_Int64,
			},
			{
				Name:     "text",
				FieldID:  101,
				DataType: schemapb.DataType_Text,
			},
			{
				Name:     "vector",
				FieldID:  102,
				DataType: schemapb.DataType_FloatVector,
				TypeParams: []*commonpb.KeyValuePair{
					{
						Key:   common.DimKey,
						Value: "128",
					},
				},
			},
		},
	}, nil, &querypb.LoadMetaInfo{
		LoadType:        querypb.LoadType_LoadCollection,
		PartitionIDs:    []int64{1001},
		SchemaBarrierTs: tsoutil.ComposeTSByTime(time.Now()),
	})
}

func (s *DelegatorDataSuite) genCollectionWithFunction() {
	s.manager.Collection.PutOrRef(s.collectionID, &schemapb.CollectionSchema{
		Name:    "TestCollection",
		Version: 1,
		Fields: []*schemapb.FieldSchema{
			{
				Name:         "id",
				FieldID:      100,
				IsPrimaryKey: true,
				DataType:     schemapb.DataType_Int64,
				AutoID:       true,
			}, {
				Name:         "vector",
				FieldID:      101,
				IsPrimaryKey: false,
				DataType:     schemapb.DataType_SparseFloatVector,
			}, {
				Name:     "text",
				FieldID:  102,
				DataType: schemapb.DataType_VarChar,
				TypeParams: []*commonpb.KeyValuePair{
					{
						Key:   common.MaxLengthKey,
						Value: "256",
					},
				},
			},
		},
		Functions: []*schemapb.FunctionSchema{{
			Type:           schemapb.FunctionType_BM25,
			InputFieldIds:  []int64{102},
			OutputFieldIds: []int64{101},
		}},
	}, nil, &querypb.LoadMetaInfo{SchemaBarrierTs: tsoutil.ComposeTSByTime(time.Now())})

	delegator, err := NewShardDelegator(context.Background(), s.collectionID, s.replicaID, s.vchannelName, s.version, s.workerManager, s.manager, s.loader, 10000, nil, s.chunkManager, NewChannelQueryView(nil, nil, nil, initialTargetVersion), nil)
	s.NoError(err)
	s.delegator = delegator.(*shardDelegator)
	s.allocFunctionRunnersForTest()
}

func (s *DelegatorDataSuite) SetupTest() {
	paramtable.Get().Save(paramtable.Get().CommonCfg.EnableGrowingSourceFlush.Key, "false")
	s.T().Cleanup(func() {
		paramtable.Get().Reset(paramtable.Get().CommonCfg.EnableGrowingSourceFlush.Key)
	})

	s.workerManager = &cluster.MockManager{}
	s.manager = segments.NewManager()
	s.loader = &segments.MockLoader{}

	// init schema
	s.genNormalCollection()
	s.mq = &msgstream.MockMsgStream{}
	s.mq.EXPECT().GetUnmarshalDispatcher().Return(nil)

	s.rootPath = s.Suite.T().Name()
	chunkManagerFactory := storage.NewTestChunkManagerFactory(paramtable.Get(), s.rootPath)
	s.chunkManager, _ = chunkManagerFactory.NewPersistentStorageChunkManager(context.Background())
	delegator, err := NewShardDelegator(context.Background(), s.collectionID, s.replicaID, s.vchannelName, s.version, s.workerManager, s.manager, s.loader, 10000, nil, s.chunkManager, NewChannelQueryView(nil, nil, nil, initialTargetVersion), nil)
	s.Require().NoError(err)
	sd, ok := delegator.(*shardDelegator)
	s.Require().True(ok)
	s.delegator = sd
}

func (s *DelegatorDataSuite) enableGrowingSourceFlush() {
	paramtable.Get().Save(paramtable.Get().CommonCfg.UseLoonFFI.Key, "true")
	paramtable.Get().Save(paramtable.Get().CommonCfg.EnableGrowingSourceFlush.Key, "true")
	s.T().Cleanup(func() {
		paramtable.Get().Reset(paramtable.Get().CommonCfg.UseLoonFFI.Key)
		paramtable.Get().Reset(paramtable.Get().CommonCfg.EnableGrowingSourceFlush.Key)
	})
}

func (s *DelegatorDataSuite) TearDownTest() {
	function.ReleaseFunctionRunners(s.collectionID, "WAL-"+s.vchannelName)
	function.ReleaseFunctionRunners(s.collectionID, delegatorFunctionRunnerKey(s.vchannelName))
}

func (s *DelegatorDataSuite) allocFunctionRunnersForTest() {
	s.Require().NoError(function.UpdateFunctionRunners(s.collectionID, delegatorFunctionRunnerKey(s.vchannelName), s.delegator.collection.Schema()))
}

func (s *DelegatorDataSuite) TestProcessInsert() {
	s.Run("normal_insert", func() {
		s.delegator.ProcessInsert(map[int64]*InsertData{
			100: {
				RowIDs:        []int64{0, 1},
				PrimaryKeys:   []storage.PrimaryKey{storage.NewInt64PrimaryKey(1), storage.NewInt64PrimaryKey(2)},
				Timestamps:    []uint64{10, 10},
				PartitionID:   500,
				StartPosition: &msgpb.MsgPosition{},
				InsertRecord: &segcorepb.InsertRecord{
					FieldsData: []*schemapb.FieldData{
						{
							Type:      schemapb.DataType_Int64,
							FieldName: "id",
							Field: &schemapb.FieldData_Scalars{
								Scalars: &schemapb.ScalarField{
									Data: &schemapb.ScalarField_LongData{
										LongData: &schemapb.LongArray{
											Data: []int64{1, 2},
										},
									},
								},
							},
							FieldId: 100,
						},
						{
							Type:      schemapb.DataType_FloatVector,
							FieldName: "vector",
							Field: &schemapb.FieldData_Vectors{
								Vectors: &schemapb.VectorField{
									Dim: 128,
									Data: &schemapb.VectorField_FloatVector{
										FloatVector: &schemapb.FloatArray{Data: make([]float32, 128*2)},
									},
								},
							},
							FieldId: 101,
						},
					},
					NumRows: 2,
				},
			},
		})

		s.NotNil(s.manager.Segment.GetGrowing(100))
	})

	s.Run("insert_bad_data", func() {
		s.Panics(func() {
			s.delegator.ProcessInsert(map[int64]*InsertData{
				100: {
					RowIDs:        []int64{0, 1},
					PrimaryKeys:   []storage.PrimaryKey{storage.NewInt64PrimaryKey(1), storage.NewInt64PrimaryKey(2)},
					Timestamps:    []uint64{10, 10},
					PartitionID:   500,
					StartPosition: &msgpb.MsgPosition{},
					InsertRecord: &segcorepb.InsertRecord{
						FieldsData: []*schemapb.FieldData{
							{
								Type:      schemapb.DataType_Int64,
								FieldName: "id",
								Field: &schemapb.FieldData_Scalars{
									Scalars: &schemapb.ScalarField{
										Data: &schemapb.ScalarField_LongData{
											LongData: &schemapb.LongArray{
												Data: []int64{1, 2},
											},
										},
									},
								},
								FieldId: 100,
							},
						},
						NumRows: 2,
					},
				},
			})
		})
	})
}

func (s *DelegatorDataSuite) TestProcessDelete() {
	s.loader.EXPECT().
		Load(mock.Anything, s.collectionID, segments.SegmentTypeGrowing, int64(0), mock.Anything).
		Call.Return(func(ctx context.Context, collectionID int64, segmentType segments.SegmentType, version int64, infos ...*querypb.SegmentLoadInfo) []segments.Segment {
		return lo.Map(infos, func(info *querypb.SegmentLoadInfo, _ int) segments.Segment {
			ms := &segments.MockSegment{}
			ms.EXPECT().ID().Return(info.GetSegmentID())
			ms.EXPECT().Type().Return(segments.SegmentTypeGrowing)
			ms.EXPECT().Collection().Return(info.GetCollectionID())
			ms.EXPECT().Partition().Return(info.GetPartitionID())
			ms.EXPECT().Indexes().Return(nil)
			ms.EXPECT().RowNum().Return(info.GetNumOfRows())
			ms.EXPECT().LoadInfo().Return(info)
			ms.EXPECT().Delete(mock.Anything, mock.Anything, mock.Anything).Return(nil)
			ms.EXPECT().MayPkExist(mock.Anything).RunAndReturn(func(lc *storage.LocationsCache) bool {
				return lc.GetPk().EQ(storage.NewInt64PrimaryKey(10))
			})
			ms.EXPECT().PkCandidateExist().Return(true)
			ms.EXPECT().BatchPkExist(mock.Anything).RunAndReturn(func(lc *storage.BatchLocationsCache) []bool {
				hits := make([]bool, lc.Size())
				for i, pk := range lc.PKs() {
					hits[i] = pk.EQ(storage.NewInt64PrimaryKey(10))
				}
				return hits
			})
			return ms
		})
	}, nil)
	s.loader.EXPECT().LoadBloomFilterSet(mock.Anything, s.collectionID, mock.Anything).
		Call.Return(func(ctx context.Context, collectionID int64, infos ...*querypb.SegmentLoadInfo) []*pkoracle.BloomFilterSet {
		return lo.Map(infos, func(info *querypb.SegmentLoadInfo, _ int) *pkoracle.BloomFilterSet {
			bfs := pkoracle.NewBloomFilterSet(info.GetSegmentID(), info.GetPartitionID(), commonpb.SegmentState_Sealed)
			bf := bloomfilter.NewBloomFilterWithType(paramtable.Get().CommonCfg.BloomFilterSize.GetAsUint(),
				paramtable.Get().CommonCfg.MaxBloomFalsePositive.GetAsFloat(),
				paramtable.Get().CommonCfg.BloomFilterType.GetValue())
			pks := &storage.PkStatistics{
				PkFilter: bf,
			}
			pks.UpdatePKRange(&storage.Int64FieldData{
				Data: []int64{10, 20, 30},
			})
			bfs.AddHistoricalStats(pks)
			return bfs
		})
	}, func(ctx context.Context, collectionID int64, infos ...*querypb.SegmentLoadInfo) error {
		return nil
	})

	workers := make(map[int64]*cluster.MockWorker)
	worker1 := &cluster.MockWorker{}
	workers[1] = worker1

	worker1.EXPECT().LoadSegments(mock.Anything, mock.AnythingOfType("*querypb.LoadSegmentsRequest")).
		Return(nil)
	worker1.EXPECT().Delete(mock.Anything, mock.AnythingOfType("*querypb.DeleteRequest")).Return(nil)
	s.workerManager.EXPECT().GetWorker(mock.Anything, mock.AnythingOfType("int64")).Call.Return(func(_ context.Context, nodeID int64) cluster.Worker {
		return workers[nodeID]
	}, nil)
	// load growing
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	err := s.delegator.LoadGrowing(ctx, []*querypb.SegmentLoadInfo{
		{
			SegmentID:     1001,
			CollectionID:  s.collectionID,
			PartitionID:   500,
			InsertChannel: fmt.Sprintf("by-dev-rootcoord-dml_0_%dv0", s.collectionID),
		},
	}, 0)
	s.Require().NoError(err)
	// load sealed
	s.delegator.LoadSegments(ctx, &querypb.LoadSegmentsRequest{
		Base:         commonpbutil.NewMsgBase(),
		DstNodeID:    1,
		CollectionID: s.collectionID,
		Infos: []*querypb.SegmentLoadInfo{
			{
				SegmentID:     1000,
				CollectionID:  s.collectionID,
				PartitionID:   500,
				StartPosition: &msgpb.MsgPosition{Timestamp: 20000},
				DeltaPosition: &msgpb.MsgPosition{Timestamp: 20000},
				InsertChannel: fmt.Sprintf("by-dev-rootcoord-dml_0_%dv0", s.collectionID),
			},
		},
	})
	s.Require().NoError(err)

	s.delegator.ProcessDelete([]*DeleteData{
		{
			PartitionID: 500,
			PrimaryKeys: []storage.PrimaryKey{storage.NewInt64PrimaryKey(10)},
			Timestamps:  []uint64{10},
			RowCount:    1,
		},
	}, 10)

	// load sealed
	s.delegator.LoadSegments(ctx, &querypb.LoadSegmentsRequest{
		Base:         commonpbutil.NewMsgBase(),
		DstNodeID:    1,
		CollectionID: s.collectionID,
		Infos: []*querypb.SegmentLoadInfo{
			{
				SegmentID:     1000,
				CollectionID:  s.collectionID,
				PartitionID:   500,
				StartPosition: &msgpb.MsgPosition{Timestamp: 5000},
				DeltaPosition: &msgpb.MsgPosition{Timestamp: 5000},
				InsertChannel: fmt.Sprintf("by-dev-rootcoord-dml_0_%dv0", s.collectionID),
			},
		},
		Version: time.Now().UnixNano(),
	})
	s.Require().NoError(err)

	// sync target version, make delegator serviceable
	s.delegator.SyncTargetVersion(&querypb.SyncAction{
		TargetVersion:   2001,
		GrowingInTarget: []int64{1001},
		SealedSegmentRowCount: map[int64]int64{
			1000: 100,
		},
		DroppedInTarget: []int64{},
		Checkpoint:      &msgpb.MsgPosition{},
		DeleteCP:        &msgpb.MsgPosition{},
	}, []int64{500, 501})
	s.delegator.ProcessDelete([]*DeleteData{
		{
			PartitionID: 500,
			PrimaryKeys: []storage.PrimaryKey{storage.NewInt64PrimaryKey(10)},
			Timestamps:  []uint64{10},
			RowCount:    1,
		},
	}, 10)
	s.True(s.delegator.distribution.Serviceable())

	// test worker return segment not loaded
	worker1.ExpectedCalls = nil
	worker1.EXPECT().Delete(mock.Anything, mock.Anything).Return(merr.ErrSegmentNotLoaded)
	s.delegator.ProcessDelete([]*DeleteData{
		{
			PartitionID: 500,
			PrimaryKeys: []storage.PrimaryKey{storage.NewInt64PrimaryKey(10)},
			Timestamps:  []uint64{10},
			RowCount:    1,
		},
	}, 10)
	s.True(s.delegator.distribution.Serviceable(), "segment not loaded shall not trigger offline")

	// test worker offline
	worker1.ExpectedCalls = nil
	worker1.EXPECT().Delete(mock.Anything, mock.Anything).Return(merr.ErrNodeNotFound)
	s.delegator.ProcessDelete([]*DeleteData{
		{
			PartitionID: 500,
			PrimaryKeys: []storage.PrimaryKey{storage.NewInt64PrimaryKey(10)},
			Timestamps:  []uint64{10},
			RowCount:    1,
		},
	}, 10)

	s.delegator.distribution.Flush()
	s.False(s.delegator.distribution.Serviceable())

	worker1.EXPECT().LoadSegments(mock.Anything, mock.AnythingOfType("*querypb.LoadSegmentsRequest")).
		Return(nil)
	// reload, refresh the state
	s.delegator.LoadSegments(ctx, &querypb.LoadSegmentsRequest{
		Base:         commonpbutil.NewMsgBase(),
		DstNodeID:    1,
		CollectionID: s.collectionID,
		Infos: []*querypb.SegmentLoadInfo{
			{
				SegmentID:     1000,
				CollectionID:  s.collectionID,
				PartitionID:   500,
				StartPosition: &msgpb.MsgPosition{Timestamp: 20000},
				DeltaPosition: &msgpb.MsgPosition{Timestamp: 20000},
				InsertChannel: fmt.Sprintf("by-dev-rootcoord-dml_0_%dv0", s.collectionID),
			},
		},
		Version: time.Now().UnixNano(),
	})
	s.Require().NoError(err)
	s.delegator.distribution.Flush()
	s.True(s.delegator.distribution.Serviceable())
	// Test normal errors with retry and fail
	worker1.ExpectedCalls = nil
	worker1.EXPECT().Delete(mock.Anything, mock.Anything).Return(merr.ErrSegcore)
	s.delegator.ProcessDelete([]*DeleteData{
		{
			PartitionID: 500,
			PrimaryKeys: []storage.PrimaryKey{storage.NewInt64PrimaryKey(10)},
			Timestamps:  []uint64{10},
			RowCount:    1,
		},
	}, 10)
	s.delegator.distribution.Flush()
	s.False(s.delegator.distribution.Serviceable(), "should retry and failed")

	// refresh
	worker1.EXPECT().LoadSegments(mock.Anything, mock.AnythingOfType("*querypb.LoadSegmentsRequest")).
		Return(nil)
	// reload, refresh the state
	s.delegator.LoadSegments(ctx, &querypb.LoadSegmentsRequest{
		Base:         commonpbutil.NewMsgBase(),
		DstNodeID:    1,
		CollectionID: s.collectionID,
		Infos: []*querypb.SegmentLoadInfo{
			{
				SegmentID:     1000,
				CollectionID:  s.collectionID,
				PartitionID:   500,
				StartPosition: &msgpb.MsgPosition{Timestamp: 20000},
				DeltaPosition: &msgpb.MsgPosition{Timestamp: 20000},
				InsertChannel: fmt.Sprintf("by-dev-rootcoord-dml_0_%dv0", s.collectionID),
			},
		},
		Version: time.Now().UnixNano(),
	})
	s.Require().NoError(err)
	s.delegator.distribution.Flush()
	s.True(s.delegator.distribution.Serviceable())

	s.delegator.Close()
	// After Close(), ProcessDelete becomes a no-op because sd.Stopped() returns true.
	// This is expected behavior - the delegator is being decommissioned.
	// Note: Serviceable() state is not changed by ProcessDelete when delegator is stopped,
	// since ProcessDelete returns early without processing any deletes.
	s.delegator.ProcessDelete([]*DeleteData{
		{
			PartitionID: 500,
			PrimaryKeys: []storage.PrimaryKey{storage.NewInt64PrimaryKey(10)},
			Timestamps:  []uint64{10},
			RowCount:    1,
		},
	}, 10)
	s.Require().NoError(err)
	// Serviceable state remains unchanged since ProcessDelete is a no-op after Close()
	s.True(s.delegator.distribution.Serviceable())
}

func (s *DelegatorDataSuite) TestProcessDeleteBatchesPreservesBatchTsInDeleteBuffer() {
	batches := []DeleteBatch{
		{
			Ts: 10,
			Data: []*DeleteData{
				{
					PartitionID: 500,
					PrimaryKeys: []storage.PrimaryKey{
						storage.NewInt64PrimaryKey(10),
					},
					Timestamps: []uint64{10},
					RowCount:   1,
				},
			},
		},
		{
			Ts: 20,
			Data: []*DeleteData{
				{
					PartitionID: 500,
					PrimaryKeys: []storage.PrimaryKey{
						storage.NewInt64PrimaryKey(20),
					},
					Timestamps: []uint64{20},
					RowCount:   1,
				},
			},
		},
	}

	s.delegator.ProcessDeleteBatches(batches)

	after15 := s.delegator.deleteBuffer.ListAfter(15)
	s.Require().Len(after15, 1)
	s.Equal(uint64(20), after15[0].Ts)
	s.Require().Len(after15[0].Data, 1)
	s.ElementsMatch([]storage.PrimaryKey{storage.NewInt64PrimaryKey(20)}, after15[0].Data[0].DeleteData.Pks)

	after0 := s.delegator.deleteBuffer.ListAfter(0)
	s.Require().Len(after0, 2)
	s.Equal(uint64(10), after0[0].Ts)
	s.Equal(uint64(20), after0[1].Ts)
}

func (s *DelegatorDataSuite) TestLoadGrowingWithBM25() {
	s.genCollectionWithFunction()
	mockSegment := segments.NewMockSegment(s.T())
	s.loader.EXPECT().Load(mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return([]segments.Segment{mockSegment}, nil)

	mockSegment.EXPECT().Partition().Return(int64(111))
	mockSegment.EXPECT().ID().Return(int64(111))
	// Note: Type() is no longer called since pkOracle was removed
	// Candidate is now stored directly in SegmentEntry
	mockSegment.EXPECT().RowNum().Return(int64(0)).Maybe()
	mockSegment.EXPECT().LoadInfo().Return(&querypb.SegmentLoadInfo{SegmentID: 111}).Maybe()
	mockSegment.EXPECT().GetBM25Stats().Return(map[int64]*storage.BM25Stats{})

	err := s.delegator.LoadGrowing(context.Background(), []*querypb.SegmentLoadInfo{{SegmentID: 1}}, 1)
	s.NoError(err)
}

func (s *DelegatorDataSuite) TestLoadSegmentsWithBm25() {
	s.genCollectionWithFunction()
	s.Run("normal_run", func() {
		defer func() {
			s.workerManager.ExpectedCalls = nil
			s.loader.ExpectedCalls = nil
		}()

		s.loader.EXPECT().GetChunkManager().Return(nil)
		s.loader.EXPECT().LoadBloomFilterSet(mock.Anything, s.collectionID, mock.Anything).
			Call.Return(func(ctx context.Context, collectionID int64, infos ...*querypb.SegmentLoadInfo) []*pkoracle.BloomFilterSet {
			return lo.Map(infos, func(info *querypb.SegmentLoadInfo, _ int) *pkoracle.BloomFilterSet {
				return pkoracle.NewBloomFilterSet(info.GetSegmentID(), info.GetPartitionID(), commonpb.SegmentState_Sealed)
			})
		}, func(ctx context.Context, collectionID int64, infos ...*querypb.SegmentLoadInfo) error {
			return nil
		})

		workers := make(map[int64]*cluster.MockWorker)
		worker1 := &cluster.MockWorker{}
		workers[1] = worker1

		worker1.EXPECT().LoadSegments(mock.Anything, mock.AnythingOfType("*querypb.LoadSegmentsRequest")).
			Return(nil)
		s.workerManager.EXPECT().GetWorker(mock.Anything, mock.AnythingOfType("int64")).Call.Return(func(_ context.Context, nodeID int64) cluster.Worker {
			return workers[nodeID]
		}, nil)

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		err := s.delegator.LoadSegments(ctx, &querypb.LoadSegmentsRequest{
			Base:         commonpbutil.NewMsgBase(),
			DstNodeID:    1,
			CollectionID: s.collectionID,
			Infos: []*querypb.SegmentLoadInfo{
				{
					SegmentID:     100,
					PartitionID:   500,
					StartPosition: &msgpb.MsgPosition{Timestamp: 20000},
					DeltaPosition: &msgpb.MsgPosition{Timestamp: 20000},
					Level:         datapb.SegmentLevel_L1,
					InsertChannel: fmt.Sprintf("by-dev-rootcoord-dml_0_%dv0", s.collectionID),
				},
			},
		})

		s.NoError(err)
		s.delegator.distribution.Flush()
		sealed, _ := s.delegator.GetSegmentInfo(false)
		s.Require().Equal(1, len(sealed))
		s.Equal(int64(1), sealed[0].NodeID)
		s.ElementsMatch([]SegmentEntry{
			{
				SegmentID:     100,
				NodeID:        1,
				PartitionID:   500,
				TargetVersion: unreadableTargetVersion,
				Level:         datapb.SegmentLevel_L1,
			},
		}, segmentEntryCoreFields(sealed[0].Segments))
	})
}

func (s *DelegatorDataSuite) TestLoadSegmentsReopenBM25Stats() {
	s.genCollectionWithFunction()

	remotePath := "bm25stats/reopen/segment_100/field_101/0"
	stats := storage.NewBM25Stats()
	for i := uint32(1); i < 4; i++ {
		stats.Append(map[uint32]float32{i: 1})
	}
	data, err := stats.Serialize()
	s.Require().NoError(err)

	cm := mocks.NewChunkManager(s.T())
	cm.EXPECT().Reader(mock.Anything, remotePath).Return(&bytesFileReader{bytes.NewReader(data)}, nil)
	s.loader.EXPECT().GetChunkManager().Return(cm)

	worker1 := &cluster.MockWorker{}
	worker1.EXPECT().LoadSegments(mock.Anything, mock.AnythingOfType("*querypb.LoadSegmentsRequest")).Return(nil)
	s.workerManager.EXPECT().GetWorker(mock.Anything, mock.AnythingOfType("int64")).Return(worker1, nil)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	err = s.delegator.LoadSegments(ctx, &querypb.LoadSegmentsRequest{
		Base:         commonpbutil.NewMsgBase(),
		DstNodeID:    1,
		CollectionID: s.collectionID,
		LoadScope:    querypb.LoadScope_Reopen,
		Infos: []*querypb.SegmentLoadInfo{
			{
				SegmentID:     100,
				PartitionID:   500,
				StartPosition: &msgpb.MsgPosition{Timestamp: 20000},
				DeltaPosition: &msgpb.MsgPosition{Timestamp: 20000},
				Level:         datapb.SegmentLevel_L1,
				InsertChannel: fmt.Sprintf("by-dev-rootcoord-dml_0_%dv0", s.collectionID),
				Bm25Logs:      bm25LogsForField(101, remotePath),
			},
		},
	})

	s.NoError(err)
	loadedStats := s.getIDFOracleForTest()
	fieldStats, err := loadedStats.current.GetStats(101)
	s.NoError(err)
	s.Equal(int64(0), fieldStats.NumRow())
	segStats, ok := loadedStats.sealed.Get(100)
	s.True(ok)
	s.True(segStats.HasField(101))

	sealed, _ := s.delegator.GetSegmentInfo(false)
	s.Empty(sealed)
}

func (s *DelegatorDataSuite) TestLoadSegmentsReopenBM25StatsMergesReadableSegment() {
	s.genCollectionWithFunction()
	s.delegator.distribution.AddDistributions(SegmentEntry{
		NodeID:      1,
		SegmentID:   100,
		PartitionID: 500,
		Version:     1,
		Level:       datapb.SegmentLevel_L1,
	})
	s.delegator.distribution.SyncTargetVersion(&querypb.SyncAction{
		TargetVersion:         1,
		SealedInTarget:        []int64{100},
		SealedSegmentRowCount: map[int64]int64{100: 3},
	}, []int64{500})
	s.True(s.delegator.distribution.IsReadableSealedSegment(100))

	remotePath := "bm25stats/reopen-readable/segment_100/field_101/0"
	stats := storage.NewBM25Stats()
	for i := uint32(1); i < 4; i++ {
		stats.Append(map[uint32]float32{i: 1})
	}
	data, err := stats.Serialize()
	s.Require().NoError(err)

	cm := mocks.NewChunkManager(s.T())
	cm.EXPECT().Reader(mock.Anything, remotePath).Return(&bytesFileReader{bytes.NewReader(data)}, nil)
	s.loader.EXPECT().GetChunkManager().Return(cm)

	worker1 := &cluster.MockWorker{}
	worker1.EXPECT().LoadSegments(mock.Anything, mock.AnythingOfType("*querypb.LoadSegmentsRequest")).Return(nil)
	s.workerManager.EXPECT().GetWorker(mock.Anything, mock.AnythingOfType("int64")).Return(worker1, nil)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	err = s.delegator.LoadSegments(ctx, &querypb.LoadSegmentsRequest{
		Base:         commonpbutil.NewMsgBase(),
		DstNodeID:    1,
		CollectionID: s.collectionID,
		LoadScope:    querypb.LoadScope_Reopen,
		Infos: []*querypb.SegmentLoadInfo{
			{
				SegmentID:     100,
				PartitionID:   500,
				StartPosition: &msgpb.MsgPosition{Timestamp: 20000},
				DeltaPosition: &msgpb.MsgPosition{Timestamp: 20000},
				Level:         datapb.SegmentLevel_L1,
				InsertChannel: fmt.Sprintf("by-dev-rootcoord-dml_0_%dv0", s.collectionID),
				Bm25Logs:      bm25LogsForField(101, remotePath),
			},
		},
	})

	s.NoError(err)
	loadedStats := s.getIDFOracleForTest()
	fieldStats, err := loadedStats.current.GetStats(101)
	s.NoError(err)
	s.Equal(int64(3), fieldStats.NumRow())
}

func (s *DelegatorDataSuite) TestLoadSegmentsReopenBM25StatsRequiresOracleForRetry() {
	worker1 := &cluster.MockWorker{}
	worker1.EXPECT().LoadSegments(mock.Anything, mock.AnythingOfType("*querypb.LoadSegmentsRequest")).Return(nil)
	s.workerManager.EXPECT().GetWorker(mock.Anything, mock.AnythingOfType("int64")).Return(worker1, nil)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	err := s.delegator.LoadSegments(ctx, &querypb.LoadSegmentsRequest{
		Base:         commonpbutil.NewMsgBase(),
		DstNodeID:    1,
		CollectionID: s.collectionID,
		LoadScope:    querypb.LoadScope_Reopen,
		Infos: []*querypb.SegmentLoadInfo{
			{
				SegmentID:     100,
				PartitionID:   500,
				StartPosition: &msgpb.MsgPosition{Timestamp: 20000},
				DeltaPosition: &msgpb.MsgPosition{Timestamp: 20000},
				Level:         datapb.SegmentLevel_L1,
				InsertChannel: fmt.Sprintf("by-dev-rootcoord-dml_0_%dv0", s.collectionID),
				Bm25Logs:      bm25LogsForField(101, "bm25stats/reopen/segment_100/field_101/0"),
			},
		},
	})

	s.Error(err)
	s.ErrorIs(err, merr.ErrServiceInternal)
}

func (s *DelegatorDataSuite) TestLoadSegmentsReopenWithoutBM25StatsNoop() {
	worker1 := &cluster.MockWorker{}
	worker1.EXPECT().LoadSegments(mock.Anything, mock.AnythingOfType("*querypb.LoadSegmentsRequest")).Return(nil)
	s.workerManager.EXPECT().GetWorker(mock.Anything, mock.AnythingOfType("int64")).Return(worker1, nil)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	err := s.delegator.LoadSegments(ctx, &querypb.LoadSegmentsRequest{
		Base:         commonpbutil.NewMsgBase(),
		DstNodeID:    1,
		CollectionID: s.collectionID,
		LoadScope:    querypb.LoadScope_Reopen,
		Infos: []*querypb.SegmentLoadInfo{
			{
				SegmentID:     100,
				PartitionID:   500,
				StartPosition: &msgpb.MsgPosition{Timestamp: 20000},
				DeltaPosition: &msgpb.MsgPosition{Timestamp: 20000},
				Level:         datapb.SegmentLevel_L1,
				InsertChannel: fmt.Sprintf("by-dev-rootcoord-dml_0_%dv0", s.collectionID),
			},
		},
	})

	s.NoError(err)
	sealed, _ := s.delegator.GetSegmentInfo(false)
	s.Empty(sealed)
}

func (s *DelegatorDataSuite) TestLoadSegments() {
	s.Run("normal_run", func() {
		defer func() {
			s.workerManager.ExpectedCalls = nil
			s.loader.ExpectedCalls = nil
		}()

		s.loader.EXPECT().LoadBloomFilterSet(mock.Anything, s.collectionID, mock.Anything).
			Call.Return(func(ctx context.Context, collectionID int64, infos ...*querypb.SegmentLoadInfo) []*pkoracle.BloomFilterSet {
			return lo.Map(infos, func(info *querypb.SegmentLoadInfo, _ int) *pkoracle.BloomFilterSet {
				return pkoracle.NewBloomFilterSet(info.GetSegmentID(), info.GetPartitionID(), commonpb.SegmentState_Sealed)
			})
		}, func(ctx context.Context, collectionID int64, infos ...*querypb.SegmentLoadInfo) error {
			return nil
		})

		workers := make(map[int64]*cluster.MockWorker)
		worker1 := &cluster.MockWorker{}
		workers[1] = worker1

		worker1.EXPECT().LoadSegments(mock.Anything, mock.AnythingOfType("*querypb.LoadSegmentsRequest")).
			Return(nil)
		s.workerManager.EXPECT().GetWorker(mock.Anything, mock.AnythingOfType("int64")).Call.Return(func(_ context.Context, nodeID int64) cluster.Worker {
			return workers[nodeID]
		}, nil)

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		err := s.delegator.LoadSegments(ctx, &querypb.LoadSegmentsRequest{
			Base:         commonpbutil.NewMsgBase(),
			DstNodeID:    1,
			CollectionID: s.collectionID,
			Infos: []*querypb.SegmentLoadInfo{
				{
					SegmentID:     100,
					PartitionID:   500,
					StartPosition: &msgpb.MsgPosition{Timestamp: 20000},
					DeltaPosition: &msgpb.MsgPosition{Timestamp: 20000},
					Level:         datapb.SegmentLevel_L1,
					InsertChannel: fmt.Sprintf("by-dev-rootcoord-dml_0_%dv0", s.collectionID),
				},
			},
		})

		s.NoError(err)
		s.delegator.distribution.Flush()
		sealed, _ := s.delegator.GetSegmentInfo(false)
		s.Require().Equal(1, len(sealed))
		s.Equal(int64(1), sealed[0].NodeID)
		s.ElementsMatch([]SegmentEntry{
			{
				SegmentID:     100,
				NodeID:        1,
				PartitionID:   500,
				TargetVersion: unreadableTargetVersion,
				Level:         datapb.SegmentLevel_L1,
			},
		}, segmentEntryCoreFields(sealed[0].Segments))
	})

	s.Run("load_segments_with_delete", func() {
		defer func() {
			s.workerManager.ExpectedCalls = nil
			s.loader.ExpectedCalls = nil
		}()

		s.loader.EXPECT().LoadBloomFilterSet(mock.Anything, s.collectionID, mock.Anything).
			Call.Return(func(ctx context.Context, collectionID int64, infos ...*querypb.SegmentLoadInfo) []*pkoracle.BloomFilterSet {
			return lo.Map(infos, func(info *querypb.SegmentLoadInfo, _ int) *pkoracle.BloomFilterSet {
				bfs := pkoracle.NewBloomFilterSet(info.GetSegmentID(), info.GetPartitionID(), commonpb.SegmentState_Sealed)
				bf := bloomfilter.NewBloomFilterWithType(
					paramtable.Get().CommonCfg.BloomFilterSize.GetAsUint(),
					paramtable.Get().CommonCfg.MaxBloomFalsePositive.GetAsFloat(),
					paramtable.Get().CommonCfg.BloomFilterType.GetValue())
				pks := &storage.PkStatistics{
					PkFilter: bf,
				}
				pks.UpdatePKRange(&storage.Int64FieldData{
					Data: []int64{10, 20, 30},
				})
				bfs.AddHistoricalStats(pks)
				return bfs
			})
		}, func(ctx context.Context, collectionID int64, infos ...*querypb.SegmentLoadInfo) error {
			return nil
		})

		workers := make(map[int64]*cluster.MockWorker)
		worker1 := &cluster.MockWorker{}
		workers[1] = worker1

		worker1.EXPECT().LoadSegments(mock.Anything, mock.AnythingOfType("*querypb.LoadSegmentsRequest")).
			Return(nil)
		worker1.EXPECT().Delete(mock.Anything, mock.AnythingOfType("*querypb.DeleteRequest")).Return(nil)
		s.workerManager.EXPECT().GetWorker(mock.Anything, mock.AnythingOfType("int64")).Call.Return(func(_ context.Context, nodeID int64) cluster.Worker {
			return workers[nodeID]
		}, nil)

		s.delegator.ProcessDelete([]*DeleteData{
			{
				PartitionID: 500,
				PrimaryKeys: []storage.PrimaryKey{
					storage.NewInt64PrimaryKey(1),
					storage.NewInt64PrimaryKey(10),
				},
				Timestamps: []uint64{10, 10},
				RowCount:   2,
			},
		}, 10)

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		err := s.delegator.LoadSegments(ctx, &querypb.LoadSegmentsRequest{
			Base:         commonpbutil.NewMsgBase(),
			DstNodeID:    1,
			CollectionID: s.collectionID,
			Infos: []*querypb.SegmentLoadInfo{
				{
					SegmentID:     200,
					PartitionID:   500,
					StartPosition: &msgpb.MsgPosition{Timestamp: 20000},
					DeltaPosition: &msgpb.MsgPosition{Timestamp: 20000},
					Deltalogs:     []*datapb.FieldBinlog{},
					Level:         datapb.SegmentLevel_L0,
					InsertChannel: fmt.Sprintf("by-dev-rootcoord-dml_0_%dv0", s.collectionID),
				},
			},
		})
		s.ErrorIs(err, merr.ErrServiceInternal)
	})

	s.Run("load_segments_with_l0_delete_failed", func() {
		s.T().Skip("skip this test for now")
		defer func() {
			s.workerManager.ExpectedCalls = nil
			s.loader.ExpectedCalls = nil
		}()

		delegator, err := NewShardDelegator(
			context.Background(),
			s.collectionID,
			s.replicaID,
			s.vchannelName,
			s.version,
			s.workerManager,
			s.manager,
			s.loader,
			10000, nil, nil, NewChannelQueryView(nil, nil, nil, initialTargetVersion), nil)
		s.NoError(err)

		growing0 := segments.NewMockSegment(s.T())
		growing0.EXPECT().ID().Return(1)
		growing0.EXPECT().Partition().Return(10)
		growing0.EXPECT().Type().Return(segments.SegmentTypeGrowing)
		growing0.EXPECT().Release(context.Background())

		growing1 := segments.NewMockSegment(s.T())
		growing1.EXPECT().ID().Return(2)
		growing1.EXPECT().Partition().Return(10)
		growing1.EXPECT().Type().Return(segments.SegmentTypeGrowing)
		growing1.EXPECT().Release(context.Background())

		mockErr := merr.WrapErrServiceInternal("mock")

		growing0.EXPECT().Delete(mock.Anything, mock.Anything, mock.Anything).Return(nil)
		growing1.EXPECT().Delete(mock.Anything, mock.Anything, mock.Anything).Return(mockErr)

		s.loader.EXPECT().Load(
			mock.Anything,
			mock.Anything,
			segments.SegmentTypeGrowing,
			mock.Anything,
			mock.Anything,
			mock.Anything,
		).Return([]segments.Segment{growing0, growing1}, nil)

		err = delegator.LoadGrowing(context.Background(), []*querypb.SegmentLoadInfo{{}, {}}, 100)
		s.ErrorIs(err, mockErr)
	})

	s.Run("get_worker_fail", func() {
		defer func() {
			s.workerManager.ExpectedCalls = nil
			s.loader.ExpectedCalls = nil
		}()

		s.workerManager.EXPECT().GetWorker(mock.Anything, mock.AnythingOfType("int64")).Return(nil, errors.New("mock error"))

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		err := s.delegator.LoadSegments(ctx, &querypb.LoadSegmentsRequest{
			Base:         commonpbutil.NewMsgBase(),
			DstNodeID:    1,
			CollectionID: s.collectionID,
			Infos: []*querypb.SegmentLoadInfo{
				{
					SegmentID:     100,
					PartitionID:   500,
					StartPosition: &msgpb.MsgPosition{Timestamp: 20000},
					DeltaPosition: &msgpb.MsgPosition{Timestamp: 20000},
					InsertChannel: fmt.Sprintf("by-dev-rootcoord-dml_0_%dv0", s.collectionID),
				},
			},
		})

		s.Error(err)
	})

	s.Run("loader_bfs_fail", func() {
		defer func() {
			s.workerManager.ExpectedCalls = nil
			s.loader.ExpectedCalls = nil
		}()

		s.loader.EXPECT().LoadBloomFilterSet(mock.Anything, s.collectionID, mock.Anything).
			Return(nil, errors.New("mocked error"))

		workers := make(map[int64]*cluster.MockWorker)
		worker1 := &cluster.MockWorker{}
		workers[1] = worker1

		worker1.EXPECT().LoadSegments(mock.Anything, mock.AnythingOfType("*querypb.LoadSegmentsRequest")).
			Return(nil)
		s.workerManager.EXPECT().GetWorker(mock.Anything, mock.AnythingOfType("int64")).Call.Return(func(_ context.Context, nodeID int64) cluster.Worker {
			return workers[nodeID]
		}, nil)

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		err := s.delegator.LoadSegments(ctx, &querypb.LoadSegmentsRequest{
			Base:         commonpbutil.NewMsgBase(),
			DstNodeID:    1,
			CollectionID: s.collectionID,
			Infos: []*querypb.SegmentLoadInfo{
				{
					SegmentID:     100,
					PartitionID:   500,
					StartPosition: &msgpb.MsgPosition{Timestamp: 20000},
					DeltaPosition: &msgpb.MsgPosition{Timestamp: 20000},
					InsertChannel: fmt.Sprintf("by-dev-rootcoord-dml_0_%dv0", s.collectionID),
				},
			},
		})

		s.Error(err)
	})

	s.Run("worker_load_fail", func() {
		defer func() {
			s.workerManager.ExpectedCalls = nil
			s.loader.ExpectedCalls = nil
		}()

		s.loader.EXPECT().LoadBloomFilterSet(mock.Anything, s.collectionID, mock.Anything).
			Call.Return(func(ctx context.Context, collectionID int64, infos ...*querypb.SegmentLoadInfo) []*pkoracle.BloomFilterSet {
			return lo.Map(infos, func(info *querypb.SegmentLoadInfo, _ int) *pkoracle.BloomFilterSet {
				return pkoracle.NewBloomFilterSet(info.GetSegmentID(), info.GetPartitionID(), commonpb.SegmentState_Sealed)
			})
		}, func(ctx context.Context, collectionID int64, infos ...*querypb.SegmentLoadInfo) error {
			return nil
		})

		workers := make(map[int64]*cluster.MockWorker)
		worker1 := &cluster.MockWorker{}
		workers[1] = worker1

		worker1.EXPECT().LoadSegments(mock.Anything, mock.AnythingOfType("*querypb.LoadSegmentsRequest")).
			Return(errors.New("mocked error"))
		s.workerManager.EXPECT().GetWorker(mock.Anything, mock.AnythingOfType("int64")).Call.Return(func(_ context.Context, nodeID int64) cluster.Worker {
			return workers[nodeID]
		}, nil)

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		err := s.delegator.LoadSegments(ctx, &querypb.LoadSegmentsRequest{
			Base:         commonpbutil.NewMsgBase(),
			DstNodeID:    1,
			CollectionID: s.collectionID,
			Infos: []*querypb.SegmentLoadInfo{
				{
					SegmentID:     100,
					PartitionID:   500,
					StartPosition: &msgpb.MsgPosition{Timestamp: 20000},
					DeltaPosition: &msgpb.MsgPosition{Timestamp: 20000},
					InsertChannel: fmt.Sprintf("by-dev-rootcoord-dml_0_%dv0", s.collectionID),
				},
			},
		})

		s.Error(err)
	})
}

func (s *DelegatorDataSuite) TestLoadSegmentsWithoutBloomFilter() {
	defer func() {
		s.workerManager.ExpectedCalls = nil
		s.loader.ExpectedCalls = nil
	}()

	paramtable.Get().Save(paramtable.Get().CommonCfg.BloomFilterEnabled.Key, "false")
	defer paramtable.Get().Reset(paramtable.Get().CommonCfg.BloomFilterEnabled.Key)

	s.loader.EXPECT().LoadBloomFilterSet(mock.Anything, s.collectionID, mock.Anything).
		Call.Return(func(ctx context.Context, collectionID int64, infos ...*querypb.SegmentLoadInfo) []*pkoracle.BloomFilterSet {
		return lo.Map(infos, func(info *querypb.SegmentLoadInfo, _ int) *pkoracle.BloomFilterSet {
			return pkoracle.NewBloomFilterSet(info.GetSegmentID(), info.GetPartitionID(), commonpb.SegmentState_Sealed)
		})
	}, func(ctx context.Context, collectionID int64, infos ...*querypb.SegmentLoadInfo) error {
		return nil
	})

	workers := make(map[int64]*cluster.MockWorker)
	worker1 := &cluster.MockWorker{}
	workers[1] = worker1

	worker1.EXPECT().LoadSegments(mock.Anything, mock.AnythingOfType("*querypb.LoadSegmentsRequest")).
		Return(nil)
	s.workerManager.EXPECT().GetWorker(mock.Anything, mock.AnythingOfType("int64")).Call.Return(func(_ context.Context, nodeID int64) cluster.Worker {
		return workers[nodeID]
	}, nil)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	err := s.delegator.LoadSegments(ctx, &querypb.LoadSegmentsRequest{
		Base:         commonpbutil.NewMsgBase(),
		DstNodeID:    1,
		CollectionID: s.collectionID,
		Infos: []*querypb.SegmentLoadInfo{
			{
				SegmentID:     100,
				PartitionID:   500,
				StartPosition: &msgpb.MsgPosition{Timestamp: 20000},
				DeltaPosition: &msgpb.MsgPosition{Timestamp: 20000},
				Level:         datapb.SegmentLevel_L1,
				InsertChannel: fmt.Sprintf("by-dev-rootcoord-dml_0_%dv0", s.collectionID),
			},
		},
	})

	s.Require().NoError(err)
	s.delegator.distribution.Flush()
	sealed, _ := s.delegator.GetSegmentInfo(false)
	s.Require().Equal(1, len(sealed))
	s.Require().Equal(1, len(sealed[0].Segments))
	s.NotNil(sealed[0].Segments[0].Candidate)
	s.False(sealed[0].Segments[0].Candidate.PkCandidateExist())
}

func (s *DelegatorDataSuite) TestPostLoadLimiter() {
	s.Run("serializes_post_load_work", func() {
		sd := &shardDelegator{
			postLoadSem: syncutil.NewSemaphore(1),
		}
		var current int32
		var maxConcurrent int32
		start := make(chan struct{})
		errCh := make(chan error, 8)
		wg := sync.WaitGroup{}
		for i := 0; i < 8; i++ {
			wg.Add(1)
			go func() {
				defer wg.Done()
				<-start
				err := sd.withPostLoadLimit(context.Background(), func() error {
					running := atomic.AddInt32(&current, 1)
					for {
						maxValue := atomic.LoadInt32(&maxConcurrent)
						if running <= maxValue || atomic.CompareAndSwapInt32(&maxConcurrent, maxValue, running) {
							break
						}
					}
					time.Sleep(5 * time.Millisecond)
					atomic.AddInt32(&current, -1)
					return nil
				})
				errCh <- err
			}()
		}
		close(start)
		wg.Wait()
		close(errCh)
		for err := range errCh {
			s.NoError(err)
		}
		s.Equal(int32(1), maxConcurrent)
		s.Equal(0, sd.postLoadSem.Current())
	})

	s.Run("returns_context_error_while_waiting", func() {
		sd := &shardDelegator{
			postLoadSem: syncutil.NewSemaphore(1),
		}
		s.NoError(sd.postLoadSem.Acquire(context.Background()))
		defer sd.postLoadSem.Release()

		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Millisecond)
		defer cancel()
		called := false
		err := sd.withPostLoadLimit(ctx, func() error {
			called = true
			return nil
		})
		s.ErrorIs(err, context.DeadlineExceeded)
		s.False(called)
	})
}

func (s *DelegatorDataSuite) waitTargetVersion(targetVersion int64) {
	for {
		if s.getIDFOracleForTest().TargetVersion() >= targetVersion {
			return
		}
		time.Sleep(time.Millisecond * 100)
	}
}

func (s *DelegatorDataSuite) TestBuildBM25IDF() {
	s.genCollectionWithFunction()

	registerSealedStats := func(oracle *idfOracle, segID int64, start uint32, end uint32) {
		stats := storage.NewBM25Stats()
		for i := start; i < end; i++ {
			stats.Append(map[uint32]float32{i: 1})
		}
		data, err := stats.Serialize()
		s.Require().NoError(err)

		cm := mocks.NewChunkManager(s.T())
		remotePath := fmt.Sprintf("bm25stats/seg_%d/field_101/0", segID)
		cm.EXPECT().Reader(mock.Anything, remotePath).Return(
			&bytesFileReader{bytes.NewReader(data)}, nil,
		).Maybe()

		bm25Logs := []*datapb.FieldBinlog{{
			FieldID: 101,
			Binlogs: []*datapb.Binlog{{LogPath: remotePath}},
		}}

		err = oracle.LoadSealed(context.Background(), segID, &querypb.SegmentLoadInfo{Bm25Logs: bm25Logs}, cm)
		s.Require().NoError(err)
	}

	genSnapShot := func(seals, grows []int64, targetVersion int64) *snapshot {
		snapshot := &snapshot{
			dist:          []SnapshotItem{{1, make([]SegmentEntry, 0)}},
			targetVersion: targetVersion,
		}

		newSeal := []SegmentEntry{}
		for _, seg := range seals {
			newSeal = append(newSeal, SegmentEntry{NodeID: 1, SegmentID: seg, TargetVersion: targetVersion})
		}

		newGrow := []SegmentEntry{}
		for _, seg := range grows {
			newGrow = append(newGrow, SegmentEntry{NodeID: 1, SegmentID: seg, TargetVersion: targetVersion})
		}

		snapshot.dist[0].Segments = newSeal
		snapshot.growing = newGrow
		return snapshot
	}

	genStringFieldData := func(strs ...string) *schemapb.FieldData {
		return &schemapb.FieldData{
			Type:    schemapb.DataType_VarChar,
			FieldId: 102,
			Field: &schemapb.FieldData_Scalars{
				Scalars: &schemapb.ScalarField{
					Data: &schemapb.ScalarField_StringData{
						StringData: &schemapb.StringArray{
							Data: strs,
						},
					},
				},
			},
		}
	}

	s.Run("normal case", func() {
		// register sealed
		sealedSegs := []int64{1, 2, 3, 4}
		for _, segID := range sealedSegs {
			// every segment stats only has one token, avgdl = 1
			registerSealedStats(s.getIDFOracleForTest(), segID, uint32(segID), uint32(segID)+1)
		}
		snapshot := genSnapShot([]int64{1, 2, 3, 4}, []int64{}, 100)

		s.getIDFOracleForTest().SetNext(snapshot)
		s.waitTargetVersion(snapshot.targetVersion)
		placeholderGroupBytes, err := funcutil.FieldDataToPlaceholderGroupBytes(genStringFieldData("test bm25 data"))
		s.NoError(err)

		plan, err := proto.Marshal(&planpb.PlanNode{
			Node: &planpb.PlanNode_VectorAnns{
				VectorAnns: &planpb.VectorANNS{
					QueryInfo: &planpb.QueryInfo{},
				},
			},
		})
		s.NoError(err)

		req := &internalpb.SearchRequest{
			PlaceholderGroup:   placeholderGroupBytes,
			SerializedExprPlan: plan,
			FieldId:            101,
		}
		avgdl, err := s.delegator.buildBM25IDF(context.Background(), req)
		s.NoError(err)
		s.Equal(float64(1), avgdl)

		// check avgdl in plan
		newplan := &planpb.PlanNode{}
		err = proto.Unmarshal(req.GetSerializedExprPlan(), newplan)
		s.NoError(err)

		annplan, ok := newplan.GetNode().(*planpb.PlanNode_VectorAnns)
		s.Require().True(ok)
		s.Equal(avgdl, annplan.VectorAnns.QueryInfo.Bm25Avgdl)

		// check idf in placeholder
		placeholder := &commonpb.PlaceholderGroup{}
		err = proto.Unmarshal(req.GetPlaceholderGroup(), placeholder)
		s.Require().NoError(err)
		s.Equal(placeholder.GetPlaceholders()[0].GetType(), commonpb.PlaceholderType_SparseFloatVector)
	})

	s.Run("invalid place holder type error", func() {
		placeholderGroupBytes, err := proto.Marshal(&commonpb.PlaceholderGroup{
			Placeholders: []*commonpb.PlaceholderValue{{Type: commonpb.PlaceholderType_SparseFloatVector}},
		})
		s.NoError(err)

		req := &internalpb.SearchRequest{
			PlaceholderGroup: placeholderGroupBytes,
			FieldId:          101,
		}
		_, err = s.delegator.buildBM25IDF(context.Background(), req)
		s.Error(err)
	})

	s.Run("no function runner error", func() {
		placeholderGroupBytes, err := funcutil.FieldDataToPlaceholderGroupBytes(genStringFieldData("test bm25 data"))
		s.NoError(err)

		req := &internalpb.SearchRequest{
			PlaceholderGroup: placeholderGroupBytes,
			FieldId:          103, // invalid field id
		}

		_, err = s.delegator.buildBM25IDF(context.Background(), req)
		s.Error(err)
	})

	s.Run("function runner run failed error", func() {
		placeholderGroupBytes, err := funcutil.FieldDataToPlaceholderGroupBytes(genStringFieldData("test bm25 data"))
		s.NoError(err)

		mockRunner := function.NewMockFunctionRunner(s.T())
		mockRunner.EXPECT().GetInputFields().Return([]*schemapb.FieldSchema{
			{
				FieldID:  101,
				Name:     "text",
				DataType: schemapb.DataType_VarChar,
			},
		})
		mockRunner.EXPECT().BatchRun(mock.Anything).Return(nil, errors.New("mock err"))
		patch := mockey.Mock(function.RunWithRunner).To(func(ctx context.Context, collectionID int64, schemaVersion int32, outputFieldID int64, run func(schemapb.FunctionType, function.FunctionRunner) error) (bool, error) {
			return true, run(schemapb.FunctionType_BM25, mockRunner)
		}).Build()
		defer patch.UnPatch()

		req := &internalpb.SearchRequest{
			PlaceholderGroup: placeholderGroupBytes,
			FieldId:          101,
		}
		_, err = s.delegator.buildBM25IDF(context.Background(), req)
		s.Error(err)
	})

	s.Run("function runner output type error", func() {
		placeholderGroupBytes, err := funcutil.FieldDataToPlaceholderGroupBytes(genStringFieldData("test bm25 data"))
		s.NoError(err)

		mockRunner := function.NewMockFunctionRunner(s.T())
		mockRunner.EXPECT().GetInputFields().Return([]*schemapb.FieldSchema{
			{
				FieldID:  101,
				Name:     "text",
				DataType: schemapb.DataType_VarChar,
			},
		})
		mockRunner.EXPECT().BatchRun(mock.Anything).Return([]interface{}{1}, nil)
		patch := mockey.Mock(function.RunWithRunner).To(func(ctx context.Context, collectionID int64, schemaVersion int32, outputFieldID int64, run func(schemapb.FunctionType, function.FunctionRunner) error) (bool, error) {
			return true, run(schemapb.FunctionType_BM25, mockRunner)
		}).Build()
		defer patch.UnPatch()

		req := &internalpb.SearchRequest{
			PlaceholderGroup: placeholderGroupBytes,
			FieldId:          101,
		}
		_, err = s.delegator.buildBM25IDF(context.Background(), req)
		s.Error(err)
	})

	s.Run("idf oracle build idf error", func() {
		placeholderGroupBytes, err := funcutil.FieldDataToPlaceholderGroupBytes(genStringFieldData("test bm25 data"))
		s.NoError(err)

		mockRunner := function.NewMockFunctionRunner(s.T())
		mockRunner.EXPECT().GetInputFields().Return([]*schemapb.FieldSchema{
			{
				FieldID:  101,
				Name:     "text",
				DataType: schemapb.DataType_VarChar,
			},
		})
		mockRunner.EXPECT().BatchRun(mock.Anything).Return([]interface{}{&schemapb.SparseFloatArray{Contents: [][]byte{typeutil.CreateAndSortSparseFloatRow(map[uint32]float32{1: 1})}}}, nil)
		patch := mockey.Mock(function.RunWithRunner).To(func(ctx context.Context, collectionID int64, schemaVersion int32, outputFieldID int64, run func(schemapb.FunctionType, function.FunctionRunner) error) (bool, error) {
			return true, run(schemapb.FunctionType_BM25, mockRunner)
		}).Build()
		defer patch.UnPatch()

		req := &internalpb.SearchRequest{
			PlaceholderGroup: placeholderGroupBytes,
			FieldId:          103, // invalid field
		}
		_, err = s.delegator.buildBM25IDF(context.Background(), req)
		s.Error(err)
		mlog.Info(context.TODO(), "test", mlog.Err(err))
	})

	s.Run("set avgdl failed", func() {
		// register sealed
		sealedSegs := []int64{1, 2, 3, 4}
		for _, segID := range sealedSegs {
			// every segment stats only has one token, avgdl = 1
			registerSealedStats(s.getIDFOracleForTest(), segID, uint32(segID), uint32(segID)+1)
		}
		snapshot := genSnapShot([]int64{1, 2, 3, 4}, []int64{}, 100)

		s.getIDFOracleForTest().SetNext(snapshot)
		s.waitTargetVersion(snapshot.targetVersion)
		placeholderGroupBytes, err := funcutil.FieldDataToPlaceholderGroupBytes(genStringFieldData("test bm25 data"))
		s.NoError(err)

		req := &internalpb.SearchRequest{
			PlaceholderGroup: placeholderGroupBytes,
			FieldId:          101,
		}
		_, err = s.delegator.buildBM25IDF(context.Background(), req)
		s.Error(err)
	})
}

func (s *DelegatorDataSuite) TestReleaseSegment() {
	s.loader.EXPECT().
		Load(mock.Anything, s.collectionID, segments.SegmentTypeGrowing, int64(0), mock.Anything).
		Call.Return(func(ctx context.Context, collectionID int64, segmentType segments.SegmentType, version int64, infos ...*querypb.SegmentLoadInfo) []segments.Segment {
		return lo.Map(infos, func(info *querypb.SegmentLoadInfo, _ int) segments.Segment {
			ms := &segments.MockSegment{}
			ms.EXPECT().ID().Return(info.GetSegmentID())
			ms.EXPECT().Type().Return(segments.SegmentTypeGrowing)
			ms.EXPECT().Partition().Return(info.GetPartitionID())
			ms.EXPECT().Collection().Return(info.GetCollectionID())
			ms.EXPECT().Indexes().Return(nil)
			ms.EXPECT().RowNum().Return(info.GetNumOfRows())
			ms.EXPECT().LoadInfo().Return(info)
			ms.EXPECT().Delete(mock.Anything, mock.Anything, mock.Anything).Return(nil)
			ms.EXPECT().MayPkExist(mock.Anything).Call.Return(func(pk storage.PrimaryKey) bool {
				return pk.EQ(storage.NewInt64PrimaryKey(10))
			})
			return ms
		})
	}, nil)
	s.loader.EXPECT().LoadBloomFilterSet(mock.Anything, s.collectionID, mock.Anything).
		Call.Return(func(ctx context.Context, collectionID int64, infos ...*querypb.SegmentLoadInfo) []*pkoracle.BloomFilterSet {
		return lo.Map(infos, func(info *querypb.SegmentLoadInfo, _ int) *pkoracle.BloomFilterSet {
			bfs := pkoracle.NewBloomFilterSet(info.GetSegmentID(), info.GetPartitionID(), commonpb.SegmentState_Sealed)
			bf := bloomfilter.NewBloomFilterWithType(
				paramtable.Get().CommonCfg.BloomFilterSize.GetAsUint(),
				paramtable.Get().CommonCfg.MaxBloomFalsePositive.GetAsFloat(),
				paramtable.Get().CommonCfg.BloomFilterType.GetValue())
			pks := &storage.PkStatistics{
				PkFilter: bf,
			}
			pks.UpdatePKRange(&storage.Int64FieldData{
				Data: []int64{10, 20, 30},
			})
			bfs.AddHistoricalStats(pks)
			return bfs
		})
	}, func(ctx context.Context, collectionID int64, infos ...*querypb.SegmentLoadInfo) error {
		return nil
	})

	workers := make(map[int64]*cluster.MockWorker)
	worker1 := &cluster.MockWorker{}
	workers[1] = worker1
	worker2 := &cluster.MockWorker{}
	workers[2] = worker2

	worker1.EXPECT().LoadSegments(mock.Anything, mock.AnythingOfType("*querypb.LoadSegmentsRequest")).
		Return(nil)
	worker1.EXPECT().ReleaseSegments(mock.Anything, mock.AnythingOfType("*querypb.ReleaseSegmentsRequest")).Return(nil)
	worker2.EXPECT().ReleaseSegments(mock.Anything, mock.AnythingOfType("*querypb.ReleaseSegmentsRequest")).Return(nil)
	s.workerManager.EXPECT().GetWorker(mock.Anything, mock.AnythingOfType("int64")).Call.Return(func(_ context.Context, nodeID int64) cluster.Worker {
		return workers[nodeID]
	}, nil)
	// load growing
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	err := s.delegator.LoadGrowing(ctx, []*querypb.SegmentLoadInfo{
		{
			SegmentID:     1001,
			CollectionID:  s.collectionID,
			PartitionID:   500,
			InsertChannel: fmt.Sprintf("by-dev-rootcoord-dml_0_%dv0", s.collectionID),
		},
	}, 0)
	s.Require().NoError(err)
	// load sealed
	s.delegator.LoadSegments(ctx, &querypb.LoadSegmentsRequest{
		Base:         commonpbutil.NewMsgBase(),
		DstNodeID:    1,
		CollectionID: s.collectionID,
		Infos: []*querypb.SegmentLoadInfo{
			{
				SegmentID:     1000,
				CollectionID:  s.collectionID,
				PartitionID:   500,
				StartPosition: &msgpb.MsgPosition{Timestamp: 20000},
				DeltaPosition: &msgpb.MsgPosition{Timestamp: 20000},
				InsertChannel: fmt.Sprintf("by-dev-rootcoord-dml_0_%dv0", s.collectionID),
			},
		},
	})
	s.Require().NoError(err)

	s.delegator.distribution.Flush()
	sealed, growing := s.delegator.GetSegmentInfo(false)
	s.Require().Equal(1, len(sealed))
	s.Equal(int64(1), sealed[0].NodeID)
	s.ElementsMatch([]SegmentEntry{
		{
			SegmentID:     1000,
			NodeID:        1,
			PartitionID:   500,
			TargetVersion: unreadableTargetVersion,
		},
	}, segmentEntryCoreFields(sealed[0].Segments))

	s.ElementsMatch([]SegmentEntry{
		{
			SegmentID:     1001,
			NodeID:        1,
			PartitionID:   500,
			TargetVersion: unreadableTargetVersion,
		},
	}, segmentEntryCoreFields(growing))

	err = s.delegator.ReleaseSegments(ctx, &querypb.ReleaseSegmentsRequest{
		Base:       commonpbutil.NewMsgBase(),
		NodeID:     1,
		SegmentIDs: []int64{1000},
		Scope:      querypb.DataScope_Historical,
	}, false)

	s.NoError(err)
	sealed, _ = s.delegator.GetSegmentInfo(false)
	s.Equal(0, len(sealed))

	err = s.delegator.ReleaseSegments(ctx, &querypb.ReleaseSegmentsRequest{
		Base:       commonpbutil.NewMsgBase(),
		NodeID:     1,
		SegmentIDs: []int64{1001},
		Scope:      querypb.DataScope_Streaming,
	}, false)

	s.NoError(err)
	_, growing = s.delegator.GetSegmentInfo(false)
	s.Equal(0, len(growing))

	err = s.delegator.ReleaseSegments(ctx, &querypb.ReleaseSegmentsRequest{
		Base:       commonpbutil.NewMsgBase(),
		NodeID:     1,
		SegmentIDs: []int64{1000},
		Scope:      querypb.DataScope_All,
	}, true)
	s.NoError(err)

	// test transfer
	req := &querypb.ReleaseSegmentsRequest{
		Base:         commonpbutil.NewMsgBase(),
		NodeID:       2,
		SegmentIDs:   []int64{1000},
		Scope:        querypb.DataScope_All,
		NeedTransfer: true,
	}
	req.Base.TargetID = 1
	err = s.delegator.ReleaseSegments(ctx, req, false)
	s.NoError(err)
}

func (s *DelegatorDataSuite) TestReleaseGrowingSourceAfterPreparedHandoff() {
	ctx := context.Background()
	s.workerManager = &cluster.MockManager{}
	s.manager = segments.NewManager()
	s.loader = &segments.MockLoader{}
	s.genTextCollection()
	s.enableGrowingSourceFlush()

	delegator, err := NewShardDelegator(
		ctx,
		s.collectionID,
		s.replicaID,
		s.vchannelName,
		s.version,
		s.workerManager,
		s.manager,
		s.loader,
		10000,
		nil,
		s.chunkManager,
		NewChannelQueryView(nil, nil, nil, initialTargetVersion),
		nil)
	s.Require().NoError(err)
	sd := delegator.(*shardDelegator)
	defer sd.Close()

	const (
		segmentID    = int64(1001)
		partitionID  = int64(500)
		targetOffset = int64(10)
	)

	segment := segments.NewMockSegment(s.T())
	segment.EXPECT().ID().Return(segmentID).Maybe()
	segment.EXPECT().Version().Return(int64(1)).Maybe()
	segment.EXPECT().Shard().Return(s.channel).Maybe()
	segment.EXPECT().Collection().Return(s.collectionID).Maybe()
	segment.EXPECT().Partition().Return(partitionID).Maybe()
	segment.EXPECT().Type().Return(segments.SegmentTypeGrowing).Maybe()
	segment.EXPECT().Level().Return(datapb.SegmentLevel_Legacy).Maybe()
	segment.EXPECT().PinIfNotReleased().Return(nil).Once()
	segment.EXPECT().InsertCount().Return(targetOffset).Once()
	segment.EXPECT().Unpin().Maybe()

	s.manager.Segment.Put(ctx, segments.SegmentTypeGrowing, segment)
	sd.distribution.AddGrowing(SegmentEntry{
		NodeID:      paramtable.GetNodeID(),
		SegmentID:   segmentID,
		PartitionID: partitionID,
		Version:     1,
	})
	worker := &cluster.MockWorker{}
	worker.EXPECT().ReleaseSegments(mock.Anything, mock.AnythingOfType("*querypb.ReleaseSegmentsRequest")).Return(nil).Once()
	s.workerManager.EXPECT().GetWorker(mock.Anything, paramtable.GetNodeID()).Return(worker, nil).Once()

	err = sd.growingSourceProvider.PrepareGrowingSourceReleaseHandoff(ctx, 10000, []syncmgr.GrowingSourceReleaseHandoffSegment{
		{
			SegmentID:    segmentID,
			TargetOffset: targetOffset,
		},
	})
	s.Require().NoError(err)

	err = sd.ReleaseSegments(ctx, &querypb.ReleaseSegmentsRequest{
		Base:         commonpbutil.NewMsgBase(),
		NodeID:       paramtable.GetNodeID(),
		CollectionID: s.collectionID,
		SegmentIDs:   []int64{segmentID},
		Scope:        querypb.DataScope_Streaming,
		Shard:        s.vchannelName,
	}, false)
	s.Require().NoError(err)
}

func (s *DelegatorDataSuite) TestReleaseGrowingSourceAfterFencePreparedHandoff() {
	ctx := context.Background()
	s.workerManager = &cluster.MockManager{}
	s.manager = segments.NewManager()
	s.loader = &segments.MockLoader{}
	s.genTextCollection()
	s.enableGrowingSourceFlush()

	delegator, err := NewShardDelegator(
		ctx,
		s.collectionID,
		s.replicaID,
		s.vchannelName,
		s.version,
		s.workerManager,
		s.manager,
		s.loader,
		10000,
		nil,
		s.chunkManager,
		NewChannelQueryView(nil, nil, nil, initialTargetVersion),
		nil)
	s.Require().NoError(err)
	sd := delegator.(*shardDelegator)
	defer sd.Close()

	const segmentID = int64(1001)
	sd.distribution.AddGrowing(SegmentEntry{
		NodeID:      paramtable.GetNodeID(),
		SegmentID:   segmentID,
		PartitionID: 500,
		Version:     1,
	})

	err = sd.growingSourceProvider.PrepareGrowingSourceReleaseHandoff(ctx, 10000, []syncmgr.GrowingSourceReleaseHandoffSegment{
		{SegmentID: segmentID},
	})
	s.Require().NoError(err)
	worker := &cluster.MockWorker{}
	worker.EXPECT().ReleaseSegments(mock.Anything, mock.AnythingOfType("*querypb.ReleaseSegmentsRequest")).Return(nil).Once()
	s.workerManager.EXPECT().GetWorker(mock.Anything, paramtable.GetNodeID()).Return(worker, nil).Once()

	err = sd.ReleaseSegments(ctx, &querypb.ReleaseSegmentsRequest{
		Base:         commonpbutil.NewMsgBase(),
		NodeID:       paramtable.GetNodeID(),
		CollectionID: s.collectionID,
		SegmentIDs:   []int64{segmentID},
		Scope:        querypb.DataScope_Streaming,
		Shard:        s.vchannelName,
		Checkpoint:   &msgpb.MsgPosition{Timestamp: 10000},
	}, false)
	s.Require().NoError(err)
}

func (s *DelegatorDataSuite) TestReleaseGrowingSourceAfterNoRetainPreparedHandoff() {
	ctx := context.Background()
	s.workerManager = &cluster.MockManager{}
	s.manager = segments.NewManager()
	s.loader = &segments.MockLoader{}
	s.genTextCollection()
	s.enableGrowingSourceFlush()

	delegator, err := NewShardDelegator(
		ctx,
		s.collectionID,
		s.replicaID,
		s.vchannelName,
		s.version,
		s.workerManager,
		s.manager,
		s.loader,
		10000,
		nil,
		s.chunkManager,
		NewChannelQueryView(nil, nil, nil, initialTargetVersion),
		nil)
	s.Require().NoError(err)
	sd := delegator.(*shardDelegator)
	defer sd.Close()

	const segmentID = int64(1001)
	sd.distribution.AddGrowing(SegmentEntry{
		NodeID:      paramtable.GetNodeID(),
		SegmentID:   segmentID,
		PartitionID: 500,
		Version:     1,
	})

	err = sd.growingSourceProvider.PrepareGrowingSourceReleaseHandoff(ctx, 10000, []syncmgr.GrowingSourceReleaseHandoffSegment{
		{SegmentID: segmentID},
	})
	s.Require().NoError(err)

	worker := &cluster.MockWorker{}
	worker.EXPECT().ReleaseSegments(mock.Anything, mock.AnythingOfType("*querypb.ReleaseSegmentsRequest")).Return(nil).Once()
	s.workerManager.EXPECT().GetWorker(mock.Anything, paramtable.GetNodeID()).Return(worker, nil).Once()

	err = sd.ReleaseSegments(ctx, &querypb.ReleaseSegmentsRequest{
		Base:         commonpbutil.NewMsgBase(),
		NodeID:       paramtable.GetNodeID(),
		CollectionID: s.collectionID,
		SegmentIDs:   []int64{segmentID},
		Scope:        querypb.DataScope_Streaming,
		Shard:        s.vchannelName,
		Checkpoint:   &msgpb.MsgPosition{Timestamp: 10000},
	}, false)
	s.Require().NoError(err)
}

func (s *DelegatorDataSuite) TestReleaseGrowingSourceWithoutPreparedHandoff() {
	ctx := context.Background()
	s.workerManager = &cluster.MockManager{}
	s.manager = segments.NewManager()
	s.loader = &segments.MockLoader{}
	s.genTextCollection()
	s.enableGrowingSourceFlush()

	delegator, err := NewShardDelegator(
		ctx,
		s.collectionID,
		s.replicaID,
		s.vchannelName,
		s.version,
		s.workerManager,
		s.manager,
		s.loader,
		10000,
		nil,
		s.chunkManager,
		NewChannelQueryView(nil, nil, nil, initialTargetVersion),
		nil)
	s.Require().NoError(err)
	sd := delegator.(*shardDelegator)
	defer sd.Close()

	const segmentID = int64(1001)
	sd.distribution.AddGrowing(SegmentEntry{
		NodeID:      paramtable.GetNodeID(),
		SegmentID:   segmentID,
		PartitionID: 500,
		Version:     1,
	})

	worker := &cluster.MockWorker{}
	worker.EXPECT().ReleaseSegments(mock.Anything, mock.AnythingOfType("*querypb.ReleaseSegmentsRequest")).Return(nil).Once()
	s.workerManager.EXPECT().GetWorker(mock.Anything, paramtable.GetNodeID()).Return(worker, nil).Once()

	err = sd.ReleaseSegments(ctx, &querypb.ReleaseSegmentsRequest{
		Base:         commonpbutil.NewMsgBase(),
		NodeID:       paramtable.GetNodeID(),
		CollectionID: s.collectionID,
		SegmentIDs:   []int64{segmentID},
		Scope:        querypb.DataScope_Streaming,
		Shard:        s.vchannelName,
		Checkpoint:   &msgpb.MsgPosition{Timestamp: 10000},
	}, false)
	s.Require().NoError(err)
}

func (s *DelegatorDataSuite) TestReleaseGrowingSourceDroppedChannelWithoutPreparedHandoff() {
	ctx := context.Background()
	s.workerManager = &cluster.MockManager{}
	s.manager = segments.NewManager()
	s.loader = &segments.MockLoader{}
	s.genTextCollection()
	s.enableGrowingSourceFlush()

	delegator, err := NewShardDelegator(
		ctx,
		s.collectionID,
		s.replicaID,
		s.vchannelName,
		s.version,
		s.workerManager,
		s.manager,
		s.loader,
		10000,
		nil,
		s.chunkManager,
		NewChannelQueryView(nil, nil, nil, initialTargetVersion),
		nil)
	s.Require().NoError(err)
	sd := delegator.(*shardDelegator)
	defer sd.Close()

	const segmentID = int64(1001)
	sd.distribution.AddGrowing(SegmentEntry{
		NodeID:      paramtable.GetNodeID(),
		SegmentID:   segmentID,
		PartitionID: 500,
		Version:     1,
	})

	worker := &cluster.MockWorker{}
	worker.EXPECT().ReleaseSegments(mock.Anything, mock.AnythingOfType("*querypb.ReleaseSegmentsRequest")).Return(nil).Once()
	s.workerManager.EXPECT().GetWorker(mock.Anything, paramtable.GetNodeID()).Return(worker, nil).Once()

	err = sd.ReleaseSegments(ctx, &querypb.ReleaseSegmentsRequest{
		Base:         commonpbutil.NewMsgBase(),
		NodeID:       paramtable.GetNodeID(),
		CollectionID: s.collectionID,
		SegmentIDs:   []int64{segmentID},
		Scope:        querypb.DataScope_Streaming,
		Shard:        s.vchannelName,
		Checkpoint:   &msgpb.MsgPosition{Timestamp: typeutil.MaxTimestamp},
	}, false)
	s.Require().NoError(err)
}

func (s *DelegatorDataSuite) TestLoadPartitionStats() {
	segStats := make(map[UniqueID]storage.SegmentStats)
	centroid := []float32{1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0}
	var segID int64 = 1
	rows := 1990
	{
		// p1 stats
		fieldStats := make([]storage.FieldStats, 0)
		fieldStat1 := storage.FieldStats{
			FieldID: 1,
			Type:    schemapb.DataType_Int64,
			Max:     storage.NewInt64FieldValue(200),
			Min:     storage.NewInt64FieldValue(100),
		}
		fieldStat2 := storage.FieldStats{
			FieldID: 2,
			Type:    schemapb.DataType_Int64,
			Max:     storage.NewInt64FieldValue(400),
			Min:     storage.NewInt64FieldValue(300),
		}
		fieldStat3 := storage.FieldStats{
			FieldID: 3,
			Type:    schemapb.DataType_FloatVector,
			Centroids: []storage.VectorFieldValue{
				&storage.FloatVectorFieldValue{
					Value: centroid,
				},
				&storage.FloatVectorFieldValue{
					Value: centroid,
				},
			},
		}
		fieldStats = append(fieldStats, fieldStat1)
		fieldStats = append(fieldStats, fieldStat2)
		fieldStats = append(fieldStats, fieldStat3)
		segStats[segID] = *storage.NewSegmentStats(fieldStats, rows)
	}
	partitionStats1 := &storage.PartitionStatsSnapshot{
		SegmentStats: segStats,
	}
	statsData1, err := storage.SerializePartitionStatsSnapshot(partitionStats1)
	s.NoError(err)
	partitionID1 := int64(1001)
	idPath1 := metautil.JoinIDPath(s.collectionID, partitionID1)
	idPath1 = filepath.Join(idPath1, s.delegator.vchannelName)
	statsPath1 := filepath.Join(s.chunkManager.RootPath(), common.PartitionStatsPath, idPath1, strconv.Itoa(1))
	s.chunkManager.Write(context.Background(), statsPath1, statsData1)
	defer s.chunkManager.Remove(context.Background(), statsPath1)

	// reload and check partition stats
	partVersions := make(map[int64]int64)
	partVersions[partitionID1] = 1
	s.delegator.loadPartitionStats(context.Background(), partVersions)
	s.Equal(1, len(s.delegator.partitionStats))
	s.NotNil(s.delegator.partitionStats[partitionID1])
	p1Stats := s.delegator.partitionStats[partitionID1]
	s.Equal(int64(1), p1Stats.GetVersion())
	s.Equal(rows, p1Stats.SegmentStats[segID].NumRows)
	s.Equal(3, len(p1Stats.SegmentStats[segID].FieldStats))

	// judge vector stats
	vecFieldStats := p1Stats.SegmentStats[segID].FieldStats[2]
	s.Equal(2, len(vecFieldStats.Centroids))
	s.Equal(8, len(vecFieldStats.Centroids[0].GetValue().([]float32)))

	// judge scalar stats
	fieldStats1 := p1Stats.SegmentStats[segID].FieldStats[0]
	s.Equal(int64(100), fieldStats1.Min.GetValue().(int64))
	s.Equal(int64(200), fieldStats1.Max.GetValue().(int64))
	fieldStats2 := p1Stats.SegmentStats[segID].FieldStats[1]
	s.Equal(int64(300), fieldStats2.Min.GetValue().(int64))
	s.Equal(int64(400), fieldStats2.Max.GetValue().(int64))
}

func (s *DelegatorDataSuite) TestSyncTargetVersion() {
	for i := int64(0); i < 5; i++ {
		ms := &segments.MockSegment{}
		ms.EXPECT().ID().Return(i)
		ms.EXPECT().StartPosition().Return(&msgpb.MsgPosition{
			Timestamp: uint64(i),
		})
		ms.EXPECT().Type().Return(segments.SegmentTypeGrowing)
		ms.EXPECT().Collection().Return(1)
		ms.EXPECT().Partition().Return(1)
		ms.EXPECT().InsertCount().Return(0)
		ms.EXPECT().Indexes().Return(nil)
		ms.EXPECT().Shard().Return(s.channel)
		ms.EXPECT().Level().Return(datapb.SegmentLevel_L1)
		s.manager.Segment.Put(context.Background(), segments.SegmentTypeGrowing, ms)
	}

	s.delegator.SyncTargetVersion(&querypb.SyncAction{
		TargetVersion:   5,
		GrowingInTarget: []int64{1},
		SealedInTarget:  []int64{2},
		DroppedInTarget: []int64{3, 4},
		Checkpoint:      &msgpb.MsgPosition{},
		DeleteCP:        &msgpb.MsgPosition{},
	}, []int64{500, 501})
	s.Equal(int64(5), s.delegator.GetChannelQueryView().GetVersion())
}

func (s *DelegatorDataSuite) TestLevel0Deletions() {
	delegator := s.delegator
	partitionID := int64(10)
	partitionDeleteData, err := storage.NewDeltaDataWithPkType(1, schemapb.DataType_Int64)
	s.Require().NoError(err)
	err = partitionDeleteData.Append(storage.NewInt64PrimaryKey(1), 100)
	s.Require().NoError(err)

	allPartitionDeleteData, err := storage.NewDeltaDataWithPkType(1, schemapb.DataType_Int64)
	s.Require().NoError(err)
	err = allPartitionDeleteData.Append(storage.NewInt64PrimaryKey(2), 101)
	s.Require().NoError(err)

	schema := mock_segcore.GenTestCollectionSchema("test_stop", schemapb.DataType_Int64, true)
	collection, err := segments.NewCollection(1, schema, nil, &querypb.LoadMetaInfo{
		LoadType: querypb.LoadType_LoadCollection,
	})
	s.NoError(err)

	l0, _ := segments.NewL0Segment(collection, segments.SegmentTypeSealed, 1, &querypb.SegmentLoadInfo{
		CollectionID:  1,
		SegmentID:     2,
		PartitionID:   partitionID,
		InsertChannel: delegator.vchannelName,
		Level:         datapb.SegmentLevel_L0,
		NumOfRows:     1,
		StartPosition: &msgpb.MsgPosition{
			Timestamp: 10,
		},
	})
	l0.LoadDeltaData(context.TODO(), partitionDeleteData)
	delegator.deleteBuffer.RegisterL0(l0)

	l0Global, _ := segments.NewL0Segment(collection, segments.SegmentTypeSealed, 2, &querypb.SegmentLoadInfo{
		CollectionID:  1,
		SegmentID:     3,
		PartitionID:   common.AllPartitionsID,
		InsertChannel: delegator.vchannelName,
		Level:         datapb.SegmentLevel_L0,
		NumOfRows:     int64(1),
		StartPosition: &msgpb.MsgPosition{
			Timestamp: 20,
		},
	})
	l0Global.LoadDeltaData(context.TODO(), allPartitionDeleteData)

	pks, _ := delegator.GetLevel0Deletions(partitionID, pkoracle.NewCandidateKey(l0.ID(), l0.Partition(), segments.SegmentTypeGrowing))
	s.True(pks.Get(0).EQ(partitionDeleteData.DeletePks().Get(0)))

	pks, _ = delegator.GetLevel0Deletions(partitionID+1, pkoracle.NewCandidateKey(l0.ID(), l0.Partition(), segments.SegmentTypeGrowing))
	s.Empty(pks)

	delegator.deleteBuffer.RegisterL0(l0Global)
	pks, _ = delegator.GetLevel0Deletions(partitionID, pkoracle.NewCandidateKey(l0.ID(), l0.Partition(), segments.SegmentTypeGrowing))
	rawPks := make([]storage.PrimaryKey, 0, pks.Len())
	for i := 0; i < pks.Len(); i++ {
		rawPks = append(rawPks, pks.Get(i))
	}
	s.ElementsMatch(rawPks, []storage.PrimaryKey{partitionDeleteData.DeletePks().Get(0), allPartitionDeleteData.DeletePks().Get(0)})

	stubCandidate := pkoracle.NewBloomFilterSet(3, l0.Partition(), commonpb.SegmentState_Sealed)
	pks, _ = delegator.GetLevel0Deletions(partitionID, stubCandidate)
	rawPks = rawPks[:0]
	for i := 0; i < pks.Len(); i++ {
		rawPks = append(rawPks, pks.Get(i))
	}
	s.ElementsMatch(rawPks, []storage.PrimaryKey{partitionDeleteData.DeletePks().Get(0), allPartitionDeleteData.DeletePks().Get(0)})

	bfs := pkoracle.NewBloomFilterSet(3, l0.Partition(), commonpb.SegmentState_Sealed)
	bfs.UpdatePkCandidate([]storage.PrimaryKey{allPartitionDeleteData.DeletePks().Get(0)})

	pks, _ = delegator.GetLevel0Deletions(partitionID, bfs)
	// bf filtered segment
	s.Equal(pks.Len(), 1)
	s.True(pks.Get(0).EQ(allPartitionDeleteData.DeletePks().Get(0)))

	delegator.deleteBuffer.UnRegister(uint64(11))
	pks, _ = delegator.GetLevel0Deletions(partitionID, pkoracle.NewCandidateKey(l0.ID(), l0.Partition(), segments.SegmentTypeGrowing))
	s.True(pks.Get(0).EQ(allPartitionDeleteData.DeletePks().Get(0)))

	pks, _ = delegator.GetLevel0Deletions(partitionID+1, pkoracle.NewCandidateKey(l0.ID(), l0.Partition(), segments.SegmentTypeGrowing))
	s.True(pks.Get(0).EQ(allPartitionDeleteData.DeletePks().Get(0)))

	delegator.deleteBuffer.UnRegister(uint64(21))
	pks, _ = delegator.GetLevel0Deletions(partitionID+1, pkoracle.NewCandidateKey(l0.ID(), l0.Partition(), segments.SegmentTypeGrowing))
	s.Empty(pks)
}

// TestLoadSegmentsDoesNotBlockProcessDelete verifies that the 3-phase loadStreamDelete
// does not block ProcessDelete for extended periods. Before the fix (#47882),
// loadStreamDelete held deleteMut.RLock for the entire duration of BF checking (30-60s),
// which blocked ProcessDelete (needs WLock), freezing tsafe and causing query timeouts.
func (s *DelegatorDataSuite) TestLoadSegmentsDoesNotBlockProcessDelete() {
	defer func() {
		s.workerManager.ExpectedCalls = nil
		s.loader.ExpectedCalls = nil
	}()

	// Set up bloom filter that will match some deletes
	s.loader.EXPECT().LoadBloomFilterSet(mock.Anything, s.collectionID, mock.Anything).
		Call.Return(func(ctx context.Context, collectionID int64, infos ...*querypb.SegmentLoadInfo) []*pkoracle.BloomFilterSet {
		return lo.Map(infos, func(info *querypb.SegmentLoadInfo, _ int) *pkoracle.BloomFilterSet {
			bfs := pkoracle.NewBloomFilterSet(info.GetSegmentID(), info.GetPartitionID(), commonpb.SegmentState_Sealed)
			bf := bloomfilter.NewBloomFilterWithType(
				paramtable.Get().CommonCfg.BloomFilterSize.GetAsUint(),
				paramtable.Get().CommonCfg.MaxBloomFalsePositive.GetAsFloat(),
				paramtable.Get().CommonCfg.BloomFilterType.GetValue())
			pks := &storage.PkStatistics{PkFilter: bf}
			pks.UpdatePKRange(&storage.Int64FieldData{Data: []int64{10, 20, 30}})
			bfs.AddHistoricalStats(pks)
			return bfs
		})
	}, func(ctx context.Context, collectionID int64, infos ...*querypb.SegmentLoadInfo) error {
		return nil
	})

	workers := make(map[int64]*cluster.MockWorker)
	worker1 := &cluster.MockWorker{}
	workers[1] = worker1

	worker1.EXPECT().LoadSegments(mock.Anything, mock.AnythingOfType("*querypb.LoadSegmentsRequest")).
		Return(nil)
	worker1.EXPECT().Delete(mock.Anything, mock.AnythingOfType("*querypb.DeleteRequest")).Return(nil)
	s.workerManager.EXPECT().GetWorker(mock.Anything, mock.AnythingOfType("int64")).Call.Return(func(_ context.Context, nodeID int64) cluster.Worker {
		return workers[nodeID]
	}, nil)

	// Pre-populate delete buffer with entries so loadStreamDelete has work to do
	for i := 0; i < 100; i++ {
		s.delegator.ProcessDelete([]*DeleteData{
			{
				PartitionID: 500,
				PrimaryKeys: []storage.PrimaryKey{
					storage.NewInt64PrimaryKey(int64(i)),
				},
				Timestamps: []uint64{uint64(10 + i)},
				RowCount:   1,
			},
		}, uint64(10+i))
	}

	// Track ProcessDelete latency during concurrent LoadSegments
	var processDeleteBlocked atomic.Bool
	var wg sync.WaitGroup

	// Start LoadSegments in background
	wg.Add(1)
	go func() {
		defer wg.Done()
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		_ = s.delegator.LoadSegments(ctx, &querypb.LoadSegmentsRequest{
			Base:         commonpbutil.NewMsgBase(),
			DstNodeID:    1,
			CollectionID: s.collectionID,
			Infos: []*querypb.SegmentLoadInfo{
				{
					SegmentID:     300,
					PartitionID:   500,
					StartPosition: &msgpb.MsgPosition{Timestamp: 5},
					DeltaPosition: &msgpb.MsgPosition{Timestamp: 5},
					Level:         datapb.SegmentLevel_L1,
					InsertChannel: fmt.Sprintf("by-dev-rootcoord-dml_0_%dv0", s.collectionID),
				},
			},
		})
	}()

	// Give LoadSegments time to start Phase 2 (unlock period)
	time.Sleep(10 * time.Millisecond)

	// Attempt ProcessDelete — with the fix, this should complete quickly
	// because Phase 2 does NOT hold deleteMut
	wg.Add(1)
	go func() {
		defer wg.Done()
		start := time.Now()
		s.delegator.ProcessDelete([]*DeleteData{
			{
				PartitionID: 500,
				PrimaryKeys: []storage.PrimaryKey{
					storage.NewInt64PrimaryKey(999),
				},
				Timestamps: []uint64{200},
				RowCount:   1,
			},
		}, 200)
		// If ProcessDelete took more than 5 seconds, something is wrong
		// (before the fix, it would block for 30-60s)
		if time.Since(start) > 5*time.Second {
			processDeleteBlocked.Store(true)
		}
	}()

	wg.Wait()
	s.False(processDeleteBlocked.Load(), "ProcessDelete was blocked for too long during LoadSegments — tsafe would freeze")
}

func (s *DelegatorDataSuite) TestDelegatorData_ExcludeSegments() {
	s.delegator.AddExcludedSegments(map[int64]uint64{
		1: 3,
	})

	s.False(s.delegator.VerifyExcludedSegments(1, 1))
	s.True(s.delegator.VerifyExcludedSegments(1, 5))

	time.Sleep(time.Second * 1)
	s.delegator.TryCleanExcludedSegments(4)
	s.True(s.delegator.VerifyExcludedSegments(1, 1))
	s.True(s.delegator.VerifyExcludedSegments(1, 5))
}

func TestSegmentEffectiveTs(t *testing.T) {
	// Import segment: commit_timestamp takes precedence over start_position.Timestamp
	info := &querypb.SegmentLoadInfo{
		SegmentID:       1,
		StartPosition:   &msgpb.MsgPosition{Timestamp: 1000},
		CommitTimestamp: 5000,
	}
	assert.Equal(t, uint64(5000), segmentEffectiveTs(info))

	// Normal segment: falls back to start_position.Timestamp
	infoNormal := &querypb.SegmentLoadInfo{
		SegmentID:     2,
		StartPosition: &msgpb.MsgPosition{Timestamp: 1000},
	}
	assert.Equal(t, uint64(1000), segmentEffectiveTs(infoNormal))
}

func TestSegmentEffectiveTs_DelegatorHelper(t *testing.T) {
	// commit_ts=0: use start_position
	info := &querypb.SegmentLoadInfo{
		StartPosition: &msgpb.MsgPosition{Timestamp: 1000},
	}
	assert.Equal(t, uint64(1000), segmentEffectiveTs(info))

	// commit_ts > 0: override
	info.CommitTimestamp = 5000
	assert.Equal(t, uint64(5000), segmentEffectiveTs(info))

	// nil start_position + commit_ts
	info2 := &querypb.SegmentLoadInfo{CommitTimestamp: 3000}
	assert.Equal(t, uint64(3000), segmentEffectiveTs(info2))

	// both zero
	info3 := &querypb.SegmentLoadInfo{}
	assert.Equal(t, uint64(0), segmentEffectiveTs(info3))
}

func TestDeleteBuffer_PinsAtCommitTs(t *testing.T) {
	// Simulate: import segment segID=42, start_position.ts=1000, commit_ts=5000.
	// After LoadSegments, the delete buffer must be pinned at commit_ts=5000.

	const segID int64 = 42
	const commitTs uint64 = 5000
	const startTs uint64 = 1000

	buf := deletebuffer.NewDoubleCacheDeleteBuffer[*deletebuffer.Item](0, 1024*1024)

	// Compute effective ts: segmentEffectiveTs must return commit_ts for import segment.
	effectiveTs := segmentEffectiveTs(&querypb.SegmentLoadInfo{
		SegmentID:       segID,
		StartPosition:   &msgpb.MsgPosition{Timestamp: startTs},
		CommitTimestamp: commitTs,
	})
	assert.Equal(t, commitTs, effectiveTs,
		"segmentEffectiveTs must return commit_ts for import segment")

	// Pin at commit_ts (as LoadSegments does).
	buf.Pin(effectiveTs, segID)

	// Unpin at wrong ts (start_position.ts): no-op, pin at commit_ts stays.
	buf.Unpin(startTs, segID)

	// Pin at commit_ts is still active — Unpin at commit_ts removes it.
	buf.Unpin(commitTs, segID)
	// After unpin, a second unpin should be safe (no panic/error).
	buf.Unpin(commitTs, segID)
}

func TestDeleteBuffer_ListAfterUsesCommitTs(t *testing.T) {
	// Invariant I10: ListAfter must use segmentEffectiveTs (commit_ts when non-zero),
	// not the raw StartPosition timestamp.
	// If ListAfter uses the stale StartPosition, it would return too many delete records
	// (including pre-import deletes that are irrelevant).

	const commitTs uint64 = 5000
	const startTs uint64 = 1000

	info := &querypb.SegmentLoadInfo{
		SegmentID:       42,
		StartPosition:   &msgpb.MsgPosition{Timestamp: startTs},
		CommitTimestamp: commitTs,
	}

	effectiveTs := segmentEffectiveTs(info)
	assert.Equal(t, commitTs, effectiveTs,
		"segmentEffectiveTs must return commit_ts for import segment, not startPosition")

	// Verify: for normal segment (commitTs=0), effectiveTs falls back to startPosition
	normalInfo := &querypb.SegmentLoadInfo{
		SegmentID:     43,
		StartPosition: &msgpb.MsgPosition{Timestamp: startTs},
	}
	normalEffective := segmentEffectiveTs(normalInfo)
	assert.Equal(t, startTs, normalEffective,
		"segmentEffectiveTs must return startPosition for normal segment")
}

func TestDeleteBuffer_CatchUpUsesCommitTs(t *testing.T) {
	// Invariant I10: catchUpTs in loadStreamDelete must start at segmentEffectiveTs,
	// not StartPosition.Timestamp. This ensures the catch-up phase doesn't
	// re-process delete records from before the import commit time.

	const commitTs uint64 = 5000
	const startTs uint64 = 1000

	info := &querypb.SegmentLoadInfo{
		SegmentID:       42,
		StartPosition:   &msgpb.MsgPosition{Timestamp: startTs},
		CommitTimestamp: commitTs,
	}

	// Simulate catchUpTs initialization logic from loadStreamDelete
	catchUpTs := segmentEffectiveTs(info)
	// snapshotMaxTs=0 means no snapshot → catchUpTs should be segmentEffectiveTs
	snapshotMaxTs := uint64(0)
	if snapshotMaxTs > 0 {
		catchUpTs = snapshotMaxTs + 1
	}
	assert.Equal(t, commitTs, catchUpTs,
		"catchUpTs must start at commit_ts (not startPosition) when no snapshot exists")

	// With snapshot: catchUpTs should advance beyond snapshot
	snapshotMaxTs = 6000
	catchUpTs2 := segmentEffectiveTs(info)
	if snapshotMaxTs > 0 {
		catchUpTs2 = snapshotMaxTs + 1
	}
	assert.Equal(t, uint64(6001), catchUpTs2,
		"catchUpTs must advance to snapshotMaxTs+1 when snapshot exists")
}

func TestDelegatorDataSuite(t *testing.T) {
	suite.Run(t, new(DelegatorDataSuite))
}