enhance: optimize ReplicaManager per-collection locking (#49389)

## Summary

- Replace ReplicaManager global locking with per-collection write locks.
- Add lock-free read indexes for replica ID and collection-scoped
replica lookups.
- Keep collection and flat replica indexes consistent across spawn,
move, recover, and removal paths.

issue: #48866

## Test Plan

- [x] Added ReplicaManager unit coverage for index consistency and
partial collection updates.
- [ ] Focused local go test is blocked by local core linker errors for
external segment symbols (`_loon_*`, `_FreeCFieldMemSizeList`,
`_SampleExternalSegmentFieldSizes`).

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Signed-off-by: sijie-ni-0214 <sijie.ni@zilliz.com>

internal/querycoordv2/job/job_update.go

@@ -142,7 +142,7 @@ func (job *UpdateLoadConfigJob) Execute() error {
 				replicaOldRG[replica.GetID()] = replica.GetResourceGroup()
 			}
 
-			if transferErr := job.meta.MoveReplica(job.ctx, rg, replicas); transferErr != nil {
+			if transferErr := job.meta.MoveReplica(job.ctx, collectionID, rg, replicas); transferErr != nil {
 				log.Warn("failed to transfer replica for collection", zap.Int64("collectionID", collectionID), zap.Error(transferErr))
 				err = transferErr
 				return err

internal/querycoordv2/meta/replica_manager.go

@@ -19,7 +19,7 @@ package meta
 import (
 	"context"
 	"fmt"
-	"sync"
+	"sort"
 	"time"
 
 	"github.com/cockroachdb/errors"
@@ -34,6 +34,7 @@ import (
 	"github.com/milvus-io/milvus/pkg/v3/proto/messagespb"
 	"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
 	"github.com/milvus-io/milvus/pkg/v3/util"
+	"github.com/milvus-io/milvus/pkg/v3/util/lock"
 	"github.com/milvus-io/milvus/pkg/v3/util/merr"
 	"github.com/milvus-io/milvus/pkg/v3/util/metricsinfo"
 	"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
@@ -50,7 +51,7 @@ type ReplicaManagerInterface interface {
 
 	// Replica manipulation
 	TransferReplica(ctx context.Context, collectionID typeutil.UniqueID, srcRGName string, dstRGName string, replicaNum int) error
-	MoveReplica(ctx context.Context, dstRGName string, toMove []*Replica) error
+	MoveReplica(ctx context.Context, collectionID typeutil.UniqueID, dstRGName string, toMove []*Replica) error
 	RemoveCollection(ctx context.Context, collectionID typeutil.UniqueID) error
 	RemoveReplicas(ctx context.Context, collectionID typeutil.UniqueID, replicas ...typeutil.UniqueID) error
 
@@ -62,8 +63,8 @@ type ReplicaManagerInterface interface {
 
 	// Node management
 	RecoverNodesInCollection(ctx context.Context, collectionID typeutil.UniqueID, rgs map[string]*ResourceGroup) error
-	RemoveNode(ctx context.Context, replicaID typeutil.UniqueID, nodes ...typeutil.UniqueID) error
-	RemoveSQNode(ctx context.Context, replicaID typeutil.UniqueID, nodes ...typeutil.UniqueID) error
+	RemoveNode(ctx context.Context, collectionID typeutil.UniqueID, replicaID typeutil.UniqueID, nodes ...typeutil.UniqueID) error
+	RemoveSQNode(ctx context.Context, collectionID typeutil.UniqueID, replicaID typeutil.UniqueID, nodes ...typeutil.UniqueID) error
 
 	// Metadata access
 	GetResourceGroupByCollection(ctx context.Context, collection typeutil.UniqueID) typeutil.Set[string]
@@ -74,46 +75,31 @@ type ReplicaManagerInterface interface {
 var _ ReplicaManagerInterface = (*ReplicaManager)(nil)
 
 type ReplicaManager struct {
-	rwmutex sync.RWMutex
+	// per-collection write lock, auto-created/recycled by KeyLock
+	collLock *lock.KeyLock[int64]
 
-	idAllocator            func() (int64, error)
-	replicas               map[typeutil.UniqueID]*Replica
-	coll2Replicas          map[typeutil.UniqueID]*collectionReplicas // typeutil.UniqueSet
-	queryInvisibleReplicas typeutil.UniqueSet
-	catalog                metastore.QueryCoordCatalog
-}
+	// flat index: replicaID -> *Replica, lock-free read for Get()
+	flatReplicas *typeutil.ConcurrentMap[int64, *Replica]
 
-// collectionReplicas maintains collection secondary index mapping
-type collectionReplicas struct {
-	id2replicas map[typeutil.UniqueID]*Replica
-	replicas    []*Replica
-}
+	// per-collection replica list, lock-free read via ConcurrentMap
+	coll2Replicas *typeutil.ConcurrentMap[int64, []*Replica]
 
-func (crs *collectionReplicas) removeReplicas(replicaIDs ...int64) (empty bool) {
-	for _, replicaID := range replicaIDs {
-		delete(crs.id2replicas, replicaID)
-	}
-	crs.replicas = lo.Values(crs.id2replicas)
-	return len(crs.replicas) == 0
-}
+	// queryInvisibleReplicas is a secondary index over Replica.queryInvisible.
+	// Replica.queryInvisible remains the source of truth; this set only avoids
+	// scanning all replicas in the load-config promotion loop.
+	queryInvisibleReplicas *typeutil.ConcurrentSet[int64]
 
-func (crs *collectionReplicas) putReplica(replica *Replica) {
-	crs.id2replicas[replica.GetID()] = replica
-	crs.replicas = lo.Values(crs.id2replicas)
-}
-
-func newCollectionReplicas() *collectionReplicas {
-	return &collectionReplicas{
-		id2replicas: make(map[typeutil.UniqueID]*Replica),
-	}
+	idAllocator func() (int64, error)
+	catalog     metastore.QueryCoordCatalog
 }
 
 func NewReplicaManager(idAllocator func() (int64, error), catalog metastore.QueryCoordCatalog) *ReplicaManager {
 	return &ReplicaManager{
+		collLock:               lock.NewKeyLock[int64](),
+		flatReplicas:           typeutil.NewConcurrentMap[int64, *Replica](),
+		coll2Replicas:          typeutil.NewConcurrentMap[int64, []*Replica](),
+		queryInvisibleReplicas: typeutil.NewConcurrentSet[int64](),
 		idAllocator:            idAllocator,
-		replicas:               make(map[int64]*Replica),
-		coll2Replicas:          make(map[int64]*collectionReplicas),
-		queryInvisibleReplicas: typeutil.NewUniqueSet(),
 		catalog:                catalog,
 	}
 }
@@ -126,6 +112,7 @@ func (m *ReplicaManager) Recover(ctx context.Context, collections []int64) error
 	}
 
 	collectionSet := typeutil.NewUniqueSet(collections...)
+	grouped := make(map[int64][]*Replica)
 	for _, replica := range replicas {
 		if len(replica.GetResourceGroup()) == 0 {
 			replica.ResourceGroup = DefaultResourceGroupName
@@ -133,7 +120,7 @@ func (m *ReplicaManager) Recover(ctx context.Context, collections []int64) error
 
 		if collectionSet.Contain(replica.GetCollectionID()) {
 			rep := NewReplicaWithPriority(replica, commonpb.LoadPriority_HIGH)
-			m.putReplicaInMemory(rep)
+			grouped[rep.GetCollectionID()] = append(grouped[rep.GetCollectionID()], rep)
 			log.Info("recover replica",
 				zap.Int64("collectionID", replica.GetCollectionID()),
 				zap.Int64("replicaID", replica.GetID()),
@@ -154,48 +141,66 @@ func (m *ReplicaManager) Recover(ctx context.Context, collections []int64) error
 			)
 		}
 	}
+	for collID, reps := range grouped {
+		m.collLock.Lock(collID)
+		m.putReplicasInMemory(collID, reps...)
+		m.collLock.Unlock(collID)
+	}
 	return nil
 }
 
 // Get returns the replica by id.
 // Replica should be read-only, do not modify it.
+// Uses lock-free flat index for fast lookup; safe because Replica is COW-immutable.
 func (m *ReplicaManager) Get(ctx context.Context, id typeutil.UniqueID) *Replica {
-	m.rwmutex.RLock()
-	defer m.rwmutex.RUnlock()
-
-	return m.replicas[id]
+	replica, _ := m.flatReplicas.Get(id)
+	return replica
 }
 
 func (m *ReplicaManager) GetQueryInvisibleReplicas(ctx context.Context) []*Replica {
-	m.rwmutex.RLock()
-	defer m.rwmutex.RUnlock()
-
-	replicas := make([]*Replica, 0, m.queryInvisibleReplicas.Len())
-	for replicaID := range m.queryInvisibleReplicas {
-		if replica := m.replicas[replicaID]; replica != nil {
-			replicas = append(replicas, replica)
+	replicaIDs := m.queryInvisibleReplicas.Collect()
+	replicas := make([]*Replica, 0, len(replicaIDs))
+	for _, replicaID := range replicaIDs {
+		replica, ok := m.flatReplicas.Get(replicaID)
+		if !ok || replica.IsQueryVisible() {
+			continue
 		}
+		replicas = append(replicas, replica)
 	}
 	return replicas
 }
 
 func (m *ReplicaManager) SetReplicasQueryVisible(ctx context.Context, replicaIDs ...typeutil.UniqueID) []typeutil.UniqueID {
-	m.rwmutex.Lock()
-	defer m.rwmutex.Unlock()
-
-	collections := typeutil.NewUniqueSet()
-	modifiedReplicas := make([]*Replica, 0, len(replicaIDs))
+	grouped := make(map[typeutil.UniqueID][]typeutil.UniqueID)
 	for _, replicaID := range replicaIDs {
-		replica := m.replicas[replicaID]
-		if replica == nil || replica.IsQueryVisible() {
+		replica, ok := m.flatReplicas.Get(replicaID)
+		if !ok || replica.IsQueryVisible() {
 			continue
 		}
-		mutableReplica := replica.CopyForWrite()
-		mutableReplica.SetQueryInvisible(false)
-		modifiedReplicas = append(modifiedReplicas, mutableReplica.IntoReplica())
-		collections.Insert(replica.GetCollectionID())
+		grouped[replica.GetCollectionID()] = append(grouped[replica.GetCollectionID()], replicaID)
+	}
+
+	collections := typeutil.NewUniqueSet()
+	for collectionID, ids := range grouped {
+		m.collLock.Lock(collectionID)
+
+		modifiedReplicas := make([]*Replica, 0, len(ids))
+		for _, replicaID := range ids {
+			replica, ok := m.flatReplicas.Get(replicaID)
+			if !ok || replica.GetCollectionID() != collectionID || replica.IsQueryVisible() {
+				continue
+			}
+			mutableReplica := replica.CopyForWrite()
+			mutableReplica.SetQueryInvisible(false)
+			modifiedReplicas = append(modifiedReplicas, mutableReplica.IntoReplica())
+		}
+		if len(modifiedReplicas) > 0 {
+			m.putReplicasInMemory(collectionID, modifiedReplicas...)
+			collections.Insert(collectionID)
+		}
+
+		m.collLock.Unlock(collectionID)
 	}
-	m.putReplicaInMemory(modifiedReplicas...)
 	return collections.Collect()
 }
 
@@ -207,14 +212,15 @@ type SpawnWithReplicaConfigParams struct {
 
 // SpawnWithReplicaConfig spawns replicas with replica config.
 func (m *ReplicaManager) SpawnWithReplicaConfig(ctx context.Context, params SpawnWithReplicaConfigParams) ([]*Replica, error) {
-	m.rwmutex.Lock()
-	defer m.rwmutex.Unlock()
+	m.collLock.Lock(params.CollectionID)
+	defer m.collLock.Unlock(params.CollectionID)
 
 	balancePolicy := paramtable.Get().QueryCoordCfg.Balancer.GetValue()
 	enableChannelExclusiveMode := balancePolicy == ChannelLevelScoreBalancerName
 	replicas := make([]*Replica, 0)
 	for _, config := range params.Configs {
-		if existedReplica, ok := m.replicas[config.GetReplicaId()]; ok {
+		if existedReplica, ok := m.flatReplicas.Get(config.GetReplicaId()); ok &&
+			existedReplica.GetCollectionID() == params.CollectionID {
 			// if the replica is already existed, just update the resource group
 			mutableReplica := existedReplica.CopyForWrite()
 			mutableReplica.SetResourceGroup(config.GetResourceGroupName())
@@ -238,7 +244,7 @@ func (m *ReplicaManager) SpawnWithReplicaConfig(ctx context.Context, params Spaw
 			zap.String("resourceGroup", config.GetResourceGroupName()),
 		)
 	}
-	if err := m.put(ctx, replicas...); err != nil {
+	if err := m.put(ctx, params.CollectionID, replicas...); err != nil {
 		return nil, errors.Wrap(err, "failed to put replicas")
 	}
 	if err := m.removeRedundantReplicas(ctx, params); err != nil {
@@ -249,12 +255,12 @@ func (m *ReplicaManager) SpawnWithReplicaConfig(ctx context.Context, params Spaw
 
 // removeRedundantReplicas removes redundant replicas that is not in the new replica config.
 func (m *ReplicaManager) removeRedundantReplicas(ctx context.Context, params SpawnWithReplicaConfigParams) error {
-	existedReplicas, ok := m.coll2Replicas[params.CollectionID]
+	existedReplicas, ok := m.coll2Replicas.Get(params.CollectionID)
 	if !ok {
 		return nil
 	}
 	toRemoveReplicas := make([]int64, 0)
-	for _, replica := range existedReplicas.replicas {
+	for _, replica := range existedReplicas {
 		found := false
 		replicaID := replica.GetID()
 		for _, channel := range params.Configs {
@@ -309,8 +315,8 @@ func (m *ReplicaManager) Spawn(ctx context.Context, collection int64, replicaNum
 		opt(cfg)
 	}
 
-	m.rwmutex.Lock()
-	defer m.rwmutex.Unlock()
+	m.collLock.Lock(collection)
+	defer m.collLock.Unlock(collection)
 
 	balancePolicy := paramtable.Get().QueryCoordCfg.Balancer.GetValue()
 	enableChannelExclusiveMode := balancePolicy == ChannelLevelScoreBalancerName
@@ -342,7 +348,7 @@ func (m *ReplicaManager) Spawn(ctx context.Context, collection int64, replicaNum
 			replicas = append(replicas, replica)
 		}
 	}
-	if err := m.put(ctx, replicas...); err != nil {
+	if err := m.put(ctx, collection, replicas...); err != nil {
 		return nil, err
 	}
 	return replicas, nil
@@ -351,13 +357,44 @@ func (m *ReplicaManager) Spawn(ctx context.Context, collection int64, replicaNum
 // Deprecated: Warning, break the consistency of ReplicaManager,
 // never use it in non-test code, use Spawn instead.
 func (m *ReplicaManager) Put(ctx context.Context, replicas ...*Replica) error {
-	m.rwmutex.Lock()
-	defer m.rwmutex.Unlock()
+	if len(replicas) == 0 {
+		return nil
+	}
+	grouped := make(map[int64][]*Replica)
+	collectionIDs := make([]int64, 0)
+	for _, r := range replicas {
+		if _, ok := grouped[r.GetCollectionID()]; !ok {
+			collectionIDs = append(collectionIDs, r.GetCollectionID())
+		}
+		grouped[r.GetCollectionID()] = append(grouped[r.GetCollectionID()], r)
+	}
+	sort.Slice(collectionIDs, func(i, j int) bool {
+		return collectionIDs[i] < collectionIDs[j]
+	})
 
-	return m.put(ctx, replicas...)
+	for _, collID := range collectionIDs {
+		m.collLock.Lock(collID)
+	}
+	defer func() {
+		for i := len(collectionIDs) - 1; i >= 0; i-- {
+			m.collLock.Unlock(collectionIDs[i])
+		}
+	}()
+
+	replicaPBs := make([]*querypb.Replica, 0, len(replicas))
+	for _, replica := range replicas {
+		replicaPBs = append(replicaPBs, replica.replicaPB)
+	}
+	if err := m.catalog.SaveReplica(ctx, replicaPBs...); err != nil {
+		return err
+	}
+	for _, collID := range collectionIDs {
+		m.putReplicasInMemory(collID, grouped[collID]...)
+	}
+	return nil
 }
 
-func (m *ReplicaManager) put(ctx context.Context, replicas ...*Replica) error {
+func (m *ReplicaManager) put(ctx context.Context, collectionID typeutil.UniqueID, replicas ...*Replica) error {
 	if len(replicas) == 0 {
 		return nil
 	}
@@ -370,35 +407,47 @@ func (m *ReplicaManager) put(ctx context.Context, replicas ...*Replica) error {
 		return err
 	}
 
-	m.putReplicaInMemory(replicas...)
+	m.putReplicasInMemory(collectionID, replicas...)
 	return nil
 }
 
-// putReplicaInMemory puts replicas into in-memory map and collIDToReplicaIDs.
-func (m *ReplicaManager) putReplicaInMemory(replicas ...*Replica) {
-	if m.queryInvisibleReplicas == nil {
-		m.queryInvisibleReplicas = typeutil.NewUniqueSet()
+// putReplicasInMemory puts replicas into in-memory indexes.
+// Caller must hold collLock for the collection.
+func (m *ReplicaManager) putReplicasInMemory(collectionID typeutil.UniqueID, replicas ...*Replica) {
+	if len(replicas) == 0 {
+		return
 	}
-	for _, replica := range replicas {
-		if oldReplica, ok := m.replicas[replica.GetID()]; ok {
+
+	old, _ := m.coll2Replicas.Get(collectionID)
+	newSlice := append([]*Replica(nil), old...)
+
+	indexes := make(map[int64]int, len(newSlice))
+	for i, r := range newSlice {
+		indexes[r.GetID()] = i
+	}
+
+	for _, r := range replicas {
+		if i, ok := indexes[r.GetID()]; ok {
+			newSlice[i] = r
+		} else {
+			newSlice = append(newSlice, r)
+		}
+	}
+
+	for _, r := range replicas {
+		if oldReplica, ok := m.flatReplicas.Get(r.GetID()); ok {
 			metrics.QueryCoordResourceGroupReplicaTotal.WithLabelValues(oldReplica.GetResourceGroup()).Dec()
 			metrics.QueryCoordReplicaRONodeTotal.Add(-float64(oldReplica.RONodesCount()))
 			m.queryInvisibleReplicas.Remove(oldReplica.GetID())
 		}
-		// update in-memory replicas.
-		m.replicas[replica.GetID()] = replica
-		if !replica.IsQueryVisible() {
-			m.queryInvisibleReplicas.Insert(replica.GetID())
+		m.flatReplicas.Insert(r.GetID(), r)
+		if !r.IsQueryVisible() {
+			m.queryInvisibleReplicas.Insert(r.GetID())
 		}
-		metrics.QueryCoordResourceGroupReplicaTotal.WithLabelValues(replica.GetResourceGroup()).Inc()
-		metrics.QueryCoordReplicaRONodeTotal.Add(float64(replica.RONodesCount()))
-
-		// update collIDToReplicaIDs.
-		if m.coll2Replicas[replica.GetCollectionID()] == nil {
-			m.coll2Replicas[replica.GetCollectionID()] = newCollectionReplicas()
-		}
-		m.coll2Replicas[replica.GetCollectionID()].putReplica(replica)
+		metrics.QueryCoordResourceGroupReplicaTotal.WithLabelValues(r.GetResourceGroup()).Inc()
+		metrics.QueryCoordReplicaRONodeTotal.Add(float64(r.RONodesCount()))
 	}
+	m.coll2Replicas.Insert(collectionID, newSlice)
 }
 
 // TransferReplica transfers N replicas from srcRGName to dstRGName.
@@ -410,8 +459,15 @@ func (m *ReplicaManager) TransferReplica(ctx context.Context, collectionID typeu
 		return merr.WrapErrParameterInvalid("NumReplica > 0", fmt.Sprintf("invalid NumReplica %d", replicaNum))
 	}
 
-	m.rwmutex.Lock()
-	defer m.rwmutex.Unlock()
+	m.collLock.Lock(collectionID)
+	defer m.collLock.Unlock(collectionID)
+
+	if _, ok := m.coll2Replicas.Get(collectionID); !ok {
+		return merr.WrapErrParameterInvalid(
+			"Collection not loaded",
+			fmt.Sprintf("collectionID %d", collectionID),
+		)
+	}
 
 	// Check if replica can be transfer.
 	srcReplicas, err := m.getSrcReplicasAndCheckIfTransferable(collectionID, srcRGName, replicaNum)
@@ -427,14 +483,24 @@ func (m *ReplicaManager) TransferReplica(ctx context.Context, collectionID typeu
 		mutableReplica.SetResourceGroup(dstRGName)
 		replicas = append(replicas, mutableReplica.IntoReplica())
 	}
-	return m.put(ctx, replicas...)
+
+	return m.put(ctx, collectionID, replicas...)
 }
 
-func (m *ReplicaManager) MoveReplica(ctx context.Context, dstRGName string, toMove []*Replica) error {
-	m.rwmutex.Lock()
-	defer m.rwmutex.Unlock()
+func (m *ReplicaManager) MoveReplica(ctx context.Context, collectionID typeutil.UniqueID, dstRGName string, toMove []*Replica) error {
+	if len(toMove) == 0 {
+		return nil
+	}
+
+	m.collLock.Lock(collectionID)
+	defer m.collLock.Unlock(collectionID)
+
+	if _, ok := m.coll2Replicas.Get(collectionID); !ok {
+		return nil
+	}
+
 	replicas := make([]*Replica, 0, len(toMove))
-	replicaIDs := make([]int64, 0)
+	replicaIDs := make([]int64, 0, len(toMove))
 	for _, replica := range toMove {
 		mutableReplica := replica.CopyForWrite()
 		mutableReplica.SetResourceGroup(dstRGName)
@@ -442,21 +508,16 @@ func (m *ReplicaManager) MoveReplica(ctx context.Context, dstRGName string, toMo
 		replicaIDs = append(replicaIDs, replica.GetID())
 	}
 	log.Info("move replicas to resource group", zap.String("dstRGName", dstRGName), zap.Int64s("replicas", replicaIDs))
-	return m.put(ctx, replicas...)
+	return m.put(ctx, collectionID, replicas...)
 }
 
-// getSrcReplicasAndCheckIfTransferable checks if the collection can be transfer from srcRGName to dstRGName.
+// getSrcReplicasAndCheckIfTransferable checks if the collection can be transferred.
+// Caller must hold collLock for this collection.
 func (m *ReplicaManager) getSrcReplicasAndCheckIfTransferable(collectionID typeutil.UniqueID, srcRGName string, replicaNum int) ([]*Replica, error) {
-	// Check if collection is loaded.
-	if m.coll2Replicas[collectionID] == nil {
-		return nil, merr.WrapErrParameterInvalid(
-			"Collection not loaded",
-			fmt.Sprintf("collectionID %d", collectionID),
-		)
-	}
-
-	// Check if replica in srcRGName is enough.
-	srcReplicas := m.getByCollectionAndRG(collectionID, srcRGName)
+	replicas, _ := m.coll2Replicas.Get(collectionID)
+	srcReplicas := lo.Filter(replicas, func(replica *Replica, _ int) bool {
+		return replica.GetResourceGroup() == srcRGName
+	})
 	if len(srcReplicas) < replicaNum {
 		err := merr.WrapErrParameterInvalid(
 			"NumReplica not greater than the number of replica in source resource group", fmt.Sprintf("only found [%d] replicas of collection [%d] in source resource group [%s], but %d require",
@@ -472,124 +533,129 @@ func (m *ReplicaManager) getSrcReplicasAndCheckIfTransferable(collectionID typeu
 // RemoveCollection removes replicas of given collection,
 // returns error if failed to remove replica from KV
 func (m *ReplicaManager) RemoveCollection(ctx context.Context, collectionID typeutil.UniqueID) error {
-	m.rwmutex.Lock()
-	defer m.rwmutex.Unlock()
+	m.collLock.Lock(collectionID)
+	defer m.collLock.Unlock(collectionID)
 
 	err := m.catalog.ReleaseReplicas(ctx, collectionID)
 	if err != nil {
 		return err
 	}
 
-	if collReplicas, ok := m.coll2Replicas[collectionID]; ok {
-		// Remove all replica of collection and remove collection from collIDToReplicaIDs.
-		for _, replica := range collReplicas.replicas {
+	if replicas, ok := m.coll2Replicas.Get(collectionID); ok {
+		// Remove all replica of collection and remove collection from coll2Replicas.
+		// coll2Replicas is updated before flatReplicas so the invariant
+		// "visible via GetByCollection => visible via Get" holds during deletion.
+		m.coll2Replicas.Remove(collectionID)
+		for _, replica := range replicas {
 			metrics.QueryCoordResourceGroupReplicaTotal.WithLabelValues(replica.GetResourceGroup()).Dec()
 			metrics.QueryCoordReplicaRONodeTotal.Add(-float64(replica.RONodesCount()))
-			delete(m.replicas, replica.GetID())
+			m.flatReplicas.Remove(replica.GetID())
 			m.queryInvisibleReplicas.Remove(replica.GetID())
 		}
-		delete(m.coll2Replicas, collectionID)
 	}
 	return nil
 }
 
-func (m *ReplicaManager) RemoveReplicas(ctx context.Context, collectionID typeutil.UniqueID, replicas ...typeutil.UniqueID) error {
-	m.rwmutex.Lock()
-	defer m.rwmutex.Unlock()
+func (m *ReplicaManager) RemoveReplicas(ctx context.Context, collectionID typeutil.UniqueID, replicaIDs ...typeutil.UniqueID) error {
+	m.collLock.Lock(collectionID)
+	defer m.collLock.Unlock(collectionID)
 
-	log.Info("release replicas", zap.Int64("collectionID", collectionID), zap.Int64s("replicas", replicas))
-
-	return m.removeReplicas(ctx, collectionID, replicas...)
-}
-
-func (m *ReplicaManager) removeReplicas(ctx context.Context, collectionID typeutil.UniqueID, replicas ...typeutil.UniqueID) error {
-	err := m.catalog.ReleaseReplica(ctx, collectionID, replicas...)
-	if err != nil {
-		return err
-	}
-
-	for _, replicaID := range replicas {
-		if replica, ok := m.replicas[replicaID]; ok {
-			metrics.QueryCoordResourceGroupReplicaTotal.WithLabelValues(replica.GetResourceGroup()).Dec()
-			metrics.QueryCoordReplicaRONodeTotal.Add(float64(-replica.RONodesCount()))
-			delete(m.replicas, replicaID)
-			m.queryInvisibleReplicas.Remove(replicaID)
-		}
-	}
-
-	if m.coll2Replicas[collectionID].removeReplicas(replicas...) {
-		delete(m.coll2Replicas, collectionID)
-	}
-
-	return nil
-}
-
-func (m *ReplicaManager) GetByCollection(ctx context.Context, collectionID typeutil.UniqueID) []*Replica {
-	m.rwmutex.RLock()
-	defer m.rwmutex.RUnlock()
-	return m.getByCollection(collectionID)
-}
-
-func (m *ReplicaManager) getByCollection(collectionID typeutil.UniqueID) []*Replica {
-	collReplicas, ok := m.coll2Replicas[collectionID]
-	if !ok {
+	if _, ok := m.coll2Replicas.Get(collectionID); !ok {
 		return nil
 	}
 
-	return collReplicas.replicas
+	log.Info("release replicas", zap.Int64("collectionID", collectionID), zap.Int64s("replicas", replicaIDs))
+	return m.removeReplicas(ctx, collectionID, replicaIDs...)
 }
 
-func (m *ReplicaManager) GetByCollectionAndNode(ctx context.Context, collectionID, nodeID typeutil.UniqueID) *Replica {
-	m.rwmutex.RLock()
-	defer m.rwmutex.RUnlock()
+// removeReplicas removes specific replicas while holding collLock.Lock.
+// coll2Replicas is updated before flatReplicas so lock-free readers never observe
+// a replica that is visible via GetByCollection but missing from Get(id).
+func (m *ReplicaManager) removeReplicas(ctx context.Context, collectionID int64, replicaIDs ...int64) error {
+	if len(replicaIDs) == 0 {
+		return nil
+	}
+	if err := m.catalog.ReleaseReplica(ctx, collectionID, replicaIDs...); err != nil {
+		return err
+	}
+	for _, replicaID := range replicaIDs {
+		if replica, ok := m.flatReplicas.Get(replicaID); ok {
+			metrics.QueryCoordResourceGroupReplicaTotal.WithLabelValues(replica.GetResourceGroup()).Dec()
+			metrics.QueryCoordReplicaRONodeTotal.Add(float64(-replica.RONodesCount()))
+		}
+	}
+	m.removeReplicasInMemory(collectionID, replicaIDs...)
+	return nil
+}
 
-	if m.coll2Replicas[collectionID] != nil {
-		for _, replica := range m.coll2Replicas[collectionID].replicas {
-			if replica.Contains(nodeID) {
-				return replica
+// removeReplicasInMemory removes replicas from both collection and flat indexes.
+// coll2Replicas is updated before flatReplicas so lock-free readers never observe
+// a replica that is visible via GetByCollection but missing from Get(id).
+// Caller must hold collLock for the collection.
+func (m *ReplicaManager) removeReplicasInMemory(collectionID typeutil.UniqueID, replicaIDs ...int64) {
+	old, ok := m.coll2Replicas.Get(collectionID)
+	if ok {
+		removeSet := typeutil.NewSet(replicaIDs...)
+		newSlice := make([]*Replica, 0, len(old))
+		for _, r := range old {
+			if !removeSet.Contain(r.GetID()) {
+				newSlice = append(newSlice, r)
 			}
 		}
+		if len(newSlice) == 0 {
+			m.coll2Replicas.Remove(collectionID)
+		} else {
+			m.coll2Replicas.Insert(collectionID, newSlice)
+		}
 	}
 
+	for _, replicaID := range replicaIDs {
+		m.flatReplicas.Remove(replicaID)
+		m.queryInvisibleReplicas.Remove(replicaID)
+	}
+}
+
+func (m *ReplicaManager) GetByCollection(ctx context.Context, collectionID typeutil.UniqueID) []*Replica {
+	replicas, _ := m.coll2Replicas.Get(collectionID)
+	return replicas
+}
+
+func (m *ReplicaManager) GetByCollectionAndNode(ctx context.Context, collectionID, nodeID typeutil.UniqueID) *Replica {
+	replicas, ok := m.coll2Replicas.Get(collectionID)
+	if !ok {
+		return nil
+	}
+	for _, replica := range replicas {
+		if replica.Contains(nodeID) {
+			return replica
+		}
+	}
 	return nil
 }
 
 func (m *ReplicaManager) GetByNode(ctx context.Context, nodeID typeutil.UniqueID) []*Replica {
-	m.rwmutex.RLock()
-	defer m.rwmutex.RUnlock()
-
 	replicas := make([]*Replica, 0)
-	for _, replica := range m.replicas {
-		if replica.Contains(nodeID) {
-			replicas = append(replicas, replica)
+	m.coll2Replicas.Range(func(_ int64, collReplicas []*Replica) bool {
+		for _, replica := range collReplicas {
+			if replica.Contains(nodeID) {
+				replicas = append(replicas, replica)
+			}
 		}
-	}
-
+		return true
+	})
 	return replicas
 }
 
-func (m *ReplicaManager) getByCollectionAndRG(collectionID int64, rgName string) []*Replica {
-	collReplicas, ok := m.coll2Replicas[collectionID]
-	if !ok {
-		return nil
-	}
-
-	return lo.Filter(collReplicas.replicas, func(replica *Replica, _ int) bool {
-		return replica.GetResourceGroup() == rgName
-	})
-}
-
 func (m *ReplicaManager) GetByResourceGroup(ctx context.Context, rgName string) []*Replica {
-	m.rwmutex.RLock()
-	defer m.rwmutex.RUnlock()
-
 	ret := make([]*Replica, 0)
-	for _, replica := range m.replicas {
-		if replica.GetResourceGroup() == rgName {
-			ret = append(ret, replica)
+	m.coll2Replicas.Range(func(_ int64, collReplicas []*Replica) bool {
+		for _, replica := range collReplicas {
+			if replica.GetResourceGroup() == rgName {
+				ret = append(ret, replica)
+			}
 		}
-	}
-
+		return true
+	})
 	return ret
 }
 
@@ -613,8 +679,12 @@ func (m *ReplicaManager) RecoverNodesInCollection(ctx context.Context, collectio
 		return err
 	}
 
-	m.rwmutex.Lock()
-	defer m.rwmutex.Unlock()
+	m.collLock.Lock(collectionID)
+	defer m.collLock.Unlock(collectionID)
+
+	if _, ok := m.coll2Replicas.Get(collectionID); !ok {
+		return errors.Errorf("collection %d not loaded", collectionID)
+	}
 
 	// create a helper to do the recover.
 	helper, err := m.getCollectionAssignmentHelper(collectionID, rgNodeSets)
@@ -626,7 +696,7 @@ func (m *ReplicaManager) RecoverNodesInCollection(ctx context.Context, collectio
 	// recover node by resource group.
 	helper.RangeOverResourceGroup(func(replicaHelper *replicasInSameRGAssignmentHelper) {
 		replicaHelper.RangeOverReplicas(func(assignment *replicaAssignmentInfo) {
-			replica := m.replicas[assignment.GetReplicaID()]
+			replica := assignment.GetReplica()
 			// For replicas with needWaitRGReady flag, skip assignment if the RG still has missing nodes.
 			if replica.NeedWaitRGReady() {
 				rgName := replica.GetResourceGroup()
@@ -676,7 +746,11 @@ func (m *ReplicaManager) RecoverNodesInCollection(ctx context.Context, collectio
 			modifiedReplicas = append(modifiedReplicas, mutableReplica.IntoReplica())
 		})
 	})
-	return m.put(ctx, modifiedReplicas...)
+
+	if len(modifiedReplicas) == 0 {
+		return nil
+	}
+	return m.put(ctx, collectionID, modifiedReplicas...)
 }
 
 // validateResourceGroups checks if the resource groups are valid.
@@ -694,56 +768,54 @@ func (m *ReplicaManager) validateResourceGroups(rgs map[string]typeutil.UniqueSe
 	return nil
 }
 
-// getCollectionAssignmentHelper checks if the collection is recoverable and group replicas by resource group.
+// getCollectionAssignmentHelper builds an assignment helper from collection replicas.
+// Caller must hold collLock for this collection.
 func (m *ReplicaManager) getCollectionAssignmentHelper(collectionID typeutil.UniqueID, rgs map[string]typeutil.UniqueSet) (*collectionAssignmentHelper, error) {
 	// check if the collection is exist.
-	collReplicas, ok := m.coll2Replicas[collectionID]
+	replicas, ok := m.coll2Replicas.Get(collectionID)
 	if !ok {
 		return nil, errors.Errorf("collection %d not loaded", collectionID)
 	}
 
 	rgToReplicas := make(map[string][]*Replica)
-	for _, replica := range collReplicas.replicas {
+	for _, replica := range replicas {
 		rgName := replica.GetResourceGroup()
 		if _, ok := rgs[rgName]; !ok {
 			return nil, errors.Errorf("lost resource group info, collectionID: %d, replicaID: %d, resourceGroup: %s", collectionID, replica.GetID(), rgName)
 		}
-		if _, ok := rgToReplicas[rgName]; !ok {
-			rgToReplicas[rgName] = make([]*Replica, 0)
-		}
 		rgToReplicas[rgName] = append(rgToReplicas[rgName], replica)
 	}
 	return newCollectionAssignmentHelper(collectionID, rgToReplicas, rgs), nil
 }
 
-// RemoveNode removes the node from all replicas of given collection.
-func (m *ReplicaManager) RemoveNode(ctx context.Context, replicaID typeutil.UniqueID, nodes ...typeutil.UniqueID) error {
-	m.rwmutex.Lock()
-	defer m.rwmutex.Unlock()
+// RemoveNode removes the node from the given replica.
+func (m *ReplicaManager) RemoveNode(ctx context.Context, collectionID typeutil.UniqueID, replicaID typeutil.UniqueID, nodes ...typeutil.UniqueID) error {
+	m.collLock.Lock(collectionID)
+	defer m.collLock.Unlock(collectionID)
 
-	replica, ok := m.replicas[replicaID]
-	if !ok {
+	replica, ok := m.flatReplicas.Get(replicaID)
+	if !ok || replica.GetCollectionID() != collectionID {
 		return merr.WrapErrReplicaNotFound(replicaID)
 	}
 
 	mutableReplica := replica.CopyForWrite()
 	mutableReplica.RemoveNode(nodes...) // ro -> unused
-	return m.put(ctx, mutableReplica.IntoReplica())
+	return m.put(ctx, collectionID, mutableReplica.IntoReplica())
 }
 
-// RemoveSQNode removes the sq node from all replicas of given collection.
-func (m *ReplicaManager) RemoveSQNode(ctx context.Context, replicaID typeutil.UniqueID, nodes ...typeutil.UniqueID) error {
-	m.rwmutex.Lock()
-	defer m.rwmutex.Unlock()
+// RemoveSQNode removes the sq node from the given replica.
+func (m *ReplicaManager) RemoveSQNode(ctx context.Context, collectionID typeutil.UniqueID, replicaID typeutil.UniqueID, nodes ...typeutil.UniqueID) error {
+	m.collLock.Lock(collectionID)
+	defer m.collLock.Unlock(collectionID)
 
-	replica, ok := m.replicas[replicaID]
-	if !ok {
+	replica, ok := m.flatReplicas.Get(replicaID)
+	if !ok || replica.GetCollectionID() != collectionID {
 		return merr.WrapErrReplicaNotFound(replicaID)
 	}
 
 	mutableReplica := replica.CopyForWrite()
 	mutableReplica.RemoveSQNode(nodes...) // ro -> unused
-	return m.put(ctx, mutableReplica.IntoReplica())
+	return m.put(ctx, collectionID, mutableReplica.IntoReplica())
 }
 
 func (m *ReplicaManager) GetResourceGroupByCollection(ctx context.Context, collection typeutil.UniqueID) typeutil.Set[string] {
@@ -753,38 +825,36 @@ func (m *ReplicaManager) GetResourceGroupByCollection(ctx context.Context, colle
 }
 
 // GetReplicasJSON returns a JSON representation of all replicas managed by the ReplicaManager.
-// It locks the ReplicaManager for reading, converts the replicas to their protobuf representation,
-// marshals them into a JSON string, and returns the result.
-// If an error occurs during marshaling, it logs a warning and returns an empty string.
 func (m *ReplicaManager) GetReplicasJSON(ctx context.Context, meta *Meta) string {
-	m.rwmutex.RLock()
-	defer m.rwmutex.RUnlock()
+	allReplicas := make([]*metricsinfo.Replica, 0)
+	m.coll2Replicas.Range(func(_ int64, collReplicas []*Replica) bool {
+		for _, r := range collReplicas {
+			channelTowRWNodes := make(map[string][]int64)
+			for k, v := range r.replicaPB.GetChannelNodeInfos() {
+				channelTowRWNodes[k] = v.GetRwNodes()
+			}
 
-	replicas := lo.MapToSlice(m.replicas, func(i typeutil.UniqueID, r *Replica) *metricsinfo.Replica {
-		channelTowRWNodes := make(map[string][]int64)
-		for k, v := range r.replicaPB.GetChannelNodeInfos() {
-			channelTowRWNodes[k] = v.GetRwNodes()
-		}
+			collectionInfo := meta.GetCollection(ctx, r.GetCollectionID())
+			dbID := util.InvalidDBID
+			if collectionInfo == nil {
+				log.Ctx(ctx).Warn("failed to get collection info", zap.Int64("collectionID", r.GetCollectionID()))
+			} else {
+				dbID = collectionInfo.GetDbID()
+			}
 
-		collectionInfo := meta.GetCollection(ctx, r.GetCollectionID())
-		dbID := util.InvalidDBID
-		if collectionInfo == nil {
-			log.Ctx(ctx).Warn("failed to get collection info", zap.Int64("collectionID", r.GetCollectionID()))
-		} else {
-			dbID = collectionInfo.GetDbID()
-		}
-
-		return &metricsinfo.Replica{
-			ID:               r.GetID(),
-			CollectionID:     r.GetCollectionID(),
-			DatabaseID:       dbID,
-			RWNodes:          r.GetNodes(),
-			ResourceGroup:    r.GetResourceGroup(),
-			RONodes:          r.GetRONodes(),
-			ChannelToRWNodes: channelTowRWNodes,
+			allReplicas = append(allReplicas, &metricsinfo.Replica{
+				ID:               r.GetID(),
+				CollectionID:     r.GetCollectionID(),
+				DatabaseID:       dbID,
+				RWNodes:          r.GetNodes(),
+				ResourceGroup:    r.GetResourceGroup(),
+				RONodes:          r.GetRONodes(),
+				ChannelToRWNodes: channelTowRWNodes,
+			})
 		}
+		return true
 	})
-	ret, err := json.Marshal(replicas)
+	ret, err := json.Marshal(allReplicas)
 	if err != nil {
 		log.Warn("failed to marshal replicas", zap.Error(err))
 		return ""
@@ -801,16 +871,16 @@ func (m *ReplicaManager) GetReplicasJSON(ctx context.Context, meta *Meta) string
 // by resource group isolation (each replica only gets streaming nodes from its own resource group).
 // Otherwise, all streaming nodes will be pooled together and assigned fairly across all replicas (fallback mode).
 func (m *ReplicaManager) RecoverSQNodesInCollection(ctx context.Context, collectionID int64, sqnNodesByRG map[string]typeutil.UniqueSet) error {
-	m.rwmutex.Lock()
-	defer m.rwmutex.Unlock()
+	m.collLock.Lock(collectionID)
+	defer m.collLock.Unlock(collectionID)
 
-	collReplicas, ok := m.coll2Replicas[collectionID]
+	replicas, ok := m.coll2Replicas.Get(collectionID)
 	if !ok {
 		return errors.Errorf("collection %d not loaded", collectionID)
 	}
 
 	// Build helpers based on whether we can use resource group isolation.
-	helpers := m.buildSQNodeAssignmentHelpers(collReplicas.replicas, sqnNodesByRG)
+	helpers := m.buildSQNodeAssignmentHelpers(replicas, sqnNodesByRG)
 
 	modifiedReplicas := make([]*Replica, 0)
 	for rgName, helper := range helpers {
@@ -821,7 +891,7 @@ func (m *ReplicaManager) RecoverSQNodesInCollection(ctx context.Context, collect
 			if len(roNodes) == 0 && len(recoverableNodes) == 0 && len(incomingNode) == 0 {
 				return
 			}
-			mutableReplica := m.replicas[assignment.GetReplicaID()].CopyForWrite()
+			mutableReplica := assignment.GetReplica().CopyForWrite()
 			mutableReplica.AddROSQNode(roNodes...)
 			mutableReplica.AddRWSQNode(recoverableNodes...)
 			mutableReplica.AddRWSQNode(incomingNode...)
@@ -839,7 +909,7 @@ func (m *ReplicaManager) RecoverSQNodesInCollection(ctx context.Context, collect
 			modifiedReplicas = append(modifiedReplicas, mutableReplica.IntoReplica())
 		})
 	}
-	return m.put(ctx, modifiedReplicas...)
+	return m.put(ctx, collectionID, modifiedReplicas...)
 }
 
 // buildSQNodeAssignmentHelpers builds assignment helpers for streaming query node recovery.

internal/querycoordv2/meta/replica_manager_helper.go

@@ -170,7 +170,7 @@ func newReplicaAssignmentInfo(replica *Replica, nodeInRG typeutil.UniqueSet) *re
 		return true
 	})
 	return &replicaAssignmentInfo{
-		replicaID:            replica.GetID(),
+		replica:              replica,
 		expectedNodeCount:    0,
 		rwNodes:              rwNodes,
 		newRONodes:           newRONodes,
@@ -204,7 +204,7 @@ func newReplicaSQNAssignmentInfo(replica *Replica, nodes typeutil.UniqueSet) *re
 		return true
 	})
 	return &replicaAssignmentInfo{
-		replicaID:            replica.GetID(),
+		replica:              replica,
 		expectedNodeCount:    0,
 		rwNodes:              rwNodes,
 		newRONodes:           newRONodes,
@@ -214,7 +214,7 @@ func newReplicaSQNAssignmentInfo(replica *Replica, nodes typeutil.UniqueSet) *re
 }
 
 type replicaAssignmentInfo struct {
-	replicaID            typeutil.UniqueID
+	replica              *Replica           // the replica snapshot this assignment was computed from
 	expectedNodeCount    int                // expected node count for each replica.
 	rwNodes              typeutil.UniqueSet // rw nodes is used by current replica. (rw -> rw)
 	newRONodes           typeutil.UniqueSet // new ro nodes for these replica. (rw -> ro)
@@ -222,9 +222,14 @@ type replicaAssignmentInfo struct {
 	unrecoverableRONodes typeutil.UniqueSet // unrecoverable ro nodes for these replica (ro node can't be put back to rw node if it's not in current resource group). (ro -> ro)
 }
 
+// GetReplica returns the replica snapshot this assignment was computed from.
+func (s *replicaAssignmentInfo) GetReplica() *Replica {
+	return s.replica
+}
+
 // GetReplicaID returns the replica id for these replica.
 func (s *replicaAssignmentInfo) GetReplicaID() typeutil.UniqueID {
-	return s.replicaID
+	return s.replica.GetID()
 }
 
 // GetNewRONodes returns the new ro nodes for these replica.
@@ -307,7 +312,7 @@ func (s replicaAssignmentInfoSortByAvailableAndRecoverable) Less(i, j int) bool
 
 	// Reach stable sort result by replica id.
 	// Otherwise unstable assignment may cause unnecessary node transfer.
-	return left < right || (left == right && s.replicaAssignmentInfoSorter[i].replicaID < s.replicaAssignmentInfoSorter[j].replicaID)
+	return left < right || (left == right && s.replicaAssignmentInfoSorter[i].GetReplicaID() < s.replicaAssignmentInfoSorter[j].GetReplicaID())
 }
 
 // newReplicaSQNAssignmentHelper creates a new replicaSQNAssignmentHelper.

internal/querycoordv2/meta/replica_manager_test.go

@@ -20,6 +20,7 @@ import (
 	"context"
 	"testing"
 
+	"github.com/cockroachdb/errors"
 	"github.com/samber/lo"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
@@ -373,7 +374,17 @@ func (suite *ReplicaManagerSuite) TestResourceGroup() {
 }
 
 func (suite *ReplicaManagerSuite) clearMemory() {
-	suite.mgr.replicas = make(map[int64]*Replica)
+	suite.mgr.coll2Replicas.Range(func(collID int64, _ []*Replica) bool {
+		suite.mgr.coll2Replicas.Remove(collID)
+		return true
+	})
+	suite.mgr.flatReplicas.Range(func(replicaID int64, _ *Replica) bool {
+		suite.mgr.flatReplicas.Remove(replicaID)
+		return true
+	})
+	for _, replicaID := range suite.mgr.queryInvisibleReplicas.Collect() {
+		suite.mgr.queryInvisibleReplicas.Remove(replicaID)
+	}
 }
 
 type ReplicaManagerV2Suite struct {
@@ -599,8 +610,8 @@ func (suite *ReplicaManagerV2Suite) recoverReplica(k int, clearOutbound bool) {
 				replicas := suite.mgr.GetByCollection(ctx, id)
 				for _, r := range replicas {
 					outboundNodes := r.GetRONodes()
-					suite.mgr.RemoveNode(ctx, r.GetID(), outboundNodes...)
-					suite.mgr.RemoveSQNode(ctx, r.GetID(), r.GetROSQNodes()...)
+					suite.mgr.RemoveNode(ctx, id, r.GetID(), outboundNodes...)
+					suite.mgr.RemoveSQNode(ctx, id, r.GetID(), r.GetROSQNodes()...)
 				}
 			}
 		}
@@ -643,7 +654,7 @@ func TestSQNodeResourceGroupIsolation(t *testing.T) {
 		Nodes:         []int64{},
 	})
 
-	err := mgr.put(ctx, replica1, replica2, replica3)
+	err := mgr.Put(ctx, replica1, replica2, replica3)
 	assert.NoError(t, err)
 
 	// Test case 1: Resource group isolation mode.
@@ -684,7 +695,7 @@ func TestSQNodeResourceGroupIsolation(t *testing.T) {
 		Nodes:         []int64{},
 	})
 
-	err = mgr.put(ctx, replica4, replica5)
+	err = mgr.Put(ctx, replica4, replica5)
 	assert.NoError(t, err)
 
 	err = mgr.RecoverSQNodesInCollection(ctx, 200, sqNodesByRG)
@@ -706,12 +717,12 @@ func TestSQNodeResourceGroupIsolation(t *testing.T) {
 
 	// Test case 3: Rolling upgrade compatibility mode.
 	// Old StreamingNodes without RG labels are reported under the default RG. Replicas
-	// in old/uncovered RGs should use only that legacy default pool, while replicas in
-	// newly covered RGs should keep strict RG isolation.
+	// in default or old/uncovered RGs should use only that legacy default pool, while
+	// replicas in newly covered RGs should keep strict RG isolation.
 	replica8 := newReplica(&querypb.Replica{
 		ID:            8,
 		CollectionID:  250,
-		ResourceGroup: "RG_OLD_REPLICA", // Not covered by labeled streaming nodes.
+		ResourceGroup: DefaultResourceGroupName,
 		Nodes:         []int64{},
 	})
 	replica9 := newReplica(&querypb.Replica{
@@ -720,12 +731,18 @@ func TestSQNodeResourceGroupIsolation(t *testing.T) {
 		ResourceGroup: "RG1",
 		Nodes:         []int64{},
 	})
+	replica10 := newReplica(&querypb.Replica{
+		ID:            10,
+		CollectionID:  250,
+		ResourceGroup: "RG_OLD_REPLICA", // Not covered by labeled streaming nodes.
+		Nodes:         []int64{},
+	})
 
-	err = mgr.put(ctx, replica8, replica9)
+	err = mgr.Put(ctx, replica8, replica9, replica10)
 	assert.NoError(t, err)
 
 	rollingUpgradeSQNodesByRG := map[string]typeutil.UniqueSet{
-		DefaultResourceGroupName: typeutil.NewUniqueSet(901, 902),
+		DefaultResourceGroupName: typeutil.NewUniqueSet(901, 902, 903, 904),
 		"RG1":                    typeutil.NewUniqueSet(101, 102),
 	}
 
@@ -734,7 +751,11 @@ func TestSQNodeResourceGroupIsolation(t *testing.T) {
 
 	updatedReplica8 := mgr.Get(ctx, 8)
 	updatedReplica9 := mgr.Get(ctx, 9)
-	assert.ElementsMatch(t, []int64{901, 902}, updatedReplica8.GetRWSQNodes())
+	updatedReplica10 := mgr.Get(ctx, 10)
+	assert.Len(t, updatedReplica8.GetRWSQNodes(), 2)
+	assert.Len(t, updatedReplica10.GetRWSQNodes(), 2)
+	defaultPoolNodes := append(updatedReplica8.GetRWSQNodes(), updatedReplica10.GetRWSQNodes()...)
+	assert.ElementsMatch(t, []int64{901, 902, 903, 904}, defaultPoolNodes)
 	assert.ElementsMatch(t, []int64{101, 102}, updatedReplica9.GetRWSQNodes())
 
 	// Test case 4: Strict isolation mode (with config enabled).
@@ -756,7 +777,7 @@ func TestSQNodeResourceGroupIsolation(t *testing.T) {
 		Nodes:         []int64{},
 	})
 
-	err = mgr.put(ctx, replica6, replica7)
+	err = mgr.Put(ctx, replica6, replica7)
 	assert.NoError(t, err)
 
 	err = mgr.RecoverSQNodesInCollection(ctx, 300, sqNodesByRG)
@@ -797,7 +818,7 @@ func TestSQNodeRecoveryWithRONodes(t *testing.T) {
 		RoSqNodes:     []int64{104},           // RO SQ node (previously removed from available set)
 	})
 
-	err := mgr.put(ctx, replica1)
+	err := mgr.Put(ctx, replica1)
 	assert.NoError(t, err)
 
 	// Available SQ nodes: only 101, 102 remain. 103 is no longer available (will become RO).
@@ -843,7 +864,7 @@ func TestSQNodeRecoveryWithUnrecoverableNodes(t *testing.T) {
 		RoSqNodes:     []int64{999}, // Node 999 is not in any available set - unrecoverable
 	})
 
-	err := mgr.put(ctx, replica1)
+	err := mgr.Put(ctx, replica1)
 	assert.NoError(t, err)
 
 	// Available SQ nodes don't include 999
@@ -889,10 +910,10 @@ func TestGetReplicasJSON(t *testing.T) {
 		Nodes:         []int64{4, 5, 6},
 	})
 
-	err := replicaManager.put(ctx, replica1)
+	err := replicaManager.Put(ctx, replica1)
 	assert.NoError(t, err)
 
-	err = replicaManager.put(ctx, replica2)
+	err = replicaManager.Put(ctx, replica2)
 	assert.NoError(t, err)
 
 	meta := &Meta{
@@ -941,3 +962,230 @@ func TestGetReplicasJSON(t *testing.T) {
 		checkResult(replica)
 	}
 }
+
+func TestReplicaManagerCollectionViewAfterPartialUpdateAndRemove(t *testing.T) {
+	catalog := mocks.NewQueryCoordCatalog(t)
+	catalog.EXPECT().SaveReplica(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil).Maybe()
+	catalog.EXPECT().SaveReplica(mock.Anything, mock.Anything, mock.Anything).Return(nil).Maybe()
+	catalog.EXPECT().SaveReplica(mock.Anything, mock.Anything).Return(nil).Maybe()
+	catalog.EXPECT().ReleaseReplica(mock.Anything, mock.Anything, mock.Anything).Return(nil).Maybe()
+	idAllocator := RandomIncrementIDAllocator()
+	mgr := NewReplicaManager(idAllocator, catalog)
+	ctx := context.Background()
+
+	collID := int64(7000)
+	replicas, err := mgr.Spawn(ctx, collID, map[string]int{"rg1": 3}, nil, commonpb.LoadPriority_LOW)
+	assert.NoError(t, err)
+	assert.Len(t, replicas, 3)
+
+	err = mgr.MoveReplica(ctx, collID, "rg2", replicas[1:])
+	assert.NoError(t, err)
+
+	updated := mgr.GetByCollection(ctx, collID)
+	assert.Len(t, updated, 3)
+	assert.Equal(t, replicas[0].GetID(), updated[0].GetID())
+	assert.Equal(t, "rg1", updated[0].GetResourceGroup())
+	assert.Equal(t, replicas[1].GetID(), updated[1].GetID())
+	assert.Equal(t, "rg2", updated[1].GetResourceGroup())
+	assert.Equal(t, replicas[2].GetID(), updated[2].GetID())
+	assert.Equal(t, "rg2", updated[2].GetResourceGroup())
+
+	err = mgr.RemoveReplicas(ctx, collID, replicas[1].GetID())
+	assert.NoError(t, err)
+
+	remaining := mgr.GetByCollection(ctx, collID)
+	assert.Len(t, remaining, 2)
+	assert.Equal(t, replicas[0].GetID(), remaining[0].GetID())
+	assert.Equal(t, replicas[2].GetID(), remaining[1].GetID())
+	assert.Nil(t, mgr.Get(ctx, replicas[1].GetID()))
+	assert.NotNil(t, mgr.Get(ctx, replicas[0].GetID()))
+	assert.NotNil(t, mgr.Get(ctx, replicas[2].GetID()))
+}
+
+func TestReplicaManagerPutMaintainsIndexesAcrossCollections(t *testing.T) {
+	catalog := mocks.NewQueryCoordCatalog(t)
+	catalog.EXPECT().SaveReplica(mock.Anything, mock.Anything, mock.Anything).Return(nil).Once()
+	idAllocator := RandomIncrementIDAllocator()
+	mgr := NewReplicaManager(idAllocator, catalog)
+	ctx := context.Background()
+
+	replica1 := newReplica(&querypb.Replica{
+		ID:            1,
+		CollectionID:  10,
+		ResourceGroup: "rg1",
+		Nodes:         []int64{101},
+	})
+	replica2 := newReplica(&querypb.Replica{
+		ID:            2,
+		CollectionID:  20,
+		ResourceGroup: "rg2",
+		Nodes:         []int64{201},
+	})
+
+	err := mgr.Put(ctx, replica1, replica2)
+	assert.NoError(t, err)
+
+	assert.Equal(t, replica1, mgr.Get(ctx, replica1.GetID()))
+	assert.Equal(t, replica2, mgr.Get(ctx, replica2.GetID()))
+	assert.Equal(t, []*Replica{replica1}, mgr.GetByCollection(ctx, replica1.GetCollectionID()))
+	assert.Equal(t, []*Replica{replica2}, mgr.GetByCollection(ctx, replica2.GetCollectionID()))
+	assert.Equal(t, replica1, mgr.GetByCollectionAndNode(ctx, replica1.GetCollectionID(), int64(101)))
+	assert.Equal(t, replica2, mgr.GetByCollectionAndNode(ctx, replica2.GetCollectionID(), int64(201)))
+	assert.ElementsMatch(t, []*Replica{replica1}, mgr.GetByNode(ctx, int64(101)))
+	assert.ElementsMatch(t, []*Replica{replica2}, mgr.GetByResourceGroup(ctx, "rg2"))
+}
+
+func TestReplicaManagerPutCrossCollectionPersistErrorIsAtomic(t *testing.T) {
+	catalog := mocks.NewQueryCoordCatalog(t)
+	saveErr := errors.New("save failed")
+	catalog.EXPECT().SaveReplica(mock.Anything, mock.Anything, mock.Anything).Return(saveErr).Once()
+	mgr := NewReplicaManager(RandomIncrementIDAllocator(), catalog)
+	ctx := context.Background()
+
+	replica1 := newReplica(&querypb.Replica{
+		ID:            1,
+		CollectionID:  10,
+		ResourceGroup: "rg1",
+	})
+	replica2 := newReplica(&querypb.Replica{
+		ID:            2,
+		CollectionID:  20,
+		ResourceGroup: "rg2",
+	})
+
+	err := mgr.Put(ctx, replica1, replica2)
+	assert.ErrorIs(t, err, saveErr)
+	assert.Nil(t, mgr.Get(ctx, replica1.GetID()))
+	assert.Nil(t, mgr.Get(ctx, replica2.GetID()))
+	assert.Empty(t, mgr.GetByCollection(ctx, replica1.GetCollectionID()))
+	assert.Empty(t, mgr.GetByCollection(ctx, replica2.GetCollectionID()))
+}
+
+func TestReplicaManagerRemoveCollectionReleasesKVWhenNotLoaded(t *testing.T) {
+	collectionID := int64(10)
+	ctx := context.Background()
+
+	t.Run("success", func(t *testing.T) {
+		catalog := mocks.NewQueryCoordCatalog(t)
+		catalog.EXPECT().ReleaseReplicas(mock.Anything, collectionID).Return(nil).Once()
+		mgr := NewReplicaManager(nil, catalog)
+
+		err := mgr.RemoveCollection(ctx, collectionID)
+		assert.NoError(t, err)
+		assert.Empty(t, mgr.GetByCollection(ctx, collectionID))
+	})
+
+	t.Run("release error", func(t *testing.T) {
+		catalog := mocks.NewQueryCoordCatalog(t)
+		releaseErr := errors.New("release failed")
+		catalog.EXPECT().ReleaseReplicas(mock.Anything, collectionID).Return(releaseErr).Once()
+		mgr := NewReplicaManager(nil, catalog)
+
+		err := mgr.RemoveCollection(ctx, collectionID)
+		assert.ErrorIs(t, err, releaseErr)
+		assert.Empty(t, mgr.GetByCollection(ctx, collectionID))
+	})
+}
+
+func TestReplicaManagerPersistErrorPaths(t *testing.T) {
+	ctx := context.Background()
+	saveErr := errors.New("save failed")
+
+	t.Run("direct persist", func(t *testing.T) {
+		catalog := mocks.NewQueryCoordCatalog(t)
+		replica := newReplica(&querypb.Replica{ID: 1, CollectionID: 10, ResourceGroup: "rg1"})
+		catalog.EXPECT().SaveReplica(mock.Anything, replica.replicaPB).Return(saveErr).Once()
+		mgr := NewReplicaManager(nil, catalog)
+
+		assert.NoError(t, mgr.put(ctx, replica.GetCollectionID()))
+		err := mgr.put(ctx, replica.GetCollectionID(), replica)
+		assert.ErrorIs(t, err, saveErr)
+	})
+
+	t.Run("remove node", func(t *testing.T) {
+		catalog := mocks.NewQueryCoordCatalog(t)
+		catalog.EXPECT().SaveReplica(mock.Anything, mock.Anything).Return(saveErr).Once()
+		mgr := NewReplicaManager(nil, catalog)
+		replica := newReplica(&querypb.Replica{
+			ID:            1,
+			CollectionID:  10,
+			ResourceGroup: "rg1",
+			Nodes:         []int64{101},
+			RoNodes:       []int64{102},
+		})
+		mgr.putReplicasInMemory(replica.GetCollectionID(), replica)
+
+		err := mgr.RemoveNode(ctx, replica.GetCollectionID(), replica.GetID(), int64(102))
+		assert.ErrorIs(t, err, saveErr)
+		assert.ElementsMatch(t, []int64{102}, mgr.Get(ctx, replica.GetID()).GetRONodes())
+	})
+
+	t.Run("remove streaming query node", func(t *testing.T) {
+		catalog := mocks.NewQueryCoordCatalog(t)
+		catalog.EXPECT().SaveReplica(mock.Anything, mock.Anything).Return(saveErr).Once()
+		mgr := NewReplicaManager(nil, catalog)
+		replica := newReplica(&querypb.Replica{
+			ID:            1,
+			CollectionID:  10,
+			ResourceGroup: "rg1",
+			RwSqNodes:     []int64{201},
+			RoSqNodes:     []int64{202},
+		})
+		mgr.putReplicasInMemory(replica.GetCollectionID(), replica)
+
+		err := mgr.RemoveSQNode(ctx, replica.GetCollectionID(), replica.GetID(), int64(202))
+		assert.ErrorIs(t, err, saveErr)
+		assert.ElementsMatch(t, []int64{202}, mgr.Get(ctx, replica.GetID()).GetROSQNodes())
+	})
+}
+
+func TestReplicaManagerSpawnWaitRGReadyRecovery(t *testing.T) {
+	catalog := mocks.NewQueryCoordCatalog(t)
+	catalog.EXPECT().SaveReplica(mock.Anything, mock.Anything).Return(nil).Twice()
+	mgr := NewReplicaManager(RandomIncrementIDAllocator(), catalog)
+	ctx := context.Background()
+	collID := int64(10)
+
+	replicas, err := mgr.Spawn(ctx, collID, map[string]int{"rg1": 1}, nil, commonpb.LoadPriority_LOW, WithNeedWaitRGReady())
+	assert.NoError(t, err)
+	assert.Len(t, replicas, 1)
+	replicaID := replicas[0].GetID()
+	assert.True(t, mgr.Get(ctx, replicaID).NeedWaitRGReady())
+
+	err = mgr.RecoverNodesInCollection(ctx, collID, map[string]*ResourceGroup{
+		"rg1": {
+			name:  "rg1",
+			nodes: typeutil.NewUniqueSet(int64(101)),
+			cfg:   newResourceGroupConfig(2, 2),
+		},
+	})
+	assert.NoError(t, err)
+	assert.Empty(t, mgr.Get(ctx, replicaID).GetRWNodes())
+	assert.True(t, mgr.Get(ctx, replicaID).NeedWaitRGReady())
+
+	err = mgr.RecoverNodesInCollection(ctx, collID, map[string]*ResourceGroup{
+		"rg1": newTestResourceGroup("rg1", typeutil.NewUniqueSet(101, 102)),
+	})
+	assert.NoError(t, err)
+	assert.ElementsMatch(t, []int64{101, 102}, mgr.Get(ctx, replicaID).GetRWNodes())
+	assert.False(t, mgr.Get(ctx, replicaID).NeedWaitRGReady())
+}
+
+func TestReplicaManagerMoveReplicaPersistError(t *testing.T) {
+	catalog := mocks.NewQueryCoordCatalog(t)
+	saveErr := errors.New("move failed")
+	catalog.EXPECT().SaveReplica(mock.Anything, mock.Anything).Return(nil).Once()
+	catalog.EXPECT().SaveReplica(mock.Anything, mock.Anything).Return(saveErr).Once()
+	mgr := NewReplicaManager(nil, catalog)
+	ctx := context.Background()
+	replica := newReplica(&querypb.Replica{
+		ID:            1,
+		CollectionID:  10,
+		ResourceGroup: "rg1",
+	})
+	assert.NoError(t, mgr.Put(ctx, replica))
+
+	err := mgr.MoveReplica(ctx, replica.GetCollectionID(), "rg2", []*Replica{replica})
+	assert.ErrorIs(t, err, saveErr)
+	assert.Equal(t, "rg1", mgr.Get(ctx, replica.GetID()).GetResourceGroup())
+}

internal/querycoordv2/meta/replica_visibility_test.go

@@ -61,7 +61,7 @@ func TestReplicaManagerQueryVisibility(t *testing.T) {
 	mutableReplica := invisibleReplica.CopyForWrite()
 	mutableReplica.SetQueryInvisible(true)
 	invisibleReplica = mutableReplica.IntoReplica()
-	manager.putReplicaInMemory(visibleReplica, invisibleReplica)
+	manager.putReplicasInMemory(100, visibleReplica, invisibleReplica)
 
 	invisibleReplicas := manager.GetQueryInvisibleReplicas(ctx)
 	assert.Len(t, invisibleReplicas, 1)
@@ -72,10 +72,10 @@ func TestReplicaManagerQueryVisibility(t *testing.T) {
 	assert.True(t, manager.Get(ctx, 11).IsQueryVisible())
 	assert.Empty(t, manager.GetQueryInvisibleReplicas(ctx))
 
-	manager.putReplicaInMemory(invisibleReplica)
+	manager.putReplicasInMemory(100, invisibleReplica)
 	assert.Len(t, manager.GetQueryInvisibleReplicas(ctx), 1)
 	mutableReplica = invisibleReplica.CopyForWrite()
 	mutableReplica.SetQueryInvisible(false)
-	manager.putReplicaInMemory(mutableReplica.IntoReplica())
+	manager.putReplicasInMemory(100, mutableReplica.IntoReplica())
 	assert.Empty(t, manager.GetQueryInvisibleReplicas(ctx))
 }

internal/querycoordv2/observers/replica_observer.go

@@ -155,7 +155,7 @@ func (ob *ReplicaObserver) checkStreamingQueryNodesInReplica(sqNodeIDsByRG map[s
 				zap.Int64s("roNodes", roSQNodes),
 				zap.Int64s("rwNodes", rwSQNodes),
 			)
-			if err := ob.meta.RemoveSQNode(ctx, replica.GetID(), removeNodes...); err != nil {
+			if err := ob.meta.RemoveSQNode(ctx, collectionID, replica.GetID(), removeNodes...); err != nil {
 				logger.Warn("fail to remove streaming query node from replica", zap.Error(err))
 				continue
 			}
@@ -213,7 +213,7 @@ func (ob *ReplicaObserver) checkNodesInReplica() {
 			if len(removeNodes) == 0 {
 				continue
 			}
-			if err := ob.meta.RemoveNode(ctx, replica.GetID(), removeNodes...); err != nil {
+			if err := ob.meta.RemoveNode(ctx, collectionID, replica.GetID(), removeNodes...); err != nil {
 				logger.Warn("fail to remove node from replica",
 					zap.Int64s("removedNodes", removeNodes),
 					zap.Error(err))

internal/querycoordv2/services_test.go

@@ -1475,8 +1475,8 @@ func (suite *ServiceSuite) TestLoadBalanceWithEmptySegmentList() {
 	defer func() {
 		for _, collection := range suite.collections {
 			replicas := suite.meta.GetByCollection(ctx, collection)
-			suite.meta.RemoveNode(ctx, replicas[0].GetID(), srcNode)
-			suite.meta.RemoveNode(ctx, replicas[0].GetID(), dstNode)
+			suite.meta.RemoveNode(ctx, collection, replicas[0].GetID(), srcNode)
+			suite.meta.RemoveNode(ctx, collection, replicas[0].GetID(), dstNode)
 		}
 		suite.nodeMgr.Remove(1001)
 		suite.nodeMgr.Remove(1002)
@@ -1619,7 +1619,7 @@ func (suite *ServiceSuite) TestLoadBalanceFailed() {
 		suite.NoError(err)
 		suite.Equal(commonpb.ErrorCode_UnexpectedError, resp.ErrorCode)
 		suite.nodeMgr.Remove(10)
-		suite.meta.RemoveNode(ctx, replicas[0].GetID(), 10)
+		suite.meta.RemoveNode(ctx, collection, replicas[0].GetID(), 10)
 	}
 }