enhance: add event-driven stats discovery

configs/milvus.yaml

@@ -762,6 +762,9 @@ indexNode:
 
 dataCoord:
   taskCheckInterval: 1
+  statsInspector:
+    discoveryMode: poll # poll: legacy discovery; shadow: read-only event checks; event: event-driven discovery. Restart required.
+    reconcileInterval: 600 # Seconds. Startup and lost-notification reconciliation use the same budget.
   channel:
     watchTimeoutInterval: 300 # Timeout on watching channels (in seconds). Datanode tickler update watch progress will reset timeout timer.
     legacyVersionWithoutRPCWatch: 2.4.1 # Datanodes <= this version are considered as legacy nodes, which doesn't have rpc based watch(). This is only used during rolling upgrade where legacy nodes won't get new channels

internal/datacoord/meta.go

@@ -102,6 +102,7 @@ type meta struct {
 	segments                  *CachedSegmentsInfo // segment id to segment info
 	dataViewManager           DataViewManager
 	queryViewLoadInfoNotifier QueryViewLoadInfoNotifier
+	statsDiscovery            atomic.Pointer[statsReconcileQueue]
 
 	channelCPs   *channelCPs // vChannel -> channel checkpoint/see position
 	chunkManager storage.ChunkManager
@@ -744,7 +745,11 @@ func (m *meta) addCollectionToCache(collection *collectionInfo) {
 // Note that collection info is just for caching and will not be set into etcd from datacoord
 func (m *meta) AddCollection(collection *collectionInfo) {
 	mlog.Info(context.TODO(), "meta update: add collection", zap.Int64("collectionID", collection.ID))
+	old := m.GetCollection(collection.ID)
 	m.addCollectionToCache(collection)
+	if q := m.statsDiscovery.Load(); q != nil && (old == nil || !proto.Equal(old.Schema, collection.Schema)) {
+		q.requestScan(collection.ID, true)
+	}
 	metrics.DataCoordNumCollections.WithLabelValues().Set(float64(m.collections.Len()))
 	mlog.Info(context.TODO(), "meta update: add collection - complete", zap.Int64("collectionID", collection.ID))
 }
@@ -753,6 +758,9 @@ func (m *meta) AddCollection(collection *collectionInfo) {
 func (m *meta) DropCollection(collectionID int64) {
 	mlog.Info(context.TODO(), "meta update: drop collection", zap.Int64("collectionID", collectionID))
 	if _, ok := m.collections.GetAndRemove(collectionID); ok {
+		if q := m.statsDiscovery.Load(); q != nil {
+			q.requestScan(collectionID, true)
+		}
 		metrics.CleanupDataCoordWithCollectionID(collectionID)
 		metrics.DataCoordNumCollections.WithLabelValues().Set(float64(m.collections.Len()))
 		mlog.Info(context.TODO(), "meta update: drop collection - complete", zap.Int64("collectionID", collectionID))
@@ -1040,6 +1048,7 @@ func (m *meta) AddSegment(ctx context.Context, segment *SegmentInfo) error {
 		return err
 	}
 	m.segments.SetSegment(segment.GetID(), segment, results[0].Version)
+	m.notifyStatsChange(nil, segment)
 
 	metrics.DataCoordNumSegments.WithLabelValues(segmentMetricLabelValues(segment)...).Inc()
 	logger.Info(ctx, "meta update: adding segment - complete", zap.Int64("segmentID", segment.GetID()))
@@ -1059,6 +1068,7 @@ func (m *meta) DropSegment(ctx context.Context, segment *SegmentInfo) error {
 		if errors.Is(err, ErrKeyNotFound) {
 			logger.Info(ctx, "meta update: dropping segment - already deleted", zap.Int64("segmentID", segmentID))
 			m.segments.DropSegment(segmentID, math.MaxInt64)
+			m.notifyStatsChange(segment, nil)
 			return nil
 		}
 		logger.Warn(ctx, "meta update: dropping segment failed",
@@ -1069,6 +1079,7 @@ func (m *meta) DropSegment(ctx context.Context, segment *SegmentInfo) error {
 	metrics.DataCoordNumSegments.WithLabelValues(segmentMetricLabelValues(segment)...).Dec()
 
 	m.segments.DropSegment(segmentID, results[0].Version)
+	m.notifyStatsChange(segment, nil)
 	logger.Info(ctx, "meta update: dropping segment - complete",
 		zap.Int64("segmentID", segmentID))
 	return nil
@@ -1138,6 +1149,7 @@ func (m *meta) DropSegments(ctx context.Context, candidates []*SegmentInfo) (int
 				if singleErr != nil {
 					if errors.Is(singleErr, ErrKeyNotFound) {
 						m.segments.DropSegment(segment.GetID(), math.MaxInt64)
+						m.notifyStatsChange(segment, nil)
 						removed++
 						continue
 					}
@@ -1152,6 +1164,7 @@ func (m *meta) DropSegments(ctx context.Context, candidates []*SegmentInfo) (int
 				}
 				metrics.DataCoordNumSegments.WithLabelValues(segmentMetricLabelValues(segment)...).Dec()
 				m.segments.DropSegment(segment.GetID(), singleResults[0].Version)
+				m.notifyStatsChange(segment, nil)
 				removed++
 			}
 			return removed, deleteErr
@@ -1167,6 +1180,7 @@ func (m *meta) DropSegments(ctx context.Context, candidates []*SegmentInfo) (int
 	for i, segment := range segments {
 		metrics.DataCoordNumSegments.WithLabelValues(segmentMetricLabelValues(segment)...).Dec()
 		m.segments.DropSegment(segment.GetID(), results[i].Version)
+		m.notifyStatsChange(segment, nil)
 	}
 	return len(segments), nil
 }
@@ -1277,6 +1291,7 @@ func (m *meta) SetState(ctx context.Context, segmentID UniqueID, targetState com
 	}
 	updatedSeg := NewSegmentInfo(results[0].Value)
 	old, existed := m.segments.SetSegment(segmentID, updatedSeg, results[0].Version)
+	m.notifyStatsChange(old, updatedSeg)
 	if existed && old.GetState() != updatedSeg.GetState() {
 		metricMutation := segMetricMutation{stateChange: make(segmentMetricStateChange)}
 		metricMutation.appendSegmentLabelChange(old, updatedSeg)
@@ -1318,7 +1333,9 @@ func (m *meta) UpdateSegment(segmentID int64, operators ...SegmentOperator) erro
 		return err
 	}
 	// Update in-memory meta.
-	m.segments.SetSegment(segmentID, NewSegmentInfo(results[0].Value), results[0].Version)
+	updated := NewSegmentInfo(results[0].Value)
+	old, _ := m.segments.SetSegment(segmentID, updated, results[0].Version)
+	m.notifyStatsChange(old, updated)
 
 	logger.Info(context.TODO(), "meta update: update segment - complete",
 		zap.Int64("segmentID", segmentID))
@@ -1692,7 +1709,8 @@ func (m *meta) UpdateSegmentsInfo(ctx context.Context, mutations map[int64][]Mut
 	type entry struct {
 		segID    int64
 		isInsert bool
-		newSeg   *SegmentInfo // only for inserts
+		newSeg   *SegmentInfo // published insert/update value
+		oldSeg   *SegmentInfo
 	}
 	var entries []entry
 
@@ -1764,13 +1782,19 @@ func (m *meta) UpdateSegmentsInfo(ctx context.Context, mutations map[int64][]Mut
 		} else {
 			newSeg := NewSegmentInfo(results[i].Value)
 			oldSeg, existed := m.segments.SetSegment(e.segID, newSeg, results[i].Version)
+			entries[i].oldSeg = oldSeg
+			entries[i].newSeg = newSeg
 			if existed && !sameSegmentMetricLabels(oldSeg, newSeg) {
 				metricMutation.appendSegmentLabelChange(oldSeg, newSeg)
 			}
 		}
 	}
 	metricMutation.commit()
 	cacheDur := time.Since(cacheStart)
+	// Publish hints only once the whole committed batch is visible in cache.
+	for _, e := range entries {
+		m.notifyStatsChange(e.oldSeg, e.newSeg)
+	}
 
 	totalDur := time.Since(start)
 	if totalDur > 40*time.Millisecond {
@@ -1929,6 +1953,10 @@ func (m *meta) UpdateDropChannelSegmentInfo(ctx context.Context, channel string,
 	}
 	metricMutation.commit()
 
+	// All cache entries in this transaction have been published.
+	for _, result := range results {
+		m.notifyStatsSegments(result.Value.GetCollectionID(), result.Value.GetID())
+	}
 	logger.Info(ctx, "meta update: update drop channel segment info - complete",
 		zap.String("channel", channel))
 	return nil
@@ -2019,7 +2047,6 @@ func (m *meta) GetFlushingSegments() []*SegmentInfo {
 
 // SelectSegments select segments with selector
 func (m *meta) SelectSegments(ctx context.Context, filters ...SegmentFilter) []*SegmentInfo {
-
 	return m.segments.GetSegmentsBySelector(filters...)
 }
 
@@ -2049,7 +2076,6 @@ func (m *meta) GetCollectionIDsByPartition(ctx context.Context, partitionIDs []i
 }
 
 func (m *meta) GetRealSegmentsForChannel(channel string) []*SegmentInfo {
-
 	return m.segments.GetRealSegmentsForChannel(channel)
 }
 
@@ -2073,14 +2099,12 @@ func (m *meta) AddAllocation(segmentID UniqueID, allocation *Allocation) error {
 }
 
 func (m *meta) SetRowCount(segmentID UniqueID, rowCount int64) {
-
 	m.segments.SetRowCount(segmentID, rowCount)
 }
 
 // SetAllocations set Segment allocations, will overwrite ALL original allocations
 // Note that allocations is not persisted in KV store
 func (m *meta) SetAllocations(segmentID UniqueID, allocations []*Allocation) {
-
 	m.segments.SetAllocations(segmentID, allocations)
 }
 
@@ -2093,26 +2117,22 @@ func (m *meta) SetLastExpire(segmentID UniqueID, lastExpire uint64) {
 // SetLastFlushTime set LastFlushTime for segment with provided `segmentID`
 // Note that lastFlushTime is not persisted in KV store
 func (m *meta) SetLastFlushTime(segmentID UniqueID, t time.Time) {
-
 	m.segments.SetFlushTime(segmentID, t)
 }
 
 // SetLastWrittenTime set LastWrittenTime for segment with provided `segmentID`
 // Note that lastWrittenTime is not persisted in KV store
 func (m *meta) SetLastWrittenTime(segmentID UniqueID) {
-
 	m.segments.SetLastWrittenTime(segmentID)
 }
 
 // SetSegmentCompacting sets compaction state for segment
 func (m *meta) SetSegmentCompacting(segmentID UniqueID, compacting bool) {
-
 	m.segments.SetIsCompacting(segmentID, compacting)
 }
 
 // IsSegmentCompacting check if segment is compacting
 func (m *meta) IsSegmentCompacting(segmentID UniqueID) bool {
-
 	seg := m.segments.GetSegment(segmentID)
 	if seg == nil {
 		return false
@@ -2124,7 +2144,6 @@ func (m *meta) IsSegmentCompacting(segmentID UniqueID) bool {
 // if true, set them compacting and return true
 // if false, skip setting and
 func (m *meta) CheckAndSetSegmentsCompacting(ctx context.Context, segmentIDs []UniqueID) (exist, canDo bool) {
-
 	var hasCompacting bool
 	exist = true
 	for _, segmentID := range segmentIDs {
@@ -2148,7 +2167,6 @@ func (m *meta) CheckAndSetSegmentsCompacting(ctx context.Context, segmentIDs []U
 }
 
 func (m *meta) SetSegmentsCompacting(ctx context.Context, segmentIDs []UniqueID, compacting bool) {
-
 	for _, segmentID := range segmentIDs {
 		m.segments.SetIsCompacting(segmentID, compacting)
 	}
@@ -2470,7 +2488,6 @@ func (m *meta) completeMixCompactionMutation(
 }
 
 func (m *meta) ValidateSegmentStateBeforeCompleteCompactionMutation(t *datapb.CompactionTask) error {
-
 	if t.GetType() != datapb.CompactionType_Level0DeleteCompaction {
 		if m.isCollectionCompactionBlocked(t.GetCollectionID()) {
 			mlog.Info(context.TODO(), "compaction rejected: collection has pending snapshot or unloaded RefIndex",
@@ -2541,6 +2558,10 @@ func (m *meta) CompleteCompactionMutation(ctx context.Context, t *datapb.Compact
 	m.publishDataViewAfterCompaction(ctx, t, lo.Map(newSegments, func(segment *SegmentInfo, _ int) int64 {
 		return segment.GetID()
 	}))
+	m.notifyStatsSegments(t.GetCollectionID(), t.GetInputSegments()...)
+	for _, segment := range newSegments {
+		m.notifyStatsSegments(segment.GetCollectionID(), segment.GetID())
+	}
 	return newSegments, metricMutation, nil
 }
 
@@ -2583,7 +2604,6 @@ func isSegmentHealthy(segment *SegmentInfo) bool {
 }
 
 func (m *meta) HasSegments(segIDs []UniqueID) (bool, error) {
-
 	for _, segID := range segIDs {
 		if m.segments.GetSegment(segID) == nil {
 			return false, fmt.Errorf("segment is not exist with ID = %d", segID)
@@ -2594,7 +2614,6 @@ func (m *meta) HasSegments(segIDs []UniqueID) (bool, error) {
 
 // GetCompactionTo returns the segment info of the segment to be compacted to.
 func (m *meta) GetCompactionTo(segmentID int64) ([]*SegmentInfo, bool) {
-
 	return m.segments.GetCompactionTo(segmentID)
 }
 
@@ -3395,7 +3414,6 @@ func (m *meta) completeBumpSchemaVersionReplacementMutation(
 }
 
 func (m *meta) getSegmentsMetrics(collectionID int64) []*metricsinfo.Segment {
-
 	allSegments := m.segments.GetSegments()
 	segments := make([]*metricsinfo.Segment, 0, len(allSegments))
 	for _, s := range allSegments {
@@ -3421,7 +3439,6 @@ func (m *meta) getSegmentsMetrics(collectionID int64) []*metricsinfo.Segment {
 }
 
 func (m *meta) DropSegmentsOfPartition(ctx context.Context, partitionIDs []int64) error {
-
 	// Collect segments to drop (read-only from cache for key construction).
 	type segRef struct {
 		id  int64
@@ -3463,6 +3480,9 @@ func (m *meta) DropSegmentsOfPartition(ctx context.Context, partitionIDs []int64
 		}
 	}
 	metricMutation.commit()
+	for _, result := range results {
+		m.notifyStatsSegments(result.Value.GetCollectionID(), result.Value.GetID())
+	}
 	return nil
 }
 
@@ -3484,7 +3504,6 @@ func (m *meta) GetFileResources(ctx context.Context, resourceIDs ...int64) ([]*i
 
 // TruncateChannelByTime drops segments of a channel that were updated before the flush timestamp
 func (m *meta) TruncateChannelByTime(ctx context.Context, vChannel string, flushTs uint64) error {
-
 	segments := m.segments.GetSegmentsBySelector(SegmentFilterFunc(isSegmentHealthy), WithChannel(vChannel))
 
 	// Collect segments to drop (read-only from cache for key construction and filtering).
@@ -3533,6 +3552,9 @@ func (m *meta) TruncateChannelByTime(ctx context.Context, vChannel string, flush
 		}
 	}
 	metricMutation.commit()
+	for _, result := range results {
+		m.notifyStatsSegments(result.Value.GetCollectionID(), result.Value.GetID())
+	}
 
 	return nil
 }

internal/datacoord/stats_discovery_benchmark_test.go

@@ -0,0 +1,60 @@
+// 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 (
+	"fmt"
+	"iter"
+	"testing"
+
+	"github.com/milvus-io/milvus/pkg/v3/common"
+	"github.com/milvus-io/milvus/pkg/v3/proto/datapb"
+)
+
+func BenchmarkStatsDiscoverySteady(b *testing.B) {
+	for _, size := range []int{1000, 100000} {
+		b.Run(fmt.Sprint(size), func(b *testing.B) {
+			f := newDiscoveryFixture(b, "event")
+			for id := int64(1); id <= int64(size); id++ {
+				segment := discoverySegment(id, true)
+				segment.TextStatsLogs = map[int64]*datapb.TextIndexStats{101: {}}
+				segment.JsonKeyStats = map[int64]*datapb.JsonKeyStats{102: {JsonKeyStatsDataFormat: common.JSONStatsDataFormatVersion}}
+				f.mt.segments.SetSegment(id, segment, 1)
+			}
+			b.Run("poll_with_direct_field_check", func(b *testing.B) {
+				b.ReportAllocs()
+				for range b.N {
+					f.si.triggerStatsTasks(0)
+				}
+			})
+			b.Run("event_no_change", func(b *testing.B) {
+				b.ReportAllocs()
+				for range b.N {
+					f.si.processStatsDiscoveryBatch()
+				}
+			})
+			b.Run("stream_first_entry", func(b *testing.B) {
+				b.ReportAllocs()
+				for range b.N {
+					next, stop := iter.Pull2(f.mt.rangeStatsSegments(0))
+					next()
+					stop()
+				}
+			})
+		})
+	}
+}