fix: handle text lob refs and flush pressure accounting (#51079)

issue: #50877

Signed-off-by: luzhang <luzhang@zilliz.com>
Co-authored-by: luzhang <luzhang@zilliz.com>

internal/flushcommon/metacache/flush_source.go

@@ -30,11 +30,8 @@ package metacache
 // phase remains). Growing segments re-derive the value through the regular
 // write-buffer source decision path during WAL replay.
 //
-// Currently consumed only by the TEXT flush owner so that the writeBuffer
-// can route between the live insert payload and an out-of-process growing
-// source provider without keeping its own shadow map. The field is generic
-// so future large-field offload paths (BM25, vector, etc.) can reuse the
-// same machinery.
+// Consumers use it to route writebuffer sync tasks and to refine StreamingNode
+// runtime flush-pressure accounting after the per-segment source is known.
 type FlushSourceMode int32
 
 const (

internal/flushcommon/pipeline/data_sync_service.go

@@ -314,10 +314,15 @@ func getServiceWithChannel(initCtx context.Context, params *util.PipelineParams,
 	// Register channel after channel pipeline is ready.
 	// This'll reject any FlushChannel and FlushSegments calls to prevent inconsistency between DN and DC over flushTs
 	// if fail to init flowgraph nodes.
-	err = params.WriteBufferManager.Register(channelName, metacache,
+	writeBufferOptions := []writebuffer.WriteBufferOption{
 		writebuffer.WithMetaWriter(syncmgr.BrokerMetaWriter(params.Broker, config.serverID)),
 		writebuffer.WithIDAllocator(params.Allocator),
-		writebuffer.WithTaskObserverCallback(wbTaskObserverCallback))
+		writebuffer.WithTaskObserverCallback(wbTaskObserverCallback),
+	}
+	if params.FlushSourceModeNotifier != nil {
+		writeBufferOptions = append(writeBufferOptions, writebuffer.WithFlushSourceModeNotifier(params.FlushSourceModeNotifier))
+	}
+	err = params.WriteBufferManager.Register(channelName, metacache, writeBufferOptions...)
 	if err != nil {
 		mlog.Warn(initCtx, "failed to register channel buffer", mlog.String("channel", channelName), mlog.Err(err))
 		return nil, err

internal/flushcommon/util/util.go

@@ -49,6 +49,9 @@ type PipelineParams struct {
 	Allocator          allocator.Interface
 	MsgHandler         MsgHandler
 	SchemaManager      metacache.SchemaManager
+	// FlushSourceModeNotifier is an optional callback invoked when writebuffer
+	// decides the sticky flush source mode for a segment.
+	FlushSourceModeNotifier writebuffer.FlushSourceModeNotifier
 }
 
 // TimeRange is a range of timestamp contains the min-timestamp and max-timestamp

internal/flushcommon/writebuffer/options.go

@@ -14,6 +14,8 @@ type WriteBufferOption func(opt *writeBufferOption)
 
 type TaskObserverCallback func(t syncmgr.Task, err error)
 
+type FlushSourceModeNotifier func(segmentID int64, mode metacache.FlushSourceMode)
+
 // GrowingSourceResolver resolves an optional in-memory growing segment source
 // for growing-source flush. GrowingSourcePending means the growing source exists but has not
 // caught up to targetOffset yet; WriteBuffer should only be used when the state
@@ -32,6 +34,7 @@ type writeBufferOption struct {
 
 	growingSourceResolver      GrowingSourceResolver
 	growingSourceRetryInterval time.Duration
+	flushSourceModeNotifier    FlushSourceModeNotifier
 }
 
 func defaultWBOption(metacache metacache.MetaCache) *writeBufferOption {
@@ -92,3 +95,9 @@ func WithGrowingSourceResolver(resolver GrowingSourceResolver) WriteBufferOption
 		opt.growingSourceResolver = resolver
 	}
 }
+
+func WithFlushSourceModeNotifier(notifier FlushSourceModeNotifier) WriteBufferOption {
+	return func(opt *writeBufferOption) {
+		opt.flushSourceModeNotifier = notifier
+	}
+}

internal/flushcommon/writebuffer/write_buffer.go

@@ -281,6 +281,7 @@ type writeBufferBase struct {
 	growingSourceRetryInterval  time.Duration
 	growingSourceRetryScheduled bool
 	growingSourceRetryTimer     *time.Timer
+	flushSourceModeNotifier     FlushSourceModeNotifier
 	closed                      bool
 
 	// pre build logger
@@ -334,6 +335,7 @@ func newWriteBufferBase(channel string, metacache metacache.MetaCache, syncMgr s
 		growingSourceResolver:      growingSourceResolver,
 		growingSourceProgress:      make(map[int64]*growingSourceProgress),
 		growingSourceRetryInterval: growingSourceRetryInterval,
+		flushSourceModeNotifier:    option.flushSourceModeNotifier,
 	}
 
 	wb.logger = mlog.With(mlog.Int64("collectionID", wb.collectionID),
@@ -807,6 +809,7 @@ func (wb *writeBufferBase) recordGrowingSourceProgress(inData *InsertData, start
 		metacache.SetFlushSourceMode(metacache.FlushSourceGrowing),
 		wb.updateGrowingSourceBufferedRows(progress),
 	), metacache.WithSegmentIDs(inData.segmentID))
+	wb.notifyFlushSourceMode(inData.segmentID)
 }
 
 func (wb *writeBufferBase) growingSourceTargetOffset(segmentID int64, rows int64) int64 {
@@ -1159,12 +1162,27 @@ func (wb *writeBufferBase) getOrCreateBuffer(segmentID int64, timetick uint64) *
 				metacache.SetFlushSourceMode(metacache.FlushSourceWriteBuffer),
 				metacache.WithSegmentIDs(segmentID),
 			)
+			wb.notifyFlushSourceMode(segmentID)
 		}
 	}
 
 	return buffer
 }
 
+func (wb *writeBufferBase) notifyFlushSourceMode(segmentID int64) {
+	if wb.flushSourceModeNotifier == nil {
+		return
+	}
+	segment, ok := wb.metaCache.GetSegmentByID(segmentID)
+	if !ok {
+		return
+	}
+	switch mode := segment.FlushSourceMode(); mode {
+	case metacache.FlushSourceWriteBuffer, metacache.FlushSourceGrowing:
+		wb.flushSourceModeNotifier(segmentID, mode)
+	}
+}
+
 func (wb *writeBufferBase) yieldBuffer(segmentID int64) ([]*storage.InsertData, map[int64]*storage.BM25Stats, *storage.DeleteData, *schemapb.CollectionSchema, *TimeRange, *msgpb.MsgPosition) {
 	buffer, ok := wb.buffers[segmentID]
 	if !ok {

internal/flushcommon/writebuffer/write_buffer_test.go

@@ -71,6 +71,7 @@ func (s *WriteBufferSuite) SetupTest() {
 func (s *WriteBufferSuite) TestHasSegment() {
 	segmentID := int64(1001)
 
+	s.wb.useGrowingSourceFlush = false
 	s.False(s.wb.HasSegment(segmentID))
 
 	s.wb.getOrCreateBuffer(segmentID, 0)
@@ -78,6 +79,51 @@ func (s *WriteBufferSuite) TestHasSegment() {
 	s.True(s.wb.HasSegment(segmentID))
 }
 
+func (s *WriteBufferSuite) TestFlushSourceModeNotifier() {
+	segmentID := int64(1001)
+	var notifiedSegmentID int64
+	var notifiedMode metacache.FlushSourceMode
+	s.wb.flushSourceModeNotifier = func(segmentID int64, mode metacache.FlushSourceMode) {
+		notifiedSegmentID = segmentID
+		notifiedMode = mode
+	}
+
+	s.Run("write_buffer_mode", func() {
+		segment := metacache.NewSegmentInfo(&datapb.SegmentInfo{ID: segmentID}, nil, nil)
+		metacache.SetFlushSourceMode(metacache.FlushSourceWriteBuffer)(segment)
+		s.metacache.EXPECT().UpdateSegments(mock.Anything, mock.Anything).Run(func(action metacache.SegmentAction, _ ...metacache.SegmentFilter) {
+			action(segment)
+		}).Return().Once()
+		s.metacache.EXPECT().GetSegmentByID(segmentID).Return(segment, true).Once()
+
+		s.wb.useGrowingSourceFlush = true
+		s.wb.getOrCreateBuffer(segmentID, 0)
+		s.Equal(segmentID, notifiedSegmentID)
+		s.Equal(metacache.FlushSourceWriteBuffer, notifiedMode)
+	})
+
+	s.Run("growing_mode", func() {
+		segmentID := int64(1002)
+		segment := metacache.NewSegmentInfo(&datapb.SegmentInfo{ID: segmentID}, nil, nil)
+		notifiedSegmentID = 0
+		notifiedMode = metacache.FlushSourceUnknown
+		s.metacache.EXPECT().GetSegmentByID(segmentID).Return(nil, false).Once()
+		s.metacache.EXPECT().AddSegment(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return().Once()
+		s.metacache.EXPECT().UpdateSegments(mock.Anything, mock.Anything).Run(func(action metacache.SegmentAction, _ ...metacache.SegmentFilter) {
+			action(segment)
+		}).Return().Once()
+		s.metacache.EXPECT().GetSegmentByID(segmentID).Return(segment, true).Once()
+
+		s.wb.recordGrowingSourceProgress(&InsertData{
+			segmentID:   segmentID,
+			partitionID: 10,
+			rowNum:      3,
+		}, &msgpb.MsgPosition{Timestamp: 100}, &msgpb.MsgPosition{Timestamp: 200}, 1, 3)
+		s.Equal(segmentID, notifiedSegmentID)
+		s.Equal(metacache.FlushSourceGrowing, notifiedMode)
+	})
+}
+
 func (s *WriteBufferSuite) TestCreateNewGrowingSegmentStorageVersion() {
 	param := paramtable.Get()
 	param.Save(param.CommonCfg.UseLoonFFI.Key, "false")

internal/storage/serde.go

@@ -112,6 +112,8 @@ type serdeEntry struct {
 	serialize func(b array.Builder, v any, elementType schemapb.DataType) error
 }
 
+type TextLobRef []byte
+
 var serdeMap = func() map[schemapb.DataType]serdeEntry {
 	m := make(map[schemapb.DataType]serdeEntry)
 	m[schemapb.DataType_Bool] = serdeEntry{
@@ -373,7 +375,52 @@ var serdeMap = func() map[schemapb.DataType]serdeEntry {
 
 	m[schemapb.DataType_VarChar] = stringEntry
 	m[schemapb.DataType_String] = stringEntry
-	m[schemapb.DataType_Text] = stringEntry
+	m[schemapb.DataType_Text] = serdeEntry{
+		arrowType: stringEntry.arrowType,
+		deserialize: func(a arrow.Array, i int, elementType schemapb.DataType, dim int, shouldCopy bool) (any, error) {
+			if a.IsNull(i) {
+				return nil, nil
+			}
+			if arr, ok := a.(*array.String); ok && i < arr.Len() {
+				value := arr.Value(i)
+				if shouldCopy {
+					return strings.Clone(value), nil
+				}
+				return value, nil
+			}
+			if arr, ok := a.(*array.Binary); ok && i < arr.Len() {
+				value := arr.Value(i)
+				if shouldCopy {
+					value = append([]byte(nil), value...)
+				}
+				return TextLobRef(value), nil
+			}
+			return nil, merr.WrapErrServiceInternalMsg("expected *array.String or *array.Binary, got %T", a)
+		},
+		serialize: func(b array.Builder, v any, elementType schemapb.DataType) error {
+			if v == nil {
+				b.AppendNull()
+				return nil
+			}
+			if builder, ok := b.(*array.StringBuilder); ok {
+				if v, ok := v.(string); ok {
+					builder.Append(v)
+					return nil
+				}
+				return merr.WrapErrServiceInternalMsg("expected string value, got %T", v)
+			}
+			if builder, ok := b.(*array.BinaryBuilder); ok {
+				switch v := v.(type) {
+				case TextLobRef:
+					builder.Append([]byte(v))
+					return nil
+				default:
+					return merr.WrapErrServiceInternalMsg("expected TEXT LOB reference value, got %T", v)
+				}
+			}
+			return merr.WrapErrServiceInternalMsg("expected *array.StringBuilder or *array.BinaryBuilder, got %T", b)
+		},
+	}
 
 	// We're not using the deserialized data in go, so we can skip the heavy pb serde.
 	// If there is need in the future, just assign it to m[schemapb.DataType_Array]

internal/storage/serde_events.go

@@ -429,8 +429,37 @@ func ValueSerializer(v []*Value, schema *schemapb.CollectionSchema) (Record, err
 		return nil, err
 	}
 
-	builders := make(map[FieldID]array.Builder, len(allFieldsSchema))
 	types := make(map[FieldID]schemapb.DataType, len(allFieldsSchema))
+	textRefFields := make(map[FieldID]struct{})
+	for _, f := range allFieldsSchema {
+		types[f.FieldID] = f.DataType
+	}
+	for _, vv := range v {
+		m := vv.Value.(map[FieldID]any)
+		for fid, value := range m {
+			if types[fid] != schemapb.DataType_Text {
+				continue
+			}
+			switch value.(type) {
+			case TextLobRef:
+				textRefFields[fid] = struct{}{}
+			}
+		}
+	}
+	if len(textRefFields) > 0 {
+		fields := make([]arrow.Field, arrowSchema.NumFields())
+		for i := 0; i < arrowSchema.NumFields(); i++ {
+			fields[i] = arrowSchema.Field(i)
+			if i < len(allFieldsSchema) {
+				if _, ok := textRefFields[allFieldsSchema[i].FieldID]; ok {
+					fields[i].Type = arrow.BinaryTypes.Binary
+				}
+			}
+		}
+		arrowSchema = arrow.NewSchema(fields, nil)
+	}
+
+	builders := make(map[FieldID]array.Builder, len(allFieldsSchema))
 	elementTypes := make(map[FieldID]schemapb.DataType, len(allFieldsSchema)) // For ArrayOfVector
 	for i, f := range allFieldsSchema {
 		if f.DataType == schemapb.DataType_ArrayOfVector {
@@ -439,7 +468,6 @@ func ValueSerializer(v []*Value, schema *schemapb.CollectionSchema) (Record, err
 
 		builders[f.FieldID] = array.NewBuilder(memory.DefaultAllocator, arrowSchema.Field(i).Type)
 		builders[f.FieldID].Reserve(len(v)) // reserve space to avoid copy
-		types[f.FieldID] = f.DataType
 	}
 
 	for _, vv := range v {

internal/storage/serde_events_test.go

@@ -926,6 +926,88 @@ func TestValueSerializerNullableDenseVectorUsesBinaryArrow(t *testing.T) {
 	}
 }
 
+func TestValueDeserializerSerializerTextLobRefUsesBinaryArrow(t *testing.T) {
+	const (
+		pkFieldID   FieldID = 100
+		textFieldID FieldID = 101
+	)
+	schema := &schemapb.CollectionSchema{
+		Name: "text_lob_ref",
+		Fields: []*schemapb.FieldSchema{
+			{FieldID: common.RowIDField, Name: "row_id", DataType: schemapb.DataType_Int64},
+			{FieldID: common.TimeStampField, Name: "Timestamp", DataType: schemapb.DataType_Int64},
+			{FieldID: pkFieldID, Name: "pk", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
+			{FieldID: textFieldID, Name: "content", DataType: schemapb.DataType_Text},
+		},
+	}
+	arrowSchema := arrow.NewSchema([]arrow.Field{
+		{Name: "row_id", Type: arrow.PrimitiveTypes.Int64},
+		{Name: "Timestamp", Type: arrow.PrimitiveTypes.Int64},
+		{Name: "pk", Type: arrow.PrimitiveTypes.Int64},
+		{Name: "content", Type: arrow.BinaryTypes.Binary},
+	}, nil)
+	builder := array.NewRecordBuilder(memory.DefaultAllocator, arrowSchema)
+	defer builder.Release()
+	builder.Field(0).(*array.Int64Builder).Append(11)
+	builder.Field(1).(*array.Int64Builder).Append(101)
+	builder.Field(2).(*array.Int64Builder).Append(1)
+	builder.Field(3).(*array.BinaryBuilder).Append([]byte("lob-ref"))
+
+	record := NewSimpleArrowRecord(builder.NewRecord(), map[FieldID]int{
+		common.RowIDField:     0,
+		common.TimeStampField: 1,
+		pkFieldID:             2,
+		textFieldID:           3,
+	})
+	defer record.Release()
+
+	values := make([]*Value, record.Len())
+	err := ValueDeserializerWithSchema(record, values, schema, true)
+	require.NoError(t, err)
+	textValue := values[0].Value.(map[FieldID]interface{})[textFieldID]
+	require.IsType(t, TextLobRef{}, textValue)
+	require.Equal(t, TextLobRef("lob-ref"), textValue)
+
+	rewrittenRecord, err := ValueSerializer(values, schema)
+	require.NoError(t, err)
+	defer rewrittenRecord.Release()
+	textColumn := rewrittenRecord.Column(textFieldID)
+	require.IsType(t, &array.Binary{}, textColumn)
+	require.Equal(t, []byte("lob-ref"), textColumn.(*array.Binary).Value(0))
+}
+
+func TestValueSerializerTextRejectsRawBytes(t *testing.T) {
+	const (
+		pkFieldID   FieldID = 100
+		textFieldID FieldID = 101
+	)
+	schema := &schemapb.CollectionSchema{
+		Name: "text_raw_bytes",
+		Fields: []*schemapb.FieldSchema{
+			{FieldID: common.RowIDField, Name: "row_id", DataType: schemapb.DataType_Int64},
+			{FieldID: common.TimeStampField, Name: "Timestamp", DataType: schemapb.DataType_Int64},
+			{FieldID: pkFieldID, Name: "pk", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
+			{FieldID: textFieldID, Name: "content", DataType: schemapb.DataType_Text},
+		},
+	}
+	values := []*Value{
+		{
+			PK:        NewInt64PrimaryKey(1),
+			Timestamp: 101,
+			Value: map[FieldID]interface{}{
+				common.RowIDField:     int64(11),
+				common.TimeStampField: int64(101),
+				pkFieldID:             int64(1),
+				textFieldID:           []byte("raw-text-bytes"),
+			},
+		},
+	}
+
+	_, err := ValueSerializer(values, schema)
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "expected string value")
+}
+
 type nullableDenseVectorSerdeCase struct {
 	name     string
 	dataType schemapb.DataType

internal/streamingnode/server/flusher/flusherimpl/flusher_components.go

@@ -70,6 +70,8 @@ func (impl *flusherComponents) WhenCreateCollection(createCollectionMsg message.
 			Allocator:          idalloc.NewMAllocator(resource.Resource().IDAllocator()),
 			MsgHandler:         newMsgHandler(resource.Resource().WriteBufferManager()),
 			SchemaManager:      newVersionedSchemaManager(createCollectionMsg.VChannel(), impl.rs),
+			FlushSourceModeNotifier: resource.Resource().SegmentStatsManager().
+				UpdateFlushSourceMode,
 		},
 		msgChan,
 		&datapb.VchannelInfo{
@@ -255,6 +257,8 @@ func (impl *flusherComponents) buildDataSyncService(ctx context.Context, recover
 			Allocator:          idalloc.NewMAllocator(resource.Resource().IDAllocator()),
 			MsgHandler:         newMsgHandler(resource.Resource().WriteBufferManager()),
 			SchemaManager:      newVersionedSchemaManager(recoverInfo.GetInfo().GetChannelName(), impl.rs),
+			FlushSourceModeNotifier: resource.Resource().SegmentStatsManager().
+				UpdateFlushSourceMode,
 		},
 		&datapb.ChannelWatchInfo{Vchan: recoverInfo.GetInfo(), Schema: schema},
 		input,

internal/streamingnode/server/wal/interceptors/shard/shards/segment_manager.go

@@ -149,7 +149,7 @@ func (s *segmentAllocManager) AllocRows(req *AssignSegmentRequest) (*AssignSegme
 		return nil, ErrNotGrowing
 	}
 
-	err := resource.Resource().SegmentStatsManager().AllocRows(s.GetSegmentID(), req.ModifiedMetrics)
+	err := resource.Resource().SegmentStatsManager().AllocRows(s.GetSegmentID(), req.ModifiedMetrics, req.RuntimeFlushSize)
 	if err != nil {
 		return nil, err
 	}

internal/streamingnode/server/wal/interceptors/shard/shards/shard_manager.go

@@ -3,14 +3,18 @@ package shards
 import (
 	"context"
 	"fmt"
+	"math"
+	"strings"
 	"sync"
 
 	"github.com/cockroachdb/errors"
 	"github.com/samber/lo"
 
+	"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
 	"github.com/milvus-io/milvus/internal/streamingnode/server/resource"
 	"github.com/milvus-io/milvus/internal/streamingnode/server/wal"
 	"github.com/milvus-io/milvus/internal/streamingnode/server/wal/interceptors/shard/stats"
+	"github.com/milvus-io/milvus/internal/streamingnode/server/wal/interceptors/shard/utils"
 	"github.com/milvus-io/milvus/internal/streamingnode/server/wal/metricsutil"
 	"github.com/milvus-io/milvus/internal/streamingnode/server/wal/recovery"
 	"github.com/milvus-io/milvus/pkg/v3/common"
@@ -111,6 +115,9 @@ func RecoverShardManager(param *ShardManagerRecoverParam) ShardManager {
 	stats := make([]*stats.SegmentStats, 0, len(belongs))
 	for _, belong := range belongs {
 		stat := m.partitionManagers[belong.PartitionUniqueKey()].segments[belong.SegmentID].GetStatFromRecovery()
+		if info := m.collections[belong.CollectionID]; info != nil {
+			stat.RuntimeFlushSize = info.RuntimeFlushSize(stat.Modified)
+		}
 		stats = append(stats, stat)
 	}
 	resource.Resource().SegmentStatsManager().RegisterSealOperator(m, belongs, stats)
@@ -242,6 +249,105 @@ func (c *CollectionInfo) HasTextField() bool {
 	return typeutil.HasTextField(c.Schema.GetSchema())
 }
 
+// RuntimeFlushSize estimates the in-memory footprint used by flush pressure decisions.
+func (c *CollectionInfo) RuntimeFlushSize(modified stats.ModifiedMetrics) uint64 {
+	if modified.Rows == 0 || modified.BinarySize == 0 {
+		return modified.BinarySize
+	}
+	if !c.shouldEstimateInterimIndexExtra() {
+		return modified.BinarySize
+	}
+
+	extra := estimateInterimIndexExtra(c.Schema.GetSchema(), modified.Rows)
+	if extra == 0 {
+		return modified.BinarySize
+	}
+	return utils.SaturatingAddUint64(modified.BinarySize, extra)
+}
+
+func (c *CollectionInfo) shouldEstimateInterimIndexExtra() bool {
+	if c == nil || c.Schema == nil || c.Schema.GetSchema() == nil || !c.UseGrowingSourceFlush() {
+		return false
+	}
+	params := paramtable.Get()
+	return params.QueryNodeCfg.EnableInterminSegmentIndex.GetAsBool() &&
+		!params.QueryNodeCfg.GrowingMmapEnabled.GetAsBool()
+}
+
+func estimateInterimIndexExtra(schema *schemapb.CollectionSchema, rows uint64) uint64 {
+	var extra uint64
+	for _, field := range schema.GetFields() {
+		switch field.GetDataType() {
+		case schemapb.DataType_FloatVector, schemapb.DataType_Float16Vector, schemapb.DataType_BFloat16Vector:
+			dim, err := typeutil.GetDim(field)
+			if err != nil || dim <= 0 {
+				continue
+			}
+			extra = utils.SaturatingAddUint64(extra, estimateDenseInterimIndexExtra(field.GetDataType(), uint64(dim), rows))
+		case schemapb.DataType_SparseFloatVector:
+			// Sparse interim indexes keep their own representation roughly at
+			// raw sparse-vector size. Modified.BinarySize already accounts for
+			// the raw insert payload, so add one more sparse estimate as index
+			// overhead when chunks are retained for growing-source flush.
+			extra = utils.SaturatingAddUint64(extra, utils.SaturatingMulUint64(rows, uint64(typeutil.GetSparseFloatVectorEstimateLength())))
+		}
+	}
+	return extra
+}
+
+func estimateDenseInterimIndexExtra(dataType schemapb.DataType, dim uint64, rows uint64) uint64 {
+	params := paramtable.Get()
+	indexType := params.QueryNodeCfg.DenseVectorInterminIndexType.GetValue()
+	switch {
+	case strings.EqualFold(indexType, "IVF_FLAT_CC"):
+		rawBytes := utils.SaturatingMulUint64(rows, denseVectorRawBytes(dataType, dim))
+		expansionRate := params.QueryNodeCfg.InterimIndexMemExpandRate.GetAsFloat()
+		if expansionRate <= 0 {
+			expansionRate = 1
+		}
+		return ceilMulFloat(rawBytes, expansionRate)
+	case strings.EqualFold(indexType, "SCANN_DVR"):
+		return utils.SaturatingMulUint64(rows, scannDVRBytesPerRow(dim))
+	default:
+		return 0
+	}
+}
+
+func denseVectorRawBytes(dataType schemapb.DataType, dim uint64) uint64 {
+	switch dataType {
+	case schemapb.DataType_FloatVector:
+		return utils.SaturatingMulUint64(dim, 4)
+	case schemapb.DataType_Float16Vector, schemapb.DataType_BFloat16Vector:
+		return utils.SaturatingMulUint64(dim, 2)
+	default:
+		return 0
+	}
+}
+
+func scannDVRBytesPerRow(dim uint64) uint64 {
+	params := paramtable.Get()
+	subDim := uint64(params.QueryNodeCfg.InterimIndexSubDim.GetAsInt64())
+	bytes := utils.SaturatingMulUint64(subDim/8, dim)
+	switch strings.ToUpper(params.QueryNodeCfg.InterimIndexRefineQuantType.GetValue()) {
+	case "UINT8":
+		bytes = utils.SaturatingAddUint64(bytes, dim)
+	case "FLOAT16", "BFLOAT16":
+		bytes = utils.SaturatingAddUint64(bytes, utils.SaturatingMulUint64(dim, 2))
+	}
+	return bytes
+}
+
+func ceilMulFloat(value uint64, factor float64) uint64 {
+	if value == 0 || factor <= 0 {
+		return 0
+	}
+	result := math.Ceil(float64(value) * factor)
+	if result >= float64(math.MaxUint64) {
+		return math.MaxUint64
+	}
+	return uint64(result)
+}
+
 func (m *shardManagerImpl) Channel() types.PChannelInfo {
 	return m.pchannel
 }

internal/streamingnode/server/wal/interceptors/shard/shards/shard_manager_segment.go

@@ -15,6 +15,7 @@ type AssignSegmentRequest struct {
 	CollectionID          int64
 	PartitionID           int64
 	ModifiedMetrics       stats.ModifiedMetrics
+	RuntimeFlushSize      uint64
 	TimeTick              uint64
 	TxnSession            TxnSession
 	SchemaVersion         int32
@@ -143,6 +144,7 @@ func (m *shardManagerImpl) AssignSegment(req *AssignSegmentRequest) (*AssignSegm
 	if info := m.collections[req.CollectionID]; info != nil {
 		req.SchemaVersion = info.SchemaVersion()
 		req.UseGrowingSourceFlush = info.UseGrowingSourceFlush()
+		req.RuntimeFlushSize = info.RuntimeFlushSize(req.ModifiedMetrics)
 	}
 
 	result, err := pm.AssignSegment(req)

internal/streamingnode/server/wal/interceptors/shard/shards/shard_manager_test.go

@@ -2,6 +2,7 @@ package shards
 
 import (
 	"context"
+	"math"
 	"testing"
 	"time"
 
@@ -10,6 +11,7 @@ import (
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/fieldmaskpb"
 
+	"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
 	"github.com/milvus-io/milvus-proto/go-api/v3/msgpb"
 	"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
 	"github.com/milvus-io/milvus/internal/mocks/streamingnode/server/mock_wal"
@@ -27,6 +29,7 @@ import (
 	"github.com/milvus-io/milvus/pkg/v3/streaming/walimpls/impls/rmq"
 	"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
 	"github.com/milvus-io/milvus/pkg/v3/util/syncutil"
+	"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
 )
 
 func TestShardManager(t *testing.T) {
@@ -878,3 +881,176 @@ func TestCollectionInfoUseGrowingSourceFlush_TextField(t *testing.T) {
 	paramtable.Get().Save(paramtable.Get().CommonCfg.UseLoonFFI.Key, "true")
 	assert.True(t, ci.UseGrowingSourceFlush())
 }
+
+func TestCollectionInfoRuntimeFlushSize(t *testing.T) {
+	paramtable.Init()
+	params := paramtable.Get()
+	params.Save(params.CommonCfg.UseLoonFFI.Key, "true")
+	params.Save(params.CommonCfg.EnableGrowingSourceFlush.Key, "true")
+	params.Save(params.QueryNodeCfg.EnableInterminSegmentIndex.Key, "true")
+	params.Save(params.QueryNodeCfg.GrowingMmapEnabled.Key, "false")
+	params.Save(params.QueryNodeCfg.DenseVectorInterminIndexType.Key, "IVF_FLAT_CC")
+	params.Save(params.QueryNodeCfg.InterimIndexMemExpandRate.Key, "1.25")
+	params.Save(params.QueryNodeCfg.InterimIndexSubDim.Key, "16")
+	params.Save(params.QueryNodeCfg.InterimIndexRefineQuantType.Key, "NONE")
+	defer params.Reset(params.CommonCfg.UseLoonFFI.Key)
+	defer params.Reset(params.CommonCfg.EnableGrowingSourceFlush.Key)
+	defer params.Reset(params.QueryNodeCfg.EnableInterminSegmentIndex.Key)
+	defer params.Reset(params.QueryNodeCfg.GrowingMmapEnabled.Key)
+	defer params.Reset(params.QueryNodeCfg.DenseVectorInterminIndexType.Key)
+	defer params.Reset(params.QueryNodeCfg.InterimIndexMemExpandRate.Key)
+	defer params.Reset(params.QueryNodeCfg.InterimIndexSubDim.Key)
+	defer params.Reset(params.QueryNodeCfg.InterimIndexRefineQuantType.Key)
+
+	metrics := stats.ModifiedMetrics{Rows: 10, BinarySize: 100}
+	vectorInfo := &CollectionInfo{
+		Schema: &streamingpb.CollectionSchemaOfVChannel{
+			Schema: &schemapb.CollectionSchema{
+				Fields: []*schemapb.FieldSchema{
+					{FieldID: 100, DataType: schemapb.DataType_Int64},
+					{
+						FieldID:  101,
+						DataType: schemapb.DataType_FloatVector,
+						TypeParams: []*commonpb.KeyValuePair{
+							{Key: common.DimKey, Value: "128"},
+						},
+					},
+				},
+			},
+		},
+	}
+	assert.Equal(t, uint64(6500), vectorInfo.RuntimeFlushSize(metrics))
+
+	scalarInfo := &CollectionInfo{
+		Schema: &streamingpb.CollectionSchemaOfVChannel{
+			Schema: &schemapb.CollectionSchema{
+				Fields: []*schemapb.FieldSchema{
+					{FieldID: 100, DataType: schemapb.DataType_Int64},
+				},
+			},
+		},
+	}
+	assert.Equal(t, metrics.BinarySize, scalarInfo.RuntimeFlushSize(metrics))
+
+	params.Save(params.QueryNodeCfg.DenseVectorInterminIndexType.Key, "SCANN_DVR")
+	assert.Equal(t, uint64(2660), vectorInfo.RuntimeFlushSize(metrics))
+	params.Save(params.QueryNodeCfg.InterimIndexRefineQuantType.Key, "FLOAT16")
+	assert.Equal(t, uint64(5220), vectorInfo.RuntimeFlushSize(metrics))
+	params.Save(params.QueryNodeCfg.InterimIndexRefineQuantType.Key, "NONE")
+	params.Save(params.QueryNodeCfg.DenseVectorInterminIndexType.Key, "IVF_FLAT_CC")
+
+	sparseInfo := &CollectionInfo{
+		Schema: &streamingpb.CollectionSchemaOfVChannel{
+			Schema: &schemapb.CollectionSchema{
+				Fields: []*schemapb.FieldSchema{
+					{FieldID: 100, DataType: schemapb.DataType_Int64},
+					{FieldID: 101, DataType: schemapb.DataType_SparseFloatVector},
+				},
+			},
+		},
+	}
+	assert.Equal(t, metrics.BinarySize+metrics.Rows*uint64(typeutil.GetSparseFloatVectorEstimateLength()), sparseInfo.RuntimeFlushSize(metrics))
+
+	params.Save(params.QueryNodeCfg.GrowingMmapEnabled.Key, "true")
+	assert.Equal(t, metrics.BinarySize, vectorInfo.RuntimeFlushSize(metrics))
+	assert.Equal(t, metrics.BinarySize, sparseInfo.RuntimeFlushSize(metrics))
+	params.Save(params.QueryNodeCfg.GrowingMmapEnabled.Key, "false")
+
+	params.Save(params.QueryNodeCfg.EnableInterminSegmentIndex.Key, "false")
+	assert.Equal(t, metrics.BinarySize, vectorInfo.RuntimeFlushSize(metrics))
+	params.Save(params.QueryNodeCfg.EnableInterminSegmentIndex.Key, "true")
+
+	params.Save(params.CommonCfg.EnableGrowingSourceFlush.Key, "false")
+	assert.Equal(t, metrics.BinarySize, vectorInfo.RuntimeFlushSize(metrics))
+}
+
+func TestCollectionInfoRuntimeFlushSizeEdgeCases(t *testing.T) {
+	paramtable.Init()
+	params := paramtable.Get()
+	params.Save(params.CommonCfg.UseLoonFFI.Key, "true")
+	params.Save(params.CommonCfg.EnableGrowingSourceFlush.Key, "true")
+	params.Save(params.QueryNodeCfg.EnableInterminSegmentIndex.Key, "true")
+	params.Save(params.QueryNodeCfg.GrowingMmapEnabled.Key, "false")
+	params.Save(params.QueryNodeCfg.DenseVectorInterminIndexType.Key, "IVF_FLAT_CC")
+	params.Save(params.QueryNodeCfg.InterimIndexMemExpandRate.Key, "2")
+	defer params.Reset(params.CommonCfg.UseLoonFFI.Key)
+	defer params.Reset(params.CommonCfg.EnableGrowingSourceFlush.Key)
+	defer params.Reset(params.QueryNodeCfg.EnableInterminSegmentIndex.Key)
+	defer params.Reset(params.QueryNodeCfg.GrowingMmapEnabled.Key)
+	defer params.Reset(params.QueryNodeCfg.DenseVectorInterminIndexType.Key)
+	defer params.Reset(params.QueryNodeCfg.InterimIndexMemExpandRate.Key)
+
+	vectorInfo := &CollectionInfo{
+		Schema: &streamingpb.CollectionSchemaOfVChannel{
+			Schema: &schemapb.CollectionSchema{
+				Fields: []*schemapb.FieldSchema{
+					{FieldID: 100, DataType: schemapb.DataType_Int64},
+					{
+						FieldID:  101,
+						DataType: schemapb.DataType_FloatVector,
+						TypeParams: []*commonpb.KeyValuePair{
+							{Key: common.DimKey, Value: "128"},
+						},
+					},
+				},
+			},
+		},
+	}
+	assert.Equal(t, uint64(100), vectorInfo.RuntimeFlushSize(stats.ModifiedMetrics{Rows: 0, BinarySize: 100}))
+	assert.Equal(t, uint64(0), vectorInfo.RuntimeFlushSize(stats.ModifiedMetrics{Rows: 10, BinarySize: 0}))
+	assert.Equal(t, uint64(math.MaxUint64), vectorInfo.RuntimeFlushSize(stats.ModifiedMetrics{Rows: math.MaxUint64, BinarySize: math.MaxUint64 - 1}))
+
+	binaryVectorInfo := &CollectionInfo{
+		Schema: &streamingpb.CollectionSchemaOfVChannel{
+			Schema: &schemapb.CollectionSchema{
+				Fields: []*schemapb.FieldSchema{
+					{FieldID: 100, DataType: schemapb.DataType_Int64},
+					{
+						FieldID:  101,
+						DataType: schemapb.DataType_BinaryVector,
+						TypeParams: []*commonpb.KeyValuePair{
+							{Key: common.DimKey, Value: "128"},
+						},
+					},
+				},
+			},
+		},
+	}
+	assert.Equal(t, uint64(100), binaryVectorInfo.RuntimeFlushSize(stats.ModifiedMetrics{Rows: 10, BinarySize: 100}))
+
+	int8VectorInfo := &CollectionInfo{
+		Schema: &streamingpb.CollectionSchemaOfVChannel{
+			Schema: &schemapb.CollectionSchema{
+				Fields: []*schemapb.FieldSchema{
+					{FieldID: 100, DataType: schemapb.DataType_Int64},
+					{
+						FieldID:  101,
+						DataType: schemapb.DataType_Int8Vector,
+						TypeParams: []*commonpb.KeyValuePair{
+							{Key: common.DimKey, Value: "128"},
+						},
+					},
+				},
+			},
+		},
+	}
+	assert.Equal(t, uint64(100), int8VectorInfo.RuntimeFlushSize(stats.ModifiedMetrics{Rows: 10, BinarySize: 100}))
+
+	invalidDimInfo := &CollectionInfo{
+		Schema: &streamingpb.CollectionSchemaOfVChannel{
+			Schema: &schemapb.CollectionSchema{
+				Fields: []*schemapb.FieldSchema{
+					{FieldID: 100, DataType: schemapb.DataType_Int64},
+					{
+						FieldID:  101,
+						DataType: schemapb.DataType_FloatVector,
+						TypeParams: []*commonpb.KeyValuePair{
+							{Key: common.DimKey, Value: "bad"},
+						},
+					},
+				},
+			},
+		},
+	}
+	assert.Equal(t, uint64(100), invalidDimInfo.RuntimeFlushSize(stats.ModifiedMetrics{Rows: 10, BinarySize: 100}))
+}

internal/streamingnode/server/wal/interceptors/shard/stats/metrics.go

@@ -11,19 +11,21 @@ import (
 // newMetricsHelper creates a new metrics helper for the WAL segment.
 func newMetricsHelper() *metricsHelper {
 	return &metricsHelper{
-		growingBytesHWM:  metrics.WALGrowingSegmentHWMBytes.With(prometheus.Labels{metrics.NodeIDLabelName: paramtable.GetStringNodeID()}),
-		growingBytesLWM:  metrics.WALGrowingSegmentLWMBytes.With(prometheus.Labels{metrics.NodeIDLabelName: paramtable.GetStringNodeID()}),
-		growingBytes:     metrics.WALGrowingSegmentBytes.MustCurryWith(prometheus.Labels{metrics.NodeIDLabelName: paramtable.GetStringNodeID()}),
-		growingRowsTotal: metrics.WALGrowingSegmentRowsTotal.MustCurryWith(prometheus.Labels{metrics.NodeIDLabelName: paramtable.GetStringNodeID()}),
+		growingBytesHWM:    metrics.WALGrowingSegmentHWMBytes.With(prometheus.Labels{metrics.NodeIDLabelName: paramtable.GetStringNodeID()}),
+		growingBytesLWM:    metrics.WALGrowingSegmentLWMBytes.With(prometheus.Labels{metrics.NodeIDLabelName: paramtable.GetStringNodeID()}),
+		flushPressureBytes: metrics.WALGrowingSegmentFlushPressureBytes.With(prometheus.Labels{metrics.NodeIDLabelName: paramtable.GetStringNodeID()}),
+		growingBytes:       metrics.WALGrowingSegmentBytes.MustCurryWith(prometheus.Labels{metrics.NodeIDLabelName: paramtable.GetStringNodeID()}),
+		growingRowsTotal:   metrics.WALGrowingSegmentRowsTotal.MustCurryWith(prometheus.Labels{metrics.NodeIDLabelName: paramtable.GetStringNodeID()}),
 	}
 }
 
 // metricsHelper is a helper struct for managing metrics related to WAL segments.
 type metricsHelper struct {
-	growingBytesHWM  prometheus.Gauge
-	growingBytesLWM  prometheus.Gauge
-	growingBytes     *prometheus.GaugeVec
-	growingRowsTotal *prometheus.GaugeVec
+	growingBytesHWM    prometheus.Gauge
+	growingBytesLWM    prometheus.Gauge
+	flushPressureBytes prometheus.Gauge
+	growingBytes       *prometheus.GaugeVec
+	growingRowsTotal   *prometheus.GaugeVec
 }
 
 // ObservePChannelBytesUpdate updates the bytes of a pchannel.
@@ -43,6 +45,11 @@ func (m *metricsHelper) ObservePChannelBytesUpdate(pchannel string, am *aggregat
 	}
 }
 
+// ObserveFlushPressureBytesUpdate updates the runtime bytes used by HWM/LWM flush decisions.
+func (m *metricsHelper) ObserveFlushPressureBytesUpdate(bytes uint64) {
+	m.flushPressureBytes.Set(float64(bytes))
+}
+
 // ObserveConfigUpdate is a update method for configuration changes.
 func (m *metricsHelper) ObserveConfigUpdate(cfg statsConfig) {
 	m.growingBytesHWM.Set(float64(cfg.growingBytesHWM))

internal/streamingnode/server/wal/interceptors/shard/stats/stats_manager.go

@@ -8,6 +8,7 @@ import (
 
 	"github.com/cockroachdb/errors"
 
+	"github.com/milvus-io/milvus/internal/flushcommon/metacache"
 	"github.com/milvus-io/milvus/internal/streamingnode/server/wal/interceptors/shard/policy"
 	"github.com/milvus-io/milvus/internal/streamingnode/server/wal/interceptors/shard/utils"
 	"github.com/milvus-io/milvus/pkg/v3/mlog"
@@ -37,10 +38,12 @@ type StatsManager struct {
 	mu                         sync.Mutex
 	cfg                        statsConfig
 	totalStats                 *aggregatedMetrics
+	totalFlushSize             uint64
 	pchannelStats              map[string]*aggregatedMetrics
 	vchannelStats              map[string]*aggregatedMetrics
-	segmentStats               map[int64]*SegmentStats       // map[SegmentID]SegmentStats
-	segmentIndex               map[int64]SegmentBelongs      // map[SegmentID]channels
+	segmentStats               map[int64]*SegmentStats  // map[SegmentID]SegmentStats
+	segmentIndex               map[int64]SegmentBelongs // map[SegmentID]channels
+	segmentFlushSourceModes    map[int64]metacache.FlushSourceMode
 	pchannelIndex              map[string]map[int64]struct{} // map[PChannel]SegmentID
 	growingL1SegmentsByChannel map[channelKey]map[int64]struct{}
 	segmentDeletePressures     map[int64]deletePressure // map[SegmentID]aggregated delete pressure
@@ -78,6 +81,7 @@ func NewStatsManager() *StatsManager {
 		vchannelStats:              make(map[string]*aggregatedMetrics),
 		segmentStats:               make(map[int64]*SegmentStats),
 		segmentIndex:               make(map[int64]SegmentBelongs),
+		segmentFlushSourceModes:    make(map[int64]metacache.FlushSourceMode),
 		pchannelIndex:              make(map[string]map[int64]struct{}),
 		growingL1SegmentsByChannel: make(map[channelKey]map[int64]struct{}),
 		segmentDeletePressures:     make(map[int64]deletePressure),
@@ -85,6 +89,7 @@ func NewStatsManager() *StatsManager {
 		metricHelper:               newMetricsHelper(),
 	}
 	m.worker = newSealWorker(m)
+	m.metricHelper.ObserveFlushPressureBytesUpdate(m.totalFlushSize)
 	go m.worker.loop()
 	return m
 }
@@ -163,6 +168,9 @@ func (m *StatsManager) registerNewGrowingSegment(belongs SegmentBelongs, stats *
 	if _, ok := m.segmentStats[segmentID]; ok {
 		panic(fmt.Sprintf("register a segment %d that already exist, critical bug", segmentID))
 	}
+	if stats.RuntimeFlushSize == 0 {
+		stats.RuntimeFlushSize = stats.Modified.BinarySize
+	}
 
 	m.segmentStats[segmentID] = stats
 	m.segmentIndex[segmentID] = belongs
@@ -178,6 +186,8 @@ func (m *StatsManager) registerNewGrowingSegment(belongs SegmentBelongs, stats *
 		m.growingL1SegmentsByChannel[key][segmentID] = struct{}{}
 	}
 	m.totalStats.Collect(stats.Level, stats.Modified)
+	m.totalFlushSize = utils.SaturatingAddUint64(m.totalFlushSize, stats.FlushSize())
+	m.metricHelper.ObserveFlushPressureBytesUpdate(m.totalFlushSize)
 	if _, ok := m.pchannelStats[belongs.PChannel]; !ok {
 		m.pchannelStats[belongs.PChannel] = newAggregatedMetrics()
 	}
@@ -194,12 +204,12 @@ func (m *StatsManager) registerNewGrowingSegment(belongs SegmentBelongs, stats *
 // AllocRows alloc number of rows on current segment.
 // AllocRows will check if the segment has enough space to insert.
 // Must be called after RegisterGrowingSegment and before UnregisterGrowingSegment.
-func (m *StatsManager) AllocRows(segmentID int64, insert ModifiedMetrics) error {
+func (m *StatsManager) AllocRows(segmentID int64, insert ModifiedMetrics, runtimeFlushSize ...uint64) error {
 	if insert.Rows == 0 || insert.BinarySize == 0 {
 		panic(fmt.Sprintf("insert rows or binary size cannot be 0, rows: %d, binary: %d", insert.Rows, insert.BinarySize))
 	}
 
-	shouldBeSealed, err := m.allocRows(segmentID, insert)
+	shouldBeSealed, err := m.allocRows(segmentID, insert, normalizeRuntimeFlushSize(insert, runtimeFlushSize...))
 	if shouldBeSealed {
 		m.worker.NotifySealSegment(segmentID, policy.PolicyCapacity())
 	}
@@ -210,8 +220,15 @@ func (m *StatsManager) AllocRows(segmentID int64, insert ModifiedMetrics) error
 	return nil
 }
 
+func normalizeRuntimeFlushSize(insert ModifiedMetrics, runtimeFlushSize ...uint64) uint64 {
+	if len(runtimeFlushSize) == 0 || runtimeFlushSize[0] < insert.BinarySize {
+		return insert.BinarySize
+	}
+	return runtimeFlushSize[0]
+}
+
 // allocRows allocates number of rows on current segment.
-func (m *StatsManager) allocRows(segmentID int64, insert ModifiedMetrics) (bool, error) {
+func (m *StatsManager) allocRows(segmentID int64, insert ModifiedMetrics, runtimeFlushSize uint64) (bool, error) {
 	m.mu.Lock()
 	defer m.mu.Unlock()
 
@@ -225,6 +242,12 @@ func (m *StatsManager) allocRows(segmentID int64, insert ModifiedMetrics) (bool,
 
 	// update the total stats if inserted.
 	if inserted {
+		if m.segmentFlushSourceModes[segmentID] == metacache.FlushSourceWriteBuffer {
+			runtimeFlushSize = insert.BinarySize
+		}
+		stat.AllocRuntimeFlushSize(runtimeFlushSize)
+		m.totalFlushSize = utils.SaturatingAddUint64(m.totalFlushSize, runtimeFlushSize)
+		m.metricHelper.ObserveFlushPressureBytesUpdate(m.totalFlushSize)
 		m.totalStats.Collect(stat.Level, insert)
 		if _, ok := m.pchannelStats[info.PChannel]; !ok {
 			m.pchannelStats[info.PChannel] = newAggregatedMetrics()
@@ -244,6 +267,37 @@ func (m *StatsManager) allocRows(segmentID int64, insert ModifiedMetrics) (bool,
 	return stat.ShouldBeSealed(), ErrNotEnoughSpace
 }
 
+// UpdateFlushSourceMode records the segment-level sticky flush source chosen by writebuffer.
+func (m *StatsManager) UpdateFlushSourceMode(segmentID int64, mode metacache.FlushSourceMode) {
+	if mode != metacache.FlushSourceWriteBuffer && mode != metacache.FlushSourceGrowing {
+		return
+	}
+
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	stat, ok := m.segmentStats[segmentID]
+	if !ok {
+		return
+	}
+	current := m.segmentFlushSourceModes[segmentID]
+	if current != metacache.FlushSourceUnknown && current != mode {
+		return
+	}
+	m.segmentFlushSourceModes[segmentID] = mode
+	if mode == metacache.FlushSourceWriteBuffer {
+		oldFlushSize := stat.FlushSize()
+		stat.RuntimeFlushSize = stat.Modified.BinarySize
+		newFlushSize := stat.FlushSize()
+		if oldFlushSize > newFlushSize {
+			m.totalFlushSize = utils.SaturatingSubUint64(m.totalFlushSize, oldFlushSize-newFlushSize)
+		} else if newFlushSize > oldFlushSize {
+			m.totalFlushSize = utils.SaturatingAddUint64(m.totalFlushSize, newFlushSize-oldFlushSize)
+		}
+		m.metricHelper.ObserveFlushPressureBytesUpdate(m.totalFlushSize)
+	}
+}
+
 // RecordDelete records local delete metrics on matching growing L1 segments.
 func (m *StatsManager) RecordDelete(pchannel string, vchannel string, timeTick uint64, rows uint64, bytes uint64) {
 	if pchannel == "" || vchannel == "" || timeTick == 0 || (rows == 0 && bytes == 0) {
@@ -269,7 +323,7 @@ func (m *StatsManager) RecordDelete(pchannel string, vchannel string, timeTick u
 // notifyIfTotalGrowingBytesOverHWM notifies if the total bytes is over the high water mark.
 func (m *StatsManager) notifyIfTotalGrowingBytesOverHWM() {
 	m.mu.Lock()
-	size := m.totalStats.Total().BinarySize
+	size := m.totalFlushSize
 	notify := size > uint64(m.cfg.growingBytesHWM)
 	m.mu.Unlock()
 
@@ -344,8 +398,11 @@ func (m *StatsManager) unregisterSealedSegment(segmentID int64) *SegmentStats {
 	stats := m.segmentStats[segmentID]
 
 	m.totalStats.Subtract(stats.Level, stats.Modified)
+	m.totalFlushSize = utils.SaturatingSubUint64(m.totalFlushSize, stats.FlushSize())
+	m.metricHelper.ObserveFlushPressureBytesUpdate(m.totalFlushSize)
 	delete(m.segmentStats, segmentID)
 	delete(m.segmentIndex, segmentID)
+	delete(m.segmentFlushSourceModes, segmentID)
 	delete(m.segmentDeletePressures, segmentID)
 	if stats.Level == datapb.SegmentLevel_L1 {
 		key := channelKey{pchannel: info.PChannel, vchannel: info.VChannel}
@@ -470,7 +527,7 @@ func (m *StatsManager) reachBlockingL0Threshold(rows uint64, bytes uint64) bool
 // selectSegmentsUntilLessThanLWM selects segments until the total size is less than the threshold.
 func (m *StatsManager) selectSegmentsUntilLessThanLWM() []int64 {
 	m.mu.Lock()
-	restSpace := int64(m.totalStats.Total().BinarySize) - m.cfg.growingBytesLWM
+	restSpace := utils.SaturatingUint64ToInt64(m.totalFlushSize) - m.cfg.growingBytesLWM
 	m.mu.Unlock()
 
 	if restSpace <= 0 {
@@ -497,10 +554,11 @@ func (m *StatsManager) createStatsSlice() []segmentWithBinarySize {
 
 	stats := make([]segmentWithBinarySize, 0, len(m.segmentStats))
 	for id, stat := range m.segmentStats {
-		if stat.Modified.BinarySize > 0 {
+		flushSize := stat.FlushSize()
+		if flushSize > 0 {
 			stats = append(stats, segmentWithBinarySize{
 				segmentID:  id,
-				binarySize: stat.Modified.BinarySize,
+				binarySize: flushSize,
 			})
 		}
 	}

internal/streamingnode/server/wal/interceptors/shard/stats/stats_manager_test.go

@@ -12,6 +12,7 @@ import (
 	"github.com/stretchr/testify/require"
 	"go.uber.org/atomic"
 
+	"github.com/milvus-io/milvus/internal/flushcommon/metacache"
 	"github.com/milvus-io/milvus/internal/mocks/streamingnode/server/wal/interceptors/shard/mock_utils"
 	"github.com/milvus-io/milvus/internal/streamingnode/server/wal/interceptors/shard/policy"
 	"github.com/milvus-io/milvus/internal/streamingnode/server/wal/interceptors/shard/utils"
@@ -148,6 +149,86 @@ func TestStatsManager(t *testing.T) {
 	assert.Empty(t, m.sealOperators)
 }
 
+func TestStatsManagerRuntimeFlushSizeForMemoryPressure(t *testing.T) {
+	paramtable.Init()
+	m := NewStatsManager()
+	m.cfg.growingBytesHWM = 10_000
+	m.cfg.growingBytesLWM = 250
+
+	sealOperator := mock_utils.NewMockSealOperator(t)
+	sealOperator.EXPECT().Channel().Return(types.PChannelInfo{Name: "pchannel"})
+	sealOperator.EXPECT().AsyncFlushSegment(mock.Anything).Return().Maybe()
+	m.RegisterSealOperator(sealOperator, nil, nil)
+
+	m.RegisterNewGrowingSegment(SegmentBelongs{PChannel: "pchannel", VChannel: "vchannel", CollectionID: 1, PartitionID: 1, SegmentID: 1}, createSegmentStats(10, 100, 1_000))
+	m.RegisterNewGrowingSegment(SegmentBelongs{PChannel: "pchannel", VChannel: "vchannel", CollectionID: 1, PartitionID: 1, SegmentID: 2}, createSegmentStats(10, 100, 1_000))
+
+	require.NoError(t, m.AllocRows(1, ModifiedMetrics{Rows: 5, BinarySize: 50}, 300))
+	require.NoError(t, m.AllocRows(2, ModifiedMetrics{Rows: 5, BinarySize: 50}, 20))
+
+	stat1 := m.GetStatsOfSegment(1)
+	assert.Equal(t, uint64(150), stat1.Modified.BinarySize)
+	assert.Equal(t, uint64(400), stat1.FlushSize())
+
+	stat2 := m.GetStatsOfSegment(2)
+	assert.Equal(t, uint64(150), stat2.Modified.BinarySize)
+	assert.Equal(t, uint64(150), stat2.FlushSize())
+
+	assert.Equal(t, uint64(300), m.totalStats.Total().BinarySize)
+	assert.Equal(t, uint64(550), m.totalFlushSize)
+	assert.Equal(t, []int64{1}, m.selectSegmentsUntilLessThanLWM())
+
+	m.UpdateFlushSourceMode(1, metacache.FlushSourceWriteBuffer)
+	stat1 = m.GetStatsOfSegment(1)
+	assert.Equal(t, uint64(150), stat1.FlushSize())
+	assert.Equal(t, uint64(300), m.totalFlushSize)
+	require.NoError(t, m.AllocRows(1, ModifiedMetrics{Rows: 5, BinarySize: 50}, 300))
+	stat1 = m.GetStatsOfSegment(1)
+	assert.Equal(t, uint64(200), stat1.Modified.BinarySize)
+	assert.Equal(t, uint64(200), stat1.FlushSize())
+	assert.Equal(t, uint64(350), m.totalFlushSize)
+
+	m.UpdateFlushSourceMode(2, metacache.FlushSourceGrowing)
+	require.NoError(t, m.AllocRows(2, ModifiedMetrics{Rows: 5, BinarySize: 50}, 300))
+	stat2 = m.GetStatsOfSegment(2)
+	assert.Equal(t, uint64(200), stat2.Modified.BinarySize)
+	assert.Equal(t, uint64(450), stat2.FlushSize())
+	assert.Equal(t, uint64(650), m.totalFlushSize)
+}
+
+func TestStatsManagerRuntimeFlushSizeUnregisterAndModeCorrection(t *testing.T) {
+	paramtable.Init()
+	m := NewStatsManager()
+
+	sealOperator := mock_utils.NewMockSealOperator(t)
+	sealOperator.EXPECT().Channel().Return(types.PChannelInfo{Name: "pchannel"})
+	sealOperator.EXPECT().AsyncFlushSegment(mock.Anything).Return().Maybe()
+	m.RegisterSealOperator(sealOperator, nil, nil)
+
+	stat := createSegmentStats(10, 100, 1_000)
+	stat.RuntimeFlushSize = 500
+	m.RegisterNewGrowingSegment(SegmentBelongs{
+		PChannel:     "pchannel",
+		VChannel:     "vchannel",
+		CollectionID: 1,
+		PartitionID:  1,
+		SegmentID:    10,
+	}, stat)
+	assert.Equal(t, uint64(500), m.totalFlushSize)
+
+	m.UpdateFlushSourceMode(10, metacache.FlushSourceWriteBuffer)
+	assert.Equal(t, uint64(100), m.totalFlushSize)
+	assert.Equal(t, uint64(100), m.GetStatsOfSegment(10).FlushSize())
+
+	m.UpdateFlushSourceMode(10, metacache.FlushSourceGrowing)
+	assert.Equal(t, uint64(100), m.totalFlushSize)
+	assert.Equal(t, uint64(100), m.GetStatsOfSegment(10).FlushSize())
+
+	m.UnregisterSealedSegment(10)
+	assert.Equal(t, uint64(0), m.totalFlushSize)
+	assert.Empty(t, m.segmentFlushSourceModes)
+}
+
 func TestConcurrentStasManager(t *testing.T) {
 	paramtable.Init()
 	params := paramtable.Get()

internal/streamingnode/server/wal/interceptors/shard/utils/math.go

@@ -0,0 +1,35 @@
+package utils
+
+import "math"
+
+// SaturatingAddUint64 adds two uint64 values and clamps overflow to MaxUint64.
+func SaturatingAddUint64(a uint64, b uint64) uint64 {
+	if b > math.MaxUint64-a {
+		return math.MaxUint64
+	}
+	return a + b
+}
+
+// SaturatingSubUint64 subtracts b from a and clamps underflow to zero.
+func SaturatingSubUint64(a uint64, b uint64) uint64 {
+	if b > a {
+		return 0
+	}
+	return a - b
+}
+
+// SaturatingMulUint64 multiplies two uint64 values and clamps overflow to MaxUint64.
+func SaturatingMulUint64(a uint64, b uint64) uint64 {
+	if a != 0 && b > math.MaxUint64/a {
+		return math.MaxUint64
+	}
+	return a * b
+}
+
+// SaturatingUint64ToInt64 converts uint64 to int64 and clamps overflow to MaxInt64.
+func SaturatingUint64ToInt64(value uint64) int64 {
+	if value > uint64(math.MaxInt64) {
+		return math.MaxInt64
+	}
+	return int64(value)
+}

internal/streamingnode/server/wal/interceptors/shard/utils/stats.go

@@ -41,6 +41,7 @@ func (s *SegmentBelongs) PartitionUniqueKey() PartitionUniqueKey {
 // SegmentStats is the usage stats of a segment.
 type SegmentStats struct {
 	Modified              ModifiedMetrics
+	RuntimeFlushSize      uint64    // runtime-only size used by StreamingNode flush HWM/LWM decisions; not persisted into recovery meta.
 	MaxRows               uint64    // MaxRows of current segment should be assigned, it's a fixed value when segment is transfer int growing.
 	MaxBinarySize         uint64    // MaxBinarySize of current segment should be assigned, it's a fixed value when segment is transfer int growing.
 	CreateTime            time.Time // created timestamp of this segment, it's a fixed value when segment is created, not a tso.
@@ -117,6 +118,23 @@ func (s *SegmentStats) AllocRows(m ModifiedMetrics) bool {
 	return true
 }
 
+// AllocRuntimeFlushSize records runtime-only size growth for flush HWM/LWM decisions.
+func (s *SegmentStats) AllocRuntimeFlushSize(size uint64) {
+	if size > math.MaxUint64-s.RuntimeFlushSize {
+		s.RuntimeFlushSize = math.MaxUint64
+		return
+	}
+	s.RuntimeFlushSize += size
+}
+
+// FlushSize returns the size used by runtime flush decisions.
+func (s *SegmentStats) FlushSize() uint64 {
+	if s.RuntimeFlushSize > 0 {
+		return s.RuntimeFlushSize
+	}
+	return s.Modified.BinarySize
+}
+
 // BinaryCanBeAssign returns the capacity of binary size can be inserted.
 func (s *SegmentStats) BinaryCanBeAssign() uint64 {
 	return s.MaxBinarySize - s.Modified.BinarySize

pkg/metrics/streaming_service_metrics.go

@@ -279,6 +279,11 @@ var (
 		Help: "Bytes of segment growing on wal",
 	}, WALChannelLabelName, WALSegmentLevelLabelName)
 
+	WALGrowingSegmentFlushPressureBytes = newWALGaugeVec(prometheus.GaugeOpts{
+		Name: "growing_segment_flush_pressure_bytes",
+		Help: "Runtime bytes used by WAL growing segment flush pressure decisions",
+	})
+
 	WALGrowingSegmentHWMBytes = newWALGaugeVec(prometheus.GaugeOpts{
 		Name: "growing_segment_hwm_bytes",
 		Help: "HWM of segment growing bytes on node",
@@ -642,6 +647,7 @@ func registerWAL(registry *prometheus.Registry) {
 	registry.MustRegister(WALInsertBytes)
 	registry.MustRegister(WALDeleteRowsTotal)
 	registry.MustRegister(WALGrowingSegmentBytes)
+	registry.MustRegister(WALGrowingSegmentFlushPressureBytes)
 	registry.MustRegister(WALGrowingSegmentRowsTotal)
 	registry.MustRegister(WALGrowingSegmentHWMBytes)
 	registry.MustRegister(WALGrowingSegmentLWMBytes)