diff --git a/internal/datacoord/compaction_policy_allocation_test.go b/internal/datacoord/compaction_policy_allocation_test.go new file mode 100644 index 00000000000..893d23d230c --- /dev/null +++ b/internal/datacoord/compaction_policy_allocation_test.go @@ -0,0 +1,111 @@ +// 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 datacoord + +import ( + "context" + "strconv" + "testing" + + "github.com/blang/semver/v4" + + "github.com/milvus-io/milvus-proto/go-api/v3/schemapb" + "github.com/milvus-io/milvus/internal/storage" + "github.com/milvus-io/milvus/pkg/v3/util/paramtable" + "github.com/milvus-io/milvus/pkg/v3/util/typeutil" +) + +type compactionCheckBenchmarkAllocator struct{} + +func (*compactionCheckBenchmarkAllocator) AllocID(context.Context) (int64, error) { + return 1, nil +} + +func (*compactionCheckBenchmarkAllocator) AllocTimestamp(context.Context) (uint64, error) { + return 1, nil +} + +func (*compactionCheckBenchmarkAllocator) AllocN(n int64) (int64, int64, error) { + return 1, n + 1, nil +} + +// Benchmark the steady-state full scan: every collection has one flushed V3 +// segment whose schema/storage versions already match the desired versions. +// The allocator is local, so its real RPC allocation cost is intentionally not +// included in this benchmark. +func BenchmarkCompactionPolicyNoCandidates(b *testing.B) { + paramtable.Init() + params := paramtable.Get() + for _, setting := range []struct { + param *paramtable.ParamItem + value string + }{ + {¶ms.CommonCfg.UseLoonFFI, "true"}, + {¶ms.DataCoordCfg.StorageVersionCompactionEnabled, "true"}, + {¶ms.DataCoordCfg.StorageFormatCompactionEnabled, "false"}, + } { + key, previous := setting.param.Key, setting.param.GetValue() + params.Save(key, setting.value) + b.Cleanup(func() { params.Save(key, previous) }) + } + + for _, policyName := range []string{"schema", "storage"} { + for _, collectionCount := range []int{1000, 10000} { + b.Run(policyName+"/"+strconv.Itoa(collectionCount), func(b *testing.B) { + m := &meta{ + ctx: context.Background(), + collections: typeutil.NewConcurrentMap[int64, *collectionInfo](), + segments: NewCachedSegmentsInfo(), + } + schema := newBumpSchemaVersionTestCollection(1, 2).Schema + for fieldID := int64(102); fieldID < 116; fieldID++ { + schema.Fields = append(schema.Fields, &schemapb.FieldSchema{ + FieldID: fieldID, Name: "field_" + strconv.FormatInt(fieldID, 10), DataType: schemapb.DataType_Int64, + }) + } + for id := int64(1); id <= int64(collectionCount); id++ { + m.collections.Insert(id, &collectionInfo{ID: id, Schema: schema}) + m.segments.SetSegment(id, newBumpSchemaVersionTestSegment(id, id, 2, storage.StorageV3, "manifest"), 1) + } + alloc := &compactionCheckBenchmarkAllocator{} + handler := &ServerHandler{s: &Server{meta: m}} + var policy CompactionPolicy + if policyName == "schema" { + policy = newBumpSchemaVersionPolicy(m, alloc, handler) + } else { + versionManager := NewMockVersionManager(b) + versionManager.EXPECT().GetMinimalSessionVer().Return(semver.MustParse("3.0.0")) + policy = newStorageVersionUpgradePolicy(m, alloc, handler, versionManager) + } + ctx := context.Background() + b.ReportAllocs() + b.ResetTimer() + for b.Loop() { + events, err := policy.Trigger(ctx) + if err != nil { + b.Fatal(err) + } + for _, views := range events { + if len(views) != 0 { + b.Fatal("up-to-date segments must not produce compaction views") + } + } + } + }) + } + } +} diff --git a/internal/datacoord/compaction_policy_bump_schema_version.go b/internal/datacoord/compaction_policy_bump_schema_version.go index bacaf3099a6..7e80f133961 100644 --- a/internal/datacoord/compaction_policy_bump_schema_version.go +++ b/internal/datacoord/compaction_policy_bump_schema_version.go @@ -102,9 +102,15 @@ func (policy *bumpSchemaVersionPolicy) Trigger(ctx context.Context) (map[Compact continue } collectionID := collection.ID - capturedSchema := proto.Clone(collection.Schema).(*schemapb.CollectionSchema) - collectionSchemaVersion := capturedSchema.GetVersion() + // The collection cache publishes replacements. Keep this schema reference + // so candidate selection and the task snapshot use the same version. + schema := collection.Schema + collectionSchemaVersion := schema.GetVersion() partSegments := policy.staleFlushedSegments(collectionID, collectionSchemaVersion) + if len(partSegments) == 0 { + continue + } + capturedSchema := proto.Clone(schema).(*schemapb.CollectionSchema) var views []CompactionView var collectionTriggerID int64 diff --git a/internal/datacoord/compaction_policy_bump_schema_version_test.go b/internal/datacoord/compaction_policy_bump_schema_version_test.go index 9e6e93c116b..e15ba79379c 100644 --- a/internal/datacoord/compaction_policy_bump_schema_version_test.go +++ b/internal/datacoord/compaction_policy_bump_schema_version_test.go @@ -18,12 +18,15 @@ package datacoord import ( "context" + "fmt" "testing" "time" + "github.com/bytedance/mockey" "github.com/cockroachdb/errors" "github.com/stretchr/testify/mock" "github.com/stretchr/testify/suite" + "google.golang.org/protobuf/proto" "github.com/milvus-io/milvus-proto/go-api/v3/commonpb" "github.com/milvus-io/milvus-proto/go-api/v3/msgpb" @@ -218,6 +221,79 @@ func (s *BumpSchemaVersionPolicySuite) TestTriggerCapturesSchemaSnapshot() { s.Len(view.schema.GetFields(), 2) } +func (s *BumpSchemaVersionPolicySuite) TestTriggerClonesSchemaOnlyForCandidates() { + for _, stale := range []bool{false, true} { + s.Run(fmt.Sprintf("stale=%t", stale), func() { + collection := newBumpSchemaVersionTestCollection(100, 2) + mockAlloc := newMockAllocator(s.T()) + policy := newBumpSchemaVersionPolicy(&meta{ + collections: typeutil.NewConcurrentMap[UniqueID, *collectionInfo](), + segments: NewCachedSegmentsInfo(), + }, mockAlloc, s.handler) + policy.meta.collections.Insert(collection.ID, collection) + segmentSchemaVersion := int32(2) + if stale { + segmentSchemaVersion = 1 + } + for _, segmentID := range []int64{101, 102} { + policy.meta.segments.SetSegment(segmentID, newBumpSchemaVersionTestSegment(collection.ID, segmentID, segmentSchemaVersion, storage.StorageV3, "manifest"), 0) + } + + cloneCount := 0 + var clone func(proto.Message) proto.Message + cloneMock := mockey.Mock(proto.Clone).To(func(message proto.Message) proto.Message { + if message == collection.Schema { + cloneCount++ + } + return clone(message) + }).Origin(&clone).Build() + defer cloneMock.UnPatch() + + events, err := policy.Trigger(context.Background()) + s.NoError(err) + views := events[TriggerTypeBumpSchemaVersion] + if !stale { + s.Empty(views) + s.Zero(cloneCount) + mockAlloc.AssertNotCalled(s.T(), "AllocID", mock.Anything) + return + } + s.Require().Len(views, 2) + s.Equal(1, cloneCount) + s.Same(views[0].(*BumpSchemaVersionView).schema, views[1].(*BumpSchemaVersionView).schema) + mockAlloc.AssertNumberOfCalls(s.T(), "AllocID", 1) + }) + } +} + +func (s *BumpSchemaVersionPolicySuite) TestTriggerRetainsSchemaAcrossCacheReplacement() { + collection := newBumpSchemaVersionTestCollection(100, 2) + policy := s.bumpSchemaVersionPolicy + policy.meta.collections.Insert(collection.ID, collection) + policy.meta.segments.SetSegment(101, newBumpSchemaVersionTestSegment(collection.ID, 101, 1, storage.StorageV3, "manifest"), 0) + // This segment only becomes stale at V3, after this scan's V2 snapshot. + policy.meta.segments.SetSegment(102, newBumpSchemaVersionTestSegment(collection.ID, 102, 2, storage.StorageV3, "manifest"), 0) + + var selectSegments func(*bumpSchemaVersionPolicy, int64, int32) []*chanPartSegments + selectMock := mockey.Mock((*bumpSchemaVersionPolicy).staleFlushedSegments).To( + func(policy *bumpSchemaVersionPolicy, collectionID int64, version int32) []*chanPartSegments { + segments := selectSegments(policy, collectionID, version) + policy.meta.collections.Insert(collectionID, newBumpSchemaVersionTestCollection(collectionID, 3)) + return segments + }).Origin(&selectSegments).Build() + defer selectMock.UnPatch() + + events, err := policy.Trigger(context.Background()) + s.NoError(err) + views := events[TriggerTypeBumpSchemaVersion] + s.Require().Len(views, 1) + view := views[0].(*BumpSchemaVersionView) + s.EqualValues(101, view.segments[0].ID) + s.EqualValues(2, view.schema.GetVersion()) + s.NotSame(collection.Schema, view.schema) + s.EqualValues(3, policy.meta.GetCollection(collection.ID).Schema.GetVersion()) +} + func (s *BumpSchemaVersionPolicySuite) TestTriggerSchedulesReadySegmentWhenCollectionHasMissingManifestFlushedDataSegment() { ctx := context.Background() collID := int64(100) diff --git a/internal/datacoord/compaction_policy_storage_version.go b/internal/datacoord/compaction_policy_storage_version.go index 390197b561c..ceff32d3b56 100644 --- a/internal/datacoord/compaction_policy_storage_version.go +++ b/internal/datacoord/compaction_policy_storage_version.go @@ -104,15 +104,17 @@ func (policy *storageVersionUpgradePolicy) Trigger(ctx context.Context) (map[Com return map[CompactionTriggerType][]CompactionView{}, nil } - collections := policy.meta.GetCollections() - if time.Since(policy.lastPeriod) > paramtable.Get().DataCoordCfg.StorageVersionCompactionRateLimitInterval.GetAsDuration(time.Second) { policy.currentCount = 0 policy.lastPeriod = time.Now() } maxCount := paramtable.Get().DataCoordCfg.StorageVersionCompactionRateLimitTokens.GetAsInt() + if policy.currentCount >= maxCount { + return map[CompactionTriggerType][]CompactionView{TriggerTypeStorageVersionUpgrade: nil}, nil + } + collections := policy.meta.GetCollections() views := make([]CompactionView, 0) for _, collection := range collections { if policy.currentCount >= maxCount { @@ -135,7 +137,6 @@ func (policy *storageVersionUpgradePolicy) Trigger(ctx context.Context) (map[Com } func (policy *storageVersionUpgradePolicy) triggerOneCollection(ctx context.Context, collectionID int64, maxCount int) ([]CompactionView, error) { - log := mlog.With(mlog.FieldCollectionID(collectionID)) collection, err := policy.handler.GetCollection(ctx, collectionID) if err != nil { mlog.Warn(ctx, "fail to apply storageVersionUpgradePolicy, unable to get collection from handler", @@ -147,7 +148,7 @@ func (policy *storageVersionUpgradePolicy) triggerOneCollection(ctx context.Cont return nil, nil } if collection.IsExternal() { - log.Info(ctx, "skip storage version compaction for external collection") + mlog.Info(ctx, "skip storage version compaction for external collection", mlog.FieldCollectionID(collectionID)) return nil, nil } @@ -157,12 +158,6 @@ func (policy *storageVersionUpgradePolicy) triggerOneCollection(ctx context.Cont return nil, err } - newTriggerID, err := policy.allocator.AllocID(ctx) - if err != nil { - mlog.Warn(ctx, "fail to apply storageVersionUpgradePolicy, unable to allocate triggerID", mlog.Err(err)) - return nil, err - } - targetVersion := policy.targetVersion() // TEXT fields require V3 manifest storage for LOB support and cannot be // downgraded. If the configured target version is lower than V3 for a @@ -195,6 +190,15 @@ func (policy *storageVersionUpgradePolicy) triggerOneCollection(ctx context.Cont segment.GetStorageVersion() == storage.StorageV3 && !segmentColumnGroupFormatsAllEqual(segment, targetFormat))) })) + if len(segments) == 0 || policy.currentCount >= maxCount { + return nil, nil + } + + newTriggerID, err := policy.allocator.AllocID(ctx) + if err != nil { + mlog.Warn(ctx, "fail to apply storageVersionUpgradePolicy, unable to allocate triggerID", mlog.Err(err)) + return nil, err + } views := make([]CompactionView, 0, len(segments)) for _, segment := range segments { diff --git a/internal/datacoord/compaction_policy_storage_version_test.go b/internal/datacoord/compaction_policy_storage_version_test.go index 12fd27e3373..04c8621f9ba 100644 --- a/internal/datacoord/compaction_policy_storage_version_test.go +++ b/internal/datacoord/compaction_policy_storage_version_test.go @@ -23,6 +23,7 @@ import ( "time" "github.com/blang/semver/v4" + "github.com/bytedance/mockey" "github.com/stretchr/testify/mock" "github.com/stretchr/testify/suite" @@ -148,6 +149,94 @@ func (s *StorageVersionUpgradePolicySuite) TestTriggerNoCollections() { s.Empty(gotViews) } +func (s *StorageVersionUpgradePolicySuite) TestTriggerNoCandidatesSkipsAllocID() { + params := paramtable.Get() + params.Save(params.CommonCfg.UseLoonFFI.Key, "true") + params.Save(params.DataCoordCfg.StorageVersionCompactionEnabled.Key, "true") + params.Save(params.DataCoordCfg.StorageFormatCompactionEnabled.Key, "false") + defer params.Reset(params.CommonCfg.UseLoonFFI.Key) + defer params.Reset(params.DataCoordCfg.StorageVersionCompactionEnabled.Key) + defer params.Reset(params.DataCoordCfg.StorageFormatCompactionEnabled.Key) + + coll := &collectionInfo{ID: 100, Schema: newTestSchema()} + s.handler.EXPECT().GetCollection(mock.Anything, coll.ID).Return(coll, nil) + for _, test := range []struct { + name string + segments map[UniqueID]*SegmentInfo + }{ + {name: "empty"}, + {name: "up to date", segments: map[UniqueID]*SegmentInfo{ + 101: newStorageVersionPolicyTestSegment(coll.ID, 101, storage.StorageV3, "parquet"), + }}, + } { + s.Run(test.name, func() { + s.setPolicyMeta(coll.ID, coll, test.segments) + s.policy.currentCount = 1 + views, err := s.policy.triggerOneCollection(context.Background(), coll.ID, 10) + s.NoError(err) + s.Empty(views) + s.Equal(1, s.policy.currentCount) + s.mockAlloc.AssertNotCalled(s.T(), "AllocID", mock.Anything) + }) + } +} + +func (s *StorageVersionUpgradePolicySuite) TestTriggerExhaustedBudgetSkipsCollectionEnumeration() { + params := paramtable.Get() + params.Save(params.DataCoordCfg.StorageVersionCompactionRateLimitInterval.Key, "3600") + params.Save(params.DataCoordCfg.StorageVersionCompactionRateLimitTokens.Key, "2") + defer params.Reset(params.DataCoordCfg.StorageVersionCompactionRateLimitInterval.Key) + defer params.Reset(params.DataCoordCfg.StorageVersionCompactionRateLimitTokens.Key) + s.versionMgr.EXPECT().GetMinimalSessionVer().Return(semver.MustParse("3.0.0")) + s.policy.currentCount = 2 + s.policy.lastPeriod = time.Now() + + enumerations := 0 + collectionsMock := mockey.Mock((*meta).GetCollections).To(func(*meta) []*collectionInfo { + enumerations++ + return nil + }).Build() + defer collectionsMock.UnPatch() + + events, err := s.policy.Trigger(context.Background()) + s.NoError(err) + s.Contains(events, TriggerTypeStorageVersionUpgrade) + s.Empty(events[TriggerTypeStorageVersionUpgrade]) + s.Zero(enumerations) + s.Equal(2, s.policy.currentCount) + s.mockAlloc.AssertNotCalled(s.T(), "AllocID", mock.Anything) + s.handler.AssertNotCalled(s.T(), "GetCollection", mock.Anything, mock.Anything) + + // The existing rate window still replenishes the budget and resumes scans. + s.policy.lastPeriod = time.Now().Add(-2 * time.Hour) + _, err = s.policy.Trigger(context.Background()) + s.NoError(err) + s.Equal(1, enumerations) + s.Zero(s.policy.currentCount) +} + +func (s *StorageVersionUpgradePolicySuite) TestTextCollectionDoesNotDowngradeOrAllocateID() { + params := paramtable.Get() + params.Save(params.CommonCfg.UseLoonFFI.Key, "false") + params.Save(params.DataCoordCfg.StorageVersionCompactionEnabled.Key, "true") + defer params.Reset(params.CommonCfg.UseLoonFFI.Key) + defer params.Reset(params.DataCoordCfg.StorageVersionCompactionEnabled.Key) + + coll := &collectionInfo{ID: 100, Schema: &schemapb.CollectionSchema{ + Fields: []*schemapb.FieldSchema{{FieldID: 100, Name: "text", DataType: schemapb.DataType_Text}}, + }} + s.handler.EXPECT().GetCollection(mock.Anything, coll.ID).Return(coll, nil).Once() + s.setPolicyMeta(coll.ID, coll, map[UniqueID]*SegmentInfo{ + 101: newStorageVersionPolicyTestSegment(coll.ID, 101, storage.StorageV3, "parquet"), + }) + + views, err := s.policy.triggerOneCollection(context.Background(), coll.ID, 10) + s.NoError(err) + s.Empty(views) + s.Zero(s.policy.currentCount) + s.mockAlloc.AssertNotCalled(s.T(), "AllocID", mock.Anything) +} + func (s *StorageVersionUpgradePolicySuite) TestTriggerWithSegments() { ctx := context.Background() collID := int64(100) @@ -377,7 +466,6 @@ func (s *StorageVersionUpgradePolicySuite) TestFormatRefreshRespectsSegmentFilte Schema: newTestSchema(), } s.handler.EXPECT().GetCollection(mock.Anything, mock.Anything).Return(coll, nil) - s.mockAlloc.EXPECT().AllocID(mock.Anything).Return(int64(1000), nil) segments := map[UniqueID]*SegmentInfo{ 101: newStorageVersionPolicyTestSegment(collID, 101, storage.StorageV3, "parquet"), @@ -479,7 +567,6 @@ func (s *StorageVersionUpgradePolicySuite) TestTriggerWithCompactingSegment() { Schema: newTestSchema(), } s.handler.EXPECT().GetCollection(mock.Anything, mock.Anything).Return(coll, nil) - s.mockAlloc.EXPECT().AllocID(mock.Anything).Return(int64(1000), nil) // Create a compacting segment (should NOT be upgraded) segments := make(map[UniqueID]*SegmentInfo) @@ -525,7 +612,6 @@ func (s *StorageVersionUpgradePolicySuite) TestTriggerWithImportingSegment() { Schema: newTestSchema(), } s.handler.EXPECT().GetCollection(mock.Anything, mock.Anything).Return(coll, nil) - s.mockAlloc.EXPECT().AllocID(mock.Anything).Return(int64(1000), nil) // Create an importing segment (should NOT be upgraded) segments := make(map[UniqueID]*SegmentInfo) @@ -612,6 +698,14 @@ func (s *StorageVersionUpgradePolicySuite) TestTriggerRateLimiting() { views, err := s.policy.triggerOneCollection(ctx, collID, 2) s.NoError(err) s.Equal(2, len(views)) + s.Equal(2, s.policy.currentCount) + s.mockAlloc.AssertNumberOfCalls(s.T(), "AllocID", 1) + + views, err = s.policy.triggerOneCollection(ctx, collID, 2) + s.NoError(err) + s.Empty(views) + s.Equal(2, s.policy.currentCount) + s.mockAlloc.AssertNumberOfCalls(s.T(), "AllocID", 1) } func (s *StorageVersionUpgradePolicySuite) TestTriggerIntervalReset() { @@ -691,25 +785,35 @@ func (s *StorageVersionUpgradePolicySuite) TestTriggerGetCollectionError() { func (s *StorageVersionUpgradePolicySuite) TestTriggerAllocIDError() { ctx := context.Background() collID := int64(100) + params := paramtable.Get() + params.Save(params.CommonCfg.UseLoonFFI.Key, "true") + params.Save(params.DataCoordCfg.StorageVersionCompactionEnabled.Key, "true") + defer params.Reset(params.CommonCfg.UseLoonFFI.Key) + defer params.Reset(params.DataCoordCfg.StorageVersionCompactionEnabled.Key) coll := &collectionInfo{ ID: collID, Schema: newTestSchema(), } s.handler.EXPECT().GetCollection(mock.Anything, mock.Anything).Return(coll, nil) - s.mockAlloc.EXPECT().AllocID(mock.Anything).Return(int64(0), context.DeadlineExceeded) - - collections := typeutil.NewConcurrentMap[UniqueID, *collectionInfo]() - collections.Insert(collID, coll) - - s.policy.meta = &meta{ - segments: NewCachedSegmentsInfo(), - collections: collections, - } + s.mockAlloc.EXPECT().AllocID(mock.Anything).Return(int64(0), context.DeadlineExceeded).Once() + s.mockAlloc.EXPECT().AllocID(mock.Anything).Return(int64(1000), nil).Once() + s.setPolicyMeta(collID, coll, map[UniqueID]*SegmentInfo{ + 101: newStorageVersionPolicyTestSegment(collID, 101, storage.StorageV2, "parquet"), + }) views, err := s.policy.triggerOneCollection(ctx, collID, 10) - s.Error(err) + s.ErrorIs(err, context.DeadlineExceeded) s.Nil(views) + s.Zero(s.policy.currentCount) + + // Failed ID allocation must leave the candidate and task budget intact. + views, err = s.policy.triggerOneCollection(ctx, collID, 10) + s.NoError(err) + s.Require().Len(views, 1) + s.EqualValues(101, views[0].GetSegmentsView()[0].ID) + s.EqualValues(1000, views[0].GetTriggerID()) + s.Equal(1, s.policy.currentCount) } func (s *StorageVersionUpgradePolicySuite) TestTriggerMultipleCollections() { @@ -856,7 +960,6 @@ func (s *StorageVersionUpgradePolicySuite) TestDroppedSegmentFiltered() { Schema: newTestSchema(), } s.handler.EXPECT().GetCollection(mock.Anything, mock.Anything).Return(coll, nil) - s.mockAlloc.EXPECT().AllocID(mock.Anything).Return(int64(1000), nil) // Create a dropped segment (should NOT be upgraded) segments := make(map[UniqueID]*SegmentInfo) @@ -901,7 +1004,6 @@ func (s *StorageVersionUpgradePolicySuite) TestGrowingSegmentFiltered() { Schema: newTestSchema(), } s.handler.EXPECT().GetCollection(mock.Anything, mock.Anything).Return(coll, nil) - s.mockAlloc.EXPECT().AllocID(mock.Anything).Return(int64(1000), nil) // Create a growing segment (should NOT be upgraded - not flushed) segments := make(map[UniqueID]*SegmentInfo)