enhance: reduce idle compaction allocations

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
+	}{
+		{&params.CommonCfg.UseLoonFFI, "true"},
+		{&params.DataCoordCfg.StorageVersionCompactionEnabled, "true"},
+		{&params.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")
+						}
+					}
+				}
+			})
+		}
+	}
+}

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

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)

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 {