Fix segment replacement bug

Signed-off-by: bigsheeper <yihao.dai@zilliz.com>

internal/querynode/collection_replica.go

@@ -513,7 +513,7 @@ func (colReplica *collectionReplicaImpl) replaceGrowingSegmentBySealedSegment(se
 		return errors.New("unexpected segment type")
 	}
 	targetSegment, err := colReplica.getSegmentByIDPrivate(segment.ID())
-	if err != nil && targetSegment != nil {
+	if err == nil && targetSegment != nil {
 		if targetSegment.segmentType != segTypeGrowing {
 			// target segment has been a sealed segment
 			return nil
@@ -521,7 +521,7 @@ func (colReplica *collectionReplicaImpl) replaceGrowingSegmentBySealedSegment(se
 		deleteSegment(targetSegment)
 	}
 
-	targetSegment = segment
+	colReplica.segments[segment.ID()] = segment
 	return nil
 }
 

internal/querynode/collection_replica_test.go

@@ -242,5 +242,36 @@ func TestCollectionReplica_freeAll(t *testing.T) {
 
 	err := node.Stop()
 	assert.NoError(t, err)
-
+}
+
+func TestReplaceGrowingSegmentBySealedSegment(t *testing.T) {
+	node := newQueryNodeMock()
+	collectionID := UniqueID(0)
+	segmentID := UniqueID(520)
+	initTestMeta(t, node, collectionID, segmentID)
+
+	_, _, segIDs := node.replica.getSealedSegmentsBySegmentType(segTypeGrowing)
+	assert.Equal(t, len(segIDs), 1)
+
+	collection, err := node.replica.getCollectionByID(collectionID)
+	assert.NoError(t, err)
+	ns := newSegment(collection, segmentID, defaultPartitionID, collectionID, segTypeSealed)
+	err = node.replica.replaceGrowingSegmentBySealedSegment(ns)
+	assert.NoError(t, err)
+
+	segmentNums := node.replica.getSegmentNum()
+	assert.Equal(t, segmentNums, 1)
+
+	segment, err := node.replica.getSegmentByID(segmentID)
+	assert.NoError(t, err)
+
+	assert.Equal(t, segment.getType(), segTypeSealed)
+
+	_, _, segIDs = node.replica.getSealedSegmentsBySegmentType(segTypeGrowing)
+	assert.Equal(t, len(segIDs), 0)
+	_, _, segIDs = node.replica.getSealedSegmentsBySegmentType(segTypeSealed)
+	assert.Equal(t, len(segIDs), 1)
+
+	err = node.Stop()
+	assert.NoError(t, err)
 }

internal/querynode/load_service.go

@@ -114,6 +114,11 @@ func (s *loadService) loadSegmentInternal(collectionID UniqueID, partitionID Uni
 	if err != nil {
 		return err
 	}
+	// replace segment
+	err = s.segLoader.replica.replaceGrowingSegmentBySealedSegment(segment)
+	if err != nil {
+		return err
+	}
 	if errIndex == nil {
 		fmt.Println("loading index...")
 		indexPaths, err := s.segLoader.indexLoader.getIndexPaths(buildID)
@@ -125,8 +130,7 @@ func (s *loadService) loadSegmentInternal(collectionID UniqueID, partitionID Uni
 			return err
 		}
 	}
-	// replace segment
-	return s.segLoader.replica.replaceGrowingSegmentBySealedSegment(segment)
+	return nil
 }
 
 func newLoadService(ctx context.Context, masterClient MasterServiceInterface, dataClient DataServiceInterface, indexClient IndexServiceInterface, replica collectionReplica, dmStream msgstream.MsgStream) *loadService {

internal/querynode/segment.go

@@ -12,6 +12,7 @@ package querynode
 */
 import "C"
 import (
+	"log"
 	"strconv"
 	"sync"
 	"unsafe"
@@ -23,13 +24,13 @@ import (
 )
 
 const (
-	segTypeInvalid  = C.Invalid
-	segTypeGrowing  = C.Growing
-	segTypeSealed   = C.Sealed
-	segTypeIndexing = C.Indexing
+	segTypeInvalid  = 0
+	segTypeGrowing  = 1
+	segTypeSealed   = 2
+	segTypeIndexing = 3
 )
 
-type segmentType = C.SegmentType
+type segmentType = int
 type indexParam = map[string]string
 
 type Segment struct {
@@ -45,7 +46,7 @@ type Segment struct {
 	recentlyModified bool
 
 	typeMu      sync.Mutex // guards builtIndex
-	segmentType C.SegmentType
+	segmentType int
 
 	paramMutex sync.RWMutex // guards index
 	indexParam map[int64]indexParam
@@ -112,7 +113,20 @@ func newSegment(collection *Collection, segmentID int64, partitionID UniqueID, c
 		NewSegment(CCollection collection, uint64_t segment_id, SegmentType seg_type);
 	*/
 	initIndexParam := make(map[int64]indexParam)
-	segmentPtr := C.NewSegment(collection.collectionPtr, C.ulong(segmentID), segType)
+	var segmentPtr C.CSegmentInterface
+	switch segType {
+	case segTypeInvalid:
+		log.Println("illegal segment type when create segment")
+		return nil
+	case segTypeSealed:
+		segmentPtr = C.NewSegment(collection.collectionPtr, C.ulong(segmentID), C.Sealed)
+	case segTypeGrowing:
+		segmentPtr = C.NewSegment(collection.collectionPtr, C.ulong(segmentID), C.Growing)
+	default:
+		log.Println("illegal segment type when create segment")
+		return nil
+	}
+
 	var newSegment = &Segment{
 		segmentPtr:   segmentPtr,
 		segmentType:  segType,