diff --git a/docs/design-docs/design_docs/qviews/balancer_design.md b/docs/design-docs/design_docs/qviews/balancer_design.md index a40a5e95298..44917e954af 100644 --- a/docs/design-docs/design_docs/qviews/balancer_design.md +++ b/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) diff --git a/internal/querycoordv2/ddl_callbacks_alter_load_info_load_collection.go b/internal/querycoordv2/ddl_callbacks_alter_load_info_load_collection.go index a80f2ee6994..81d61d6b1a0 100644 --- a/internal/querycoordv2/ddl_callbacks_alter_load_info_load_collection.go +++ b/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 { diff --git a/internal/querycoordv2/scoped_load_reads_test.go b/internal/querycoordv2/scoped_load_reads_test.go new file mode 100644 index 00000000000..f3b5eb1dfa7 --- /dev/null +++ b/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) +} diff --git a/internal/querycoordv2/services.go b/internal/querycoordv2/services.go index cea8f629e43..54f69ce3187 100644 --- a/internal/querycoordv2/services.go +++ b/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 diff --git a/internal/views/coord/balancer/scoped_load_config_test.go b/internal/views/coord/balancer/scoped_load_config_test.go new file mode 100644 index 00000000000..18229f870f1 --- /dev/null +++ b/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)) + }) + } +} diff --git a/internal/views/coord/balancer/snapshot_builder.go b/internal/views/coord/balancer/snapshot_builder.go index e07a3f517e1..2ed4a574af2 100644 --- a/internal/views/coord/balancer/snapshot_builder.go +++ b/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) diff --git a/internal/views/coord/balancer/trigger.go b/internal/views/coord/balancer/trigger.go index f1f682d7590..62fe920bf0c 100644 --- a/internal/views/coord/balancer/trigger.go +++ b/internal/views/coord/balancer/trigger.go @@ -43,6 +43,9 @@ type triggerBatch struct { // reconcileScope is the scoped DataView-read and Policy-planning boundary // resolved from one trigger batch. type reconcileScope struct { + // full requires expansion from all configured collections and resident shards. + full bool + // collectionIDs selects collections whose DataViews are fetched. collectionIDs map[int64]struct{} @@ -60,13 +63,13 @@ func (b triggerBatch) empty() bool { // resolveScope converts queued collection and shard events into scoped reads // and Policy targets. Collection events include resident residual shards; -// malformed shard channels conservatively fall back to a full reconcile. +// malformed shard channels request a full scope, expanded after load configs +// have been captured by the snapshot builder. func (b triggerBatch) resolveScope( - loadSnapshot *loadmgr.LoadConfigSnapshot, registry *coordview.ShardViewRegistry, ) reconcileScope { if b.full { - return fullReconcileScope(loadSnapshot, registry) + return reconcileScope{full: true} } scope := newReconcileScope() @@ -80,12 +83,12 @@ func (b triggerBatch) resolveScope( for nodeID := range b.dirtyNodes { if registry == nil { - return fullReconcileScope(loadSnapshot, registry) + return reconcileScope{full: true} } for _, shardID := range registry.NodeShards(nodeID) { collectionID, ok := parseShardCollection(shardID) if !ok { - return fullReconcileScope(loadSnapshot, registry) + return reconcileScope{full: true} } scope.collectionIDs[collectionID] = struct{}{} scope.targetShards[shardID] = struct{}{} @@ -95,7 +98,7 @@ func (b triggerBatch) resolveScope( for shardID := range b.dirtyShards { collectionID, ok := parseShardCollection(shardID) if !ok { - return fullReconcileScope(loadSnapshot, registry) + return reconcileScope{full: true} } scope.collectionIDs[collectionID] = struct{}{} scope.targetShards[shardID] = struct{}{} @@ -148,6 +151,7 @@ func fullReconcileScope( registry *coordview.ShardViewRegistry, ) reconcileScope { scope := newReconcileScope() + scope.full = true if loadSnapshot != nil { for collectionID := range loadSnapshot.ConfigsMap() { scope.collectionIDs[collectionID] = struct{}{} diff --git a/internal/views/coord/balancer/trigger_test.go b/internal/views/coord/balancer/trigger_test.go index 2ba6bcf4b0a..e44df0af467 100644 --- a/internal/views/coord/balancer/trigger_test.go +++ b/internal/views/coord/balancer/trigger_test.go @@ -21,7 +21,7 @@ func TestTriggerBatchResolveDirtyCollection(t *testing.T) { registry.Ensure(shardB) registry.Ensure(unrelated) - scope := (triggerBatch{dirtyColls: setOf[int64](1)}).resolveScope(nil, registry) + scope := (triggerBatch{dirtyColls: setOf[int64](1)}).resolveScope(registry) assert.Equal(t, setOf[int64](1), scope.collectionIDs) assert.Equal(t, setOf(shardB, shardA), scope.targetShards) @@ -29,7 +29,7 @@ func TestTriggerBatchResolveDirtyCollection(t *testing.T) { func TestTriggerBatchResolveDirtyShard(t *testing.T) { shard := triggerShard(10, 2, 0) - scope := (triggerBatch{dirtyShards: setOf(shard)}).resolveScope(nil, nil) + scope := (triggerBatch{dirtyShards: setOf(shard)}).resolveScope(nil) assert.Equal(t, setOf[int64](2), scope.collectionIDs) assert.Equal(t, setOf(shard), scope.targetShards) @@ -50,7 +50,7 @@ func TestTriggerBatchResolveDirtyNode(t *testing.T) { 200: {30: {301}}, }) - scope := (triggerBatch{dirtyNodes: setOf[int64](100)}).resolveScope(nil, registry) + scope := (triggerBatch{dirtyNodes: setOf[int64](100)}).resolveScope(registry) assert.Equal(t, setOf[int64](1, 2), scope.collectionIDs) assert.Empty(t, scope.collectionWideIDs) @@ -69,7 +69,9 @@ func TestTriggerBatchMalformedNodeShardFallsBackToFull(t *testing.T) { 100: {10: {101}}, }) - scope := (triggerBatch{dirtyNodes: setOf[int64](100)}).resolveScope(loadSnapshot, registry) + scope := (triggerBatch{dirtyNodes: setOf[int64](100)}).resolveScope(registry) + assert.True(t, scope.full) + scope = fullReconcileScope(loadSnapshot, registry) assert.Equal(t, setOf[int64](1, 2), scope.collectionIDs) assert.Equal(t, setOf(registry.ShardIDs()...), scope.targetShards) @@ -84,7 +86,7 @@ func TestTriggerBatchResolveMergesScopes(t *testing.T) { scope := (triggerBatch{ dirtyColls: setOf[int64](1), dirtyShards: setOf(dirtyShard), - }).resolveScope(nil, registry) + }).resolveScope(registry) assert.Equal(t, setOf[int64](1, 2), scope.collectionIDs) assert.Equal(t, setOf(collectionShard, dirtyShard), scope.targetShards) @@ -105,7 +107,9 @@ func TestTriggerBatchResolveFull(t *testing.T) { queue := newTriggerQueue() queue.add(TriggerScope{NodeChanged: true}) - scope := queue.takePending().resolveScope(loadSnapshot, registry) + scope := queue.takePending().resolveScope(registry) + assert.True(t, scope.full) + scope = fullReconcileScope(loadSnapshot, registry) assert.Equal(t, setOf[int64](1, 2, 3), scope.collectionIDs) assert.Equal(t, setOf(malformedShard, shardA, residualShard), scope.targetShards) @@ -120,7 +124,9 @@ func TestTriggerBatchMalformedDirtyShardFallsBackToFull(t *testing.T) { registry.Ensure(validShard) registry.Ensure(malformedShard) - scope := (triggerBatch{dirtyShards: setOf(malformedShard)}).resolveScope(loadSnapshot, registry) + scope := (triggerBatch{dirtyShards: setOf(malformedShard)}).resolveScope(registry) + assert.True(t, scope.full) + scope = fullReconcileScope(loadSnapshot, registry) assert.Equal(t, setOf[int64](1, 2), scope.collectionIDs) assert.Equal(t, setOf(registry.ShardIDs()...), scope.targetShards) @@ -133,10 +139,7 @@ func TestTriggerBatchReleasedCollectionIncludesResidualShards(t *testing.T) { registry.Ensure(residualShard) registry.Ensure(unrelatedShard) - scope := (triggerBatch{dirtyColls: setOf[int64](1)}).resolveScope( - triggerLoadSnapshot(triggerLoadConfig(2, 20)), - registry, - ) + scope := (triggerBatch{dirtyColls: setOf[int64](1)}).resolveScope(registry) assert.Equal(t, setOf[int64](1), scope.collectionIDs) assert.Equal(t, setOf(residualShard), scope.targetShards) @@ -146,9 +149,7 @@ func TestTriggerBatchDirtyCollectionUsesCollectionIndexOnly(t *testing.T) { registry := triggerTestRegistry(t) malformedResidual := qviews.ShardID{ReplicaID: 20, VChannel: "malformed"} registry.Ensure(malformedResidual) - loadSnapshot := triggerLoadSnapshot(triggerLoadConfig(1, 10)) - - scope := (triggerBatch{dirtyColls: setOf[int64](99)}).resolveScope(loadSnapshot, registry) + scope := (triggerBatch{dirtyColls: setOf[int64](99)}).resolveScope(registry) assert.Equal(t, setOf[int64](99), scope.collectionIDs) assert.Empty(t, scope.targetShards) @@ -178,7 +179,7 @@ func TestReconcileScopeAddDataViewShards(t *testing.T) { Shards: []*viewpb.DataViewOfShard{{Vchannel: "by-dev-rootcoord-dml_0_2v0"}}, }, }, nil) - scope := (triggerBatch{dirtyColls: setOf[int64](1)}).resolveScope(loadSnapshot, nil) + scope := (triggerBatch{dirtyColls: setOf[int64](1)}).resolveScope(nil) scope.AddDataViewShards(loadSnapshot, dataSnapshot) @@ -210,7 +211,7 @@ func TestReconcileScopeDoesNotExpandDirtyShard(t *testing.T) { }, nil) dirtyShard := triggerShard(10, 1, 0) - scope := (triggerBatch{dirtyShards: setOf(dirtyShard)}).resolveScope(loadSnapshot, nil) + scope := (triggerBatch{dirtyShards: setOf(dirtyShard)}).resolveScope(nil) scope.AddDataViewShards(loadSnapshot, dataSnapshot) diff --git a/internal/views/coord/coordview/shard_view_collection_snapshot_test.go b/internal/views/coord/coordview/shard_view_collection_snapshot_test.go new file mode 100644 index 00000000000..e7de517aeb3 --- /dev/null +++ b/internal/views/coord/coordview/shard_view_collection_snapshot_test.go @@ -0,0 +1,94 @@ +// 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 coordview + +import ( + "fmt" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/milvus-io/milvus/internal/views/qviews" +) + +func TestRegistry_SnapshotForCollection(t *testing.T) { + reg := newTestRegistry(t, newMockCatalog(), newMockSyncer()) + first := qviews.ShardID{ReplicaID: 1, VChannel: "by-dev-rootcoord-dml_100v0"} + second := qviews.ShardID{ReplicaID: 2, VChannel: "by-dev-rootcoord-dml_100v1"} + other := qviews.ShardID{ReplicaID: 3, VChannel: "by-dev-rootcoord-dml_200v0"} + for _, shard := range []qviews.ShardID{first, second, other} { + reg.Ensure(shard) + } + resident := reg.Snapshot() + stats := shardStatsForNodes(map[int64][]int64{101: {1}}) + reg.onShardStatsChanged(first, stats) + scoped := reg.SnapshotForCollection(100) + assert.Same(t, resident, reg.snapshot) + assert.Equal(t, reg.version, scoped.Version()) + require.Len(t, scoped.StatsMap(), 2) + assert.Same(t, stats, scoped.StatsMap()[first]) + assert.Contains(t, scoped.StatsMap(), second, "include resident shards that have no Up view") + assert.NotContains(t, scoped.StatsMap(), other) + assert.Empty(t, reg.SnapshotForCollection(300).StatsMap()) + + reg.onShardStatsChanged(first, emptyShardStats()) + assert.Same(t, stats, scoped.StatsMap()[first], "later updates must not change a captured snapshot") + delete(scoped.StatsMap(), second) + assert.Contains(t, reg.SnapshotForCollection(100).StatsMap(), second, "scoped snapshots own their outer maps") +} + +func BenchmarkCollectionShardReads(b *testing.B) { + for _, count := range []int{10000, 50000, 150000} { + reg := &ShardViewRegistry{ + version: 1, + stats: make(map[qviews.ShardID]*ShardStats, count), + collectionShards: make(map[int64]map[qviews.ShardID]struct{}, count), + } + version := qviews.QueryViewVersion{} + for id := int64(1); id <= int64(count); id++ { + shard := qviews.ShardID{ReplicaID: id, VChannel: fmt.Sprintf("by-dev-rootcoord-dml_%dv0", id)} + reg.stats[shard] = &ShardStats{UpVersion: &version} + reg.collectionShards[id] = map[qviews.ShardID]struct{}{shard: {}} + } + for _, mode := range []string{"full", "scoped"} { + b.Run(fmt.Sprintf("%d/%s", count, mode), func(b *testing.B) { + b.ReportAllocs() + for b.Loop() { + reg.mu.Lock() + reg.version++ + reg.mu.Unlock() + var snapshot *ShardViewSnapshot + if mode == "full" { + snapshot = reg.Snapshot() + } else { + snapshot = reg.SnapshotForCollection(1) + } + loaded := 0 + for shard, stats := range snapshot.StatsMap() { + if shard.ReplicaID == 1 && stats.UpVersion != nil { + loaded++ + } + } + if loaded != 1 { + b.Fatalf("expected one loaded shard, got %d", loaded) + } + } + }) + } + } +} diff --git a/internal/views/coord/coordview/shard_view_registry.go b/internal/views/coord/coordview/shard_view_registry.go index 70889c9f6fc..a3eb60b068f 100644 --- a/internal/views/coord/coordview/shard_view_registry.go +++ b/internal/views/coord/coordview/shard_view_registry.go @@ -231,6 +231,22 @@ func (r *ShardViewRegistry) SnapshotForShards(shardIDs []qviews.ShardID) *ShardV } } +// SnapshotForCollection captures the collection index and its current stats +// under one lock, without refreshing the cached full snapshot. +func (r *ShardViewRegistry) SnapshotForCollection(collectionID int64) *ShardViewSnapshot { + r.mu.RLock() + defer r.mu.RUnlock() + + shards := r.collectionShards[collectionID] + stats := make(map[qviews.ShardID]*ShardStats, len(shards)) + for shardID := range shards { + if shardStats, ok := r.stats[shardID]; ok { + stats[shardID] = shardStats + } + } + return &ShardViewSnapshot{version: r.version, stats: stats} +} + // CollectionShards returns the resident shards belonging to collectionID. func (r *ShardViewRegistry) CollectionShards(collectionID int64) []qviews.ShardID { r.mu.RLock() diff --git a/internal/views/coord/coordview/shard_view_registry_test.go b/internal/views/coord/coordview/shard_view_registry_test.go index 2828a3f0e0b..cb29bdcff6c 100644 --- a/internal/views/coord/coordview/shard_view_registry_test.go +++ b/internal/views/coord/coordview/shard_view_registry_test.go @@ -95,6 +95,8 @@ func TestRegistry_RemovesManagerAfterLastViewDropped(t *testing.T) { require.NoError(t, mgr.AddPreparing(context.Background(), builder)) require.NoError(t, reg.flushScheduler.Flush(context.Background())) + scopedBeforeDrop := reg.SnapshotForCollection(100) + require.Contains(t, scopedBeforeDrop.StatsMap(), shardID) require.NoError(t, mgr.RequestRelease(context.Background())) require.NoError(t, reg.flushScheduler.Flush(context.Background())) @@ -114,9 +116,12 @@ func TestRegistry_RemovesManagerAfterLastViewDropped(t *testing.T) { assert.Nil(t, reg.Get(shardID)) assert.Empty(t, reg.CollectionShards(100)) + assert.Empty(t, reg.SnapshotForCollection(100).StatsMap()) + assert.Contains(t, scopedBeforeDrop.StatsMap(), shardID) assert.Empty(t, reg.ShardIDs()) assert.NotContains(t, reg.Snapshot().StatsMap(), shardID) assert.NotSame(t, mgr, reg.Ensure(shardID)) + assert.Contains(t, reg.SnapshotForCollection(100).StatsMap(), shardID) } func TestRegistry_RecoverWithPersistedViews(t *testing.T) { diff --git a/internal/views/coord/loadmgr/load_config_reads_test.go b/internal/views/coord/loadmgr/load_config_reads_test.go new file mode 100644 index 00000000000..8694ad33668 --- /dev/null +++ b/internal/views/coord/loadmgr/load_config_reads_test.go @@ -0,0 +1,144 @@ +// 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 loadmgr + +import ( + "context" + "fmt" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/require" + + "github.com/milvus-io/milvus/pkg/v3/proto/messagespb" +) + +func TestScopedConfigReadsPreserveVersionsAndSnapshots(t *testing.T) { + ctx := context.Background() + store, catalog := newTestStore(t) + resident := store.Snapshot() + expectFullSave(catalog, 4) + first := sampleConfig() + require.NoError(t, store.Put(ctx, first)) + other := sampleConfig() + other.CollectionID++ + other.Replicas = []*ReplicaAssignment{{ReplicaID: 2000}} + require.NoError(t, store.Put(ctx, other)) + + entry := store.Get(first.CollectionID) + require.NotNil(t, entry.Config) + assert.Equal(t, uint64(2), entry.ConfigVersion) + assert.Equal(t, uint64(3), entry.StoreVersion) + scoped := store.SnapshotForCollections([]int64{first.CollectionID, -1, first.CollectionID}) + assert.Same(t, resident, store.snapshot, "point and scoped reads must not rebuild the full snapshot") + assert.Equal(t, entry.StoreVersion, scoped.Version()) + assert.Equal(t, entry.ConfigVersion, scoped.ConfigVersion(first.CollectionID)) + assert.Equal(t, uint64(0), scoped.ConfigVersion(-1)) + require.Len(t, scoped.ConfigsMap(), 1) + assert.Same(t, entry.Config, scoped.ConfigsMap()[first.CollectionID]) + assert.Len(t, scoped.ReplicaToConfigMap(), 2) + assert.NotContains(t, scoped.ReplicaToConfigMap(), int64(2000)) + assert.Empty(t, store.SnapshotForCollections(nil).ConfigsMap()) + + updated := first.Clone() + updated.LoadFields[0].IndexId++ + require.NoError(t, store.Put(ctx, updated)) + assert.Equal(t, first.LoadFields[0].IndexId, entry.Config.LoadFields[0].IndexId) + assert.Equal(t, first.LoadFields[0].IndexId, scoped.ConfigsMap()[first.CollectionID].LoadFields[0].IndexId) + assert.Equal(t, uint64(4), store.Get(first.CollectionID).ConfigVersion) + + catalog.EXPECT().ReleaseReplicas(mock.Anything, first.CollectionID).Return(nil).Once() + catalog.EXPECT().ReleaseCollection(mock.Anything, first.CollectionID).Return(nil).Once() + require.NoError(t, store.Remove(ctx, first.CollectionID)) + absent := store.Get(first.CollectionID) + assert.Nil(t, absent.Config) + assert.Zero(t, absent.ConfigVersion) + assert.Equal(t, uint64(5), absent.StoreVersion) + assert.Empty(t, store.SnapshotForCollections([]int64{first.CollectionID}).ConfigsMap()) + + require.NoError(t, store.Put(ctx, updated)) + assert.Equal(t, uint64(6), store.Get(first.CollectionID).ConfigVersion) + assert.NotEqual(t, entry.ConfigVersion, store.Get(first.CollectionID).ConfigVersion) + assert.NotEqual(t, absent.ConfigVersion, store.Get(first.CollectionID).ConfigVersion) + assert.Same(t, resident, store.snapshot) +} + +func TestScopedConfigReadsDuringPut(t *testing.T) { + store, catalog := newTestStore(t) + const updates = 200 + catalog.EXPECT().SaveCollection(mock.Anything, mock.Anything).Return(nil).Times(updates + 1) + put := func(revision int64) error { + return store.Put(context.Background(), &LoadConfig{ + CollectionID: 1, + LoadFields: []*messagespb.LoadFieldConfig{{FieldId: 100, IndexId: revision}}, + }) + } + require.NoError(t, put(1)) + done := make(chan struct{}) + go func() { + defer close(done) + for revision := int64(2); revision <= updates+1; revision++ { + assert.NoError(t, put(revision)) + } + }() + for range updates { + entry := store.Get(1) + assert.Equal(t, uint64(entry.Config.LoadFields[0].IndexId+1), entry.ConfigVersion) + assert.Equal(t, entry.ConfigVersion, entry.StoreVersion) + snapshot := store.SnapshotForCollections([]int64{1}) + assert.Equal(t, uint64(snapshot.ConfigsMap()[1].LoadFields[0].IndexId+1), snapshot.ConfigVersion(1)) + assert.Equal(t, snapshot.ConfigVersion(1), snapshot.Version()) + } + <-done +} + +// Compare hot-path reads after a config update invalidates the full snapshot. +// Persistence is excluded so this measures the in-memory work that scales with +// the number of loaded collections. +func BenchmarkLoadConfigReads(b *testing.B) { + for _, count := range []int{10000, 50000, 150000} { + store := &LoadConfigStore{ + version: 1, + configs: make(map[int64]*LoadConfig, count), + versions: make(map[int64]uint64, count), + } + for id := int64(1); id <= int64(count); id++ { + store.configs[id] = &LoadConfig{CollectionID: id, Replicas: []*ReplicaAssignment{{ReplicaID: id}}} + store.versions[id] = 1 + } + for _, mode := range []string{"full", "point", "scoped"} { + b.Run(fmt.Sprintf("%d/%s", count, mode), func(b *testing.B) { + b.ReportAllocs() + for b.Loop() { + store.mu.Lock() + store.version++ + store.versions[1] = store.version + store.mu.Unlock() + switch mode { + case "full": + store.Snapshot() + case "point": + store.Get(1) + case "scoped": + store.SnapshotForCollections([]int64{1}) + } + } + }) + } + } +} diff --git a/internal/views/coord/loadmgr/load_config_store.go b/internal/views/coord/loadmgr/load_config_store.go index e1bd612bd03..6a845ee2399 100644 --- a/internal/views/coord/loadmgr/load_config_store.go +++ b/internal/views/coord/loadmgr/load_config_store.go @@ -17,7 +17,7 @@ import ( // // # Copy-On-Write semantics // -// Snapshot returns pointers into the store's copy-on-write state for zero-copy +// Reads return pointers into the store's copy-on-write state for zero-copy // efficiency in this read-heavy path. Callers MUST treat the returned // LoadConfig / ReplicaAssignment values as read-only. // @@ -41,6 +41,14 @@ type LoadConfigStore struct { snapshot *LoadConfigSnapshot } +// LoadConfigEntry captures one immutable config and its versions in one read. +// Config is nil and ConfigVersion is zero when the collection is absent. +type LoadConfigEntry struct { + Config *LoadConfig + ConfigVersion uint64 + StoreVersion uint64 +} + // RecoverLoadConfigStore constructs a LoadConfigStore and rebuilds its // in-memory state from ETCD via the catalog. It is the sole constructor: // the store is always fully recovered before any operation. @@ -184,6 +192,43 @@ func (s *LoadConfigStore) Contains(collectionID int64) bool { return s.configs[collectionID] != nil } +// Get reads one collection without materializing the full snapshot. +func (s *LoadConfigStore) Get(collectionID int64) LoadConfigEntry { + s.mu.RLock() + defer s.mu.RUnlock() + return LoadConfigEntry{ + Config: s.configs[collectionID], + ConfigVersion: s.versions[collectionID], + StoreVersion: s.version, + } +} + +// SnapshotForCollections captures only the requested collections. An empty list +// selects none. It does not refresh the cached full snapshot. +func (s *LoadConfigStore) SnapshotForCollections(collectionIDs []int64) *LoadConfigSnapshot { + snapshot := &LoadConfigSnapshot{ + configs: make(map[int64]*LoadConfig, len(collectionIDs)), + configVersions: make(map[int64]uint64, len(collectionIDs)), + replicaToConfig: make(map[int64]*LoadConfig), + } + s.mu.RLock() + snapshot.version = s.version + for _, collectionID := range collectionIDs { + if cfg := s.configs[collectionID]; cfg != nil { + snapshot.configs[collectionID] = cfg + snapshot.configVersions[collectionID] = s.versions[collectionID] + } + } + s.mu.RUnlock() + // Configs are immutable; build the replica index outside the store lock. + for _, cfg := range snapshot.configs { + for _, replica := range cfg.Replicas { + snapshot.replicaToConfig[replica.ReplicaID] = cfg + } + } + return snapshot +} + // Snapshot returns the current immutable load-config view. It refreshes the // resident snapshot lazily when the live version has advanced. The returned // maps point at the store's copy-on-write snapshots and must be treated as