enhance: scope collection load and state reads

docs/design-docs/design_docs/qviews/balancer_design.md

@@ -158,6 +158,13 @@ a scoped trigger cannot be narrowed safely, it uses the same full planning
 scope; only an explicit or periodic full trigger also rebuilds the row-count
 ledger.
 
+Scope selection precedes load-config reads too. Scoped cycles capture only
+`collectionIDs` via `LoadConfigStore.SnapshotForCollections`; full cycles and
+conservative full-scope fallbacks use `Snapshot()`. The scoped config snapshot
+retains each selected collection's version, including version zero for absent
+configs. Node row totals remain cluster-wide through the existing incremental
+ledger.
+
 The DataView provider exposes both full and collection-scoped reads:
 
 ```go
@@ -261,6 +268,14 @@ func (s *LoadConfigStore) Put(ctx context.Context, cfg *LoadConfig) error
 func (s *LoadConfigStore) Remove(ctx context.Context, collectionID int64) error
 
 func (s *LoadConfigStore) Snapshot() *LoadConfigSnapshot
+func (s *LoadConfigStore) SnapshotForCollections(collectionIDs []int64) *LoadConfigSnapshot
+func (s *LoadConfigStore) Get(collectionID int64) LoadConfigEntry
+
+type LoadConfigEntry struct {
+    Config        *LoadConfig
+    ConfigVersion uint64
+    StoreVersion  uint64
+}
 
 // LoadConfig is the complete load configuration for a collection.
 type LoadConfig struct {
@@ -297,6 +312,14 @@ Legacy proto fields are kept for wire compatibility but ignored by the new desig
 
 **Copy-On-Write semantics**: Put clones its input before storing, so callers may reuse/mutate their input freely. Snapshot returns pointers into the store's immutable view — callers must call `.Clone()` before any mutation. The store never modifies published snapshots in place; updates advance the live version and the next Snapshot call lazily publishes a new immutable view.
 
+Point reads and scoped snapshots capture config pointers and their versions
+under one read lock, without rebuilding the resident full snapshot. An empty
+collection list selects none. The scoped replica index is built after releasing
+the lock, using the captured immutable configs. Single-collection metadata reads
+use `Get`; explicit collection lists use `SnapshotForCollections`.
+`GetQueryViewLoadInfo` preserves its existing global `StoreVersion`
+response; the collection-specific load-info version remains `ConfigVersion`.
+
 **Write amplification**: Put always writes the full config (no diff). Orphan partitions / replicas (present in previous state but absent from new config) are deleted. This is intentionally simple — dedup / diff optimization can be added later if write volume becomes a concern.
 
 #### ShardViewRegistry
@@ -316,6 +339,7 @@ func (r *ShardViewRegistry) Ensure(shardID qviews.ShardID) *ShardViewManager
 func (r *ShardViewRegistry) Get(shardID qviews.ShardID) *ShardViewManager
 func (r *ShardViewRegistry) Snapshot() *ShardViewSnapshot
 func (r *ShardViewRegistry) SnapshotForShards(shardIDs []qviews.ShardID) *ShardViewSnapshot
+func (r *ShardViewRegistry) SnapshotForCollection(collectionID int64) *ShardViewSnapshot
 func (r *ShardViewRegistry) CollectionShards(collectionID int64) []qviews.ShardID
 func (r *ShardViewRegistry) NodeShards(nodeID int64) []qviews.ShardID
 func (r *ShardViewRegistry) ShardIDs() []qviews.ShardID
@@ -328,11 +352,15 @@ Maintains live per-shard stats via callbacks from each `ShardViewManager`.
 publishes the resident full snapshot lazily; `SnapshotForShards()` copies only
 the requested `ShardID -> *ShardStats` entries.
 
-Shard managers remain resident for the lifetime of the Registry. A QueryView
-reaching Dropped removes that state machine from its manager, but does not
-remove the manager. Collection index entries are maintained independently of
-view state transitions. Node index entries follow the latest stats and
-disappear when the shard no longer references that node.
+Load-progress reads use `SnapshotForCollection`, which captures the collection
+index and its stats under one read lock. Progress still filters the collection's
+current replica IDs and counts resident shards with an Up view.
+
+After the last QueryView completes durable removal, the Registry reclaims its
+manager and removes the shard's stats and collection/node index entries. The
+callback rechecks manager identity and emptiness before removal; a later
+`Ensure` can create a fresh manager. Node index entries also follow the latest
+stats and disappear when the shard no longer references that node.
 
 #### CollectionLoadManager (Facade)
 

internal/querycoordv2/ddl_callbacks_alter_load_info_load_collection.go

@@ -66,7 +66,7 @@ func (s *Server) broadcastAlterLoadConfigCollectionV2ForLoadCollection(ctx conte
 		return err
 	}
 
-	currentLoadConfig := s.qviewsRuntime.loadConfigStore.Snapshot().ConfigsMap()[req.GetCollectionID()]
+	currentLoadConfig := s.qviewsRuntime.loadConfigStore.Get(req.GetCollectionID()).Config
 	// only check node number when the collection is not loaded
 	expectedReplicasNumber, err := utils.AssignReplica(ctx, s.meta, resourceGroups, replicaNumber, currentLoadConfig == nil)
 	if err != nil {

internal/querycoordv2/scoped_load_reads_test.go

@@ -0,0 +1,112 @@
+// 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 querycoordv2
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+
+	"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
+	metastoremocks "github.com/milvus-io/milvus/internal/metastore/mocks"
+	"github.com/milvus-io/milvus/internal/querycoordv2/meta"
+	"github.com/milvus-io/milvus/internal/views/coord/coordview"
+	"github.com/milvus-io/milvus/internal/views/coord/coordview/syncer"
+	"github.com/milvus-io/milvus/internal/views/coord/loadmgr"
+	"github.com/milvus-io/milvus/internal/views/qviews"
+	"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
+	"github.com/milvus-io/milvus/pkg/v3/proto/viewpb"
+	"github.com/milvus-io/milvus/pkg/v3/util/merr"
+)
+
+type loadProgressSyncer struct{}
+
+func (*loadProgressSyncer) SyncViews(context.Context, syncer.SyncGroup) error { return nil }
+func (*loadProgressSyncer) Close() error                                      { return nil }
+
+func TestScopedLoadReadsPreserveProgressAndResponseVersion(t *testing.T) {
+	ctx := context.Background()
+	previousCache := meta.GlobalFailedLoadCache
+	meta.GlobalFailedLoadCache = meta.NewFailedLoadCache()
+	t.Cleanup(func() { meta.GlobalFailedLoadCache = previousCache })
+	catalog := metastoremocks.NewQueryCoordCatalog(t)
+	catalog.EXPECT().GetCollections(mock.Anything).Return([]*querypb.CollectionLoadInfo{
+		{CollectionID: 1, LoadFields: []int64{100}, FieldIndexID: map[int64]int64{100: 200}},
+		{CollectionID: 2},
+	}, nil).Once()
+	catalog.EXPECT().GetPartitions(mock.Anything, mock.Anything).Return(map[int64][]*querypb.PartitionLoadInfo{
+		1: {{CollectionID: 1, PartitionID: 101}},
+	}, nil).Once()
+	catalog.EXPECT().GetReplicas(mock.Anything).Return([]*querypb.Replica{
+		{ID: 10, CollectionID: 1}, {ID: 20, CollectionID: 2},
+	}, nil).Once()
+	store, err := loadmgr.RecoverLoadConfigStore(ctx, catalog)
+	require.NoError(t, err)
+	first := qviews.ShardID{ReplicaID: 10, VChannel: "by-dev-rootcoord-dml_1v0"}
+	missing := qviews.ShardID{ReplicaID: 10, VChannel: "by-dev-rootcoord-dml_1v1"}
+	oldReplica := qviews.ShardID{ReplicaID: 11, VChannel: first.VChannel}
+	other := qviews.ShardID{ReplicaID: 20, VChannel: "by-dev-rootcoord-dml_2v0"}
+	up := testPersistedQueryView(1, first)
+	registry, err := coordview.RecoverShardViewRegistry(ctx, &fakeQueryViewCatalog{
+		views: []*viewpb.QueryViewOfShard{up, testPersistedQueryView(2, other)},
+	}, &loadProgressSyncer{})
+	require.NoError(t, err)
+	t.Cleanup(registry.Close)
+	registry.Ensure(oldReplica) // A residual replica must not lower the current replica's progress.
+	server := &Server{ctx: ctx, qviewsRuntime: &qviewsRuntime{loadConfigStore: store, shardViewRegistry: registry}}
+	server.UpdateStateCode(commonpb.StateCode_Healthy)
+	cfg := store.Get(1).Config
+	require.EqualValues(t, 100, server.qviewsLoadPercentage(cfg))
+	require.Zero(t, server.qviewsLoadPercentage(&loadmgr.LoadConfig{CollectionID: 3}))
+
+	registry.Ensure(missing)
+	assert.EqualValues(t, 50, server.qviewsLoadPercentage(cfg))
+
+	// Advance only the other collection: the RPC's global version must remain
+	// distinct from the selected collection's load-info version.
+	catalog.EXPECT().SaveCollection(mock.Anything, mock.Anything).Return(nil).Once()
+	catalog.EXPECT().SaveReplica(mock.Anything, mock.Anything).Return(nil).Once()
+	require.NoError(t, store.Put(ctx, store.Get(2).Config))
+	entry := store.Get(1)
+	require.NotEqual(t, entry.ConfigVersion, entry.StoreVersion)
+	info, err := server.GetQueryViewLoadInfo(ctx, &querypb.GetQueryViewLoadInfoRequest{CollectionID: 1})
+	require.NoError(t, merr.CheckRPCCall(info, err))
+	assert.Equal(t, entry.StoreVersion, info.GetVersion())
+	assert.Equal(t, []int64{101}, info.GetPartitionIDs())
+	info.PartitionIDs[0] = -1
+	info.LoadFields[0].IndexId = -1
+	assert.Equal(t, []int64{101}, store.Get(1).Config.PartitionIDs)
+	assert.EqualValues(t, 200, store.Get(1).Config.LoadFields[0].IndexId)
+
+	selected, err := server.ShowLoadCollections(ctx, &querypb.ShowCollectionsRequest{CollectionIDs: []int64{1, 1}})
+	require.NoError(t, merr.CheckRPCCall(selected, err))
+	assert.Equal(t, []int64{1}, selected.CollectionIDs)
+	assert.Equal(t, []int64{50}, selected.InMemoryPercentages)
+	all, err := server.ShowLoadCollections(ctx, &querypb.ShowCollectionsRequest{})
+	require.NoError(t, merr.CheckRPCCall(all, err))
+	assert.ElementsMatch(t, []int64{1, 2}, all.CollectionIDs)
+	partitions, err := server.ShowLoadPartitions(ctx, &querypb.ShowPartitionsRequest{CollectionID: 1})
+	require.NoError(t, merr.CheckRPCCall(partitions, err))
+	assert.Equal(t, []int64{101}, partitions.PartitionIDs)
+	assert.Equal(t, []int64{50}, partitions.InMemoryPercentages)
+	absent, err := server.ShowLoadCollections(ctx, &querypb.ShowCollectionsRequest{CollectionIDs: []int64{3}})
+	require.NoError(t, err)
+	assert.ErrorIs(t, merr.Error(absent.GetStatus()), merr.ErrCollectionNotLoaded)
+}

internal/querycoordv2/services.go

@@ -72,13 +72,16 @@ func (s *Server) ShowLoadCollections(ctx context.Context, req *querypb.ShowColle
 	defer meta.GlobalFailedLoadCache.TryExpire()
 
 	isGetAll := false
-	configs := s.qviewsRuntime.loadConfigStore.Snapshot().ConfigsMap()
+	var configs map[int64]*loadmgr.LoadConfig
 	collectionSet := typeutil.NewUniqueSet(req.GetCollectionIDs()...)
 	if len(req.GetCollectionIDs()) == 0 {
+		configs = s.qviewsRuntime.loadConfigStore.Snapshot().ConfigsMap()
 		for collectionID := range configs {
 			collectionSet.Insert(collectionID)
 		}
 		isGetAll = true
+	} else {
+		configs = s.qviewsRuntime.loadConfigStore.SnapshotForCollections(req.GetCollectionIDs()).ConfigsMap()
 	}
 	collections := collectionSet.Collect()
 
@@ -137,7 +140,7 @@ func (s *Server) ShowLoadPartitions(ctx context.Context, req *querypb.ShowPartit
 	}
 	defer meta.GlobalFailedLoadCache.TryExpire()
 
-	cfg := s.qviewsRuntime.loadConfigStore.Snapshot().ConfigsMap()[req.GetCollectionID()]
+	cfg := s.qviewsRuntime.loadConfigStore.Get(req.GetCollectionID()).Config
 	if cfg == nil {
 		err := meta.GlobalFailedLoadCache.Get(req.GetCollectionID())
 		if err != nil {
@@ -204,7 +207,7 @@ func (s *Server) qviewsLoadPercentage(cfg *loadmgr.LoadConfig) int64 {
 	}
 	total := int64(0)
 	loaded := int64(0)
-	for shardID, stats := range s.qviewsRuntime.shardViewRegistry.Snapshot().StatsMap() {
+	for shardID, stats := range s.qviewsRuntime.shardViewRegistry.SnapshotForCollection(cfg.CollectionID).StatsMap() {
 		if !replicaIDs.Contain(shardID.ReplicaID) {
 			continue
 		}
@@ -255,13 +258,13 @@ func (s *Server) GetQueryViewLoadInfo(ctx context.Context, req *querypb.GetQuery
 		resp.Status = merr.Status(merr.WrapErrServiceInternalMsg("query view runtime is nil"))
 		return resp, nil
 	}
-	snapshot := s.qviewsRuntime.loadConfigStore.Snapshot()
-	cfg := snapshot.ConfigsMap()[req.GetCollectionID()]
+	entry := s.qviewsRuntime.loadConfigStore.Get(req.GetCollectionID())
+	cfg := entry.Config
 	if cfg == nil {
 		resp.Status = merr.Status(merr.WrapErrCollectionNotLoaded(req.GetCollectionID()))
 		return resp, nil
 	}
-	resp.Version = snapshot.Version()
+	resp.Version = entry.StoreVersion
 	resp.PartitionIDs = append([]int64(nil), cfg.PartitionIDs...)
 	resp.LoadFields = cloneLoadFields(cfg.LoadFields)
 	return resp, nil

internal/views/coord/balancer/scoped_load_config_test.go

@@ -0,0 +1,87 @@
+// 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 balancer
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+
+	"github.com/milvus-io/milvus/internal/metastore/mocks"
+	"github.com/milvus-io/milvus/internal/views/coord/loadmgr"
+	"github.com/milvus-io/milvus/internal/views/qviews"
+	"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
+	"github.com/milvus-io/milvus/pkg/v3/proto/viewpb"
+)
+
+func TestSnapshotBuilder_LoadConfigsFollowResolvedScope(t *testing.T) {
+	catalog := mocks.NewQueryCoordCatalog(t)
+	catalog.EXPECT().GetCollections(mock.Anything).Return([]*querypb.CollectionLoadInfo{
+		{CollectionID: 1}, {CollectionID: 2},
+	}, nil).Once()
+	catalog.EXPECT().GetPartitions(mock.Anything, mock.Anything).Return(nil, nil).Once()
+	catalog.EXPECT().GetReplicas(mock.Anything).Return([]*querypb.Replica{
+		{ID: 10, CollectionID: 1}, {ID: 20, CollectionID: 2},
+	}, nil).Once()
+	store, err := loadmgr.RecoverLoadConfigStore(context.Background(), catalog)
+	require.NoError(t, err)
+	registry := triggerTestRegistry(t)
+	first, second, residual := triggerShard(10, 1, 0), triggerShard(20, 2, 0), triggerShard(30, 3, 0)
+	addShardWithPreparingView(t, registry, first, map[int64]map[int64][]int64{100: {1: {101}}})
+	addShardWithPreparingView(t, registry, second, map[int64]map[int64][]int64{200: {2: {201}}})
+	registry.Ensure(residual)
+	malformed := qviews.ShardID{ReplicaID: 40, VChannel: "malformed"}
+	addShardWithPreparingView(t, registry, malformed, map[int64]map[int64][]int64{300: {4: {401}}})
+	allShards := []qviews.ShardID{first, second, residual, malformed}
+
+	for _, test := range []struct {
+		name    string
+		batch   triggerBatch
+		configs []int64
+		shards  []qviews.ShardID
+	}{
+		{"collection", triggerBatch{dirtyColls: setOf[int64](1)}, []int64{1}, []qviews.ShardID{first}},
+		{"shard", triggerBatch{dirtyShards: setOf(second)}, []int64{2}, []qviews.ShardID{second}},
+		{"node", triggerBatch{dirtyNodes: setOf[int64](100)}, []int64{1}, []qviews.ShardID{first}},
+		{"released", triggerBatch{dirtyColls: setOf[int64](3)}, nil, []qviews.ShardID{residual}},
+		{"empty", triggerBatch{}, nil, nil},
+		{"full", triggerBatch{full: true}, []int64{1, 2}, allShards},
+		{"malformed shard", triggerBatch{dirtyShards: setOf(malformed)}, []int64{1, 2}, allShards},
+		{"malformed node shard", triggerBatch{dirtyNodes: setOf[int64](300)}, []int64{1, 2}, allShards},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			provider := &fakeDataViewProvider{collections: []*viewpb.DataViewOfCollection{
+				{CollectionId: 1, DataVersion: &viewpb.DataVersion{}, Shards: []*viewpb.DataViewOfShard{{Vchannel: first.VChannel}}},
+				{CollectionId: 2, DataVersion: &viewpb.DataVersion{}, Shards: []*viewpb.DataViewOfShard{{Vchannel: second.VChannel}}},
+			}}
+			builder := NewSnapshotBuilder(store, registry, &fakeNodeProvider{}, provider, policyTestConfig())
+			snapshot, targets := builder.build(context.Background(), test.batch)
+			assert.Len(t, snapshot.ConfigsMap(), len(test.configs))
+			for _, id := range test.configs {
+				require.Contains(t, snapshot.ConfigsMap(), id)
+				assert.Equal(t, store.Get(id).ConfigVersion, snapshot.LoadConfigSnapshot.ConfigVersion(id))
+				assert.Same(t, snapshot.ConfigsMap()[id], snapshot.ConfigForShard(triggerShard(id*10, id, 0)))
+			}
+			assert.ElementsMatch(t, test.shards, targets)
+			assert.Nil(t, snapshot.ConfigForShard(residual))
+			assert.Zero(t, snapshot.LoadConfigSnapshot.ConfigVersion(3))
+		})
+	}
+}

internal/views/coord/balancer/snapshot_builder.go

@@ -110,10 +110,15 @@ func (b *SnapshotBuilder) ObserveShardStats(shardID qviews.ShardID, _ *coordview
 // refreshes the incremental row-count ledger, and returns the exact shard list
 // that BalancePolicy should plan in this cycle.
 func (b *SnapshotBuilder) build(ctx context.Context, pending triggerBatch) (*BalancerSnapshot, []qviews.ShardID) {
-	// 1. Capture load configs and resolve the preliminary trigger scope.
-	loadSnapshot := b.configStore.Snapshot()
-
-	scope := pending.resolveScope(loadSnapshot, b.viewRegistry)
+	// 1. Resolve the trigger scope before reading load configs.
+	scope := pending.resolveScope(b.viewRegistry)
+	var loadSnapshot *loadmgr.LoadConfigSnapshot
+	if scope.full {
+		loadSnapshot = b.configStore.Snapshot()
+		scope = fullReconcileScope(loadSnapshot, b.viewRegistry)
+	} else {
+		loadSnapshot = b.configStore.SnapshotForCollections(maps.Keys(scope.collectionIDs))
+	}
 
 	// 2. Read scoped DataViews and expand collection triggers into target shards.
 	dataViewSnapshot := b.dataViewProvider.DataViewSnapshotForCollections(ctx, scope.collectionIDs)