enhance: add delete rows into delete msg header and more metric

- add delete rows into delete messsage header
- add more insert/delete metrics
- fix non-broadcast message has broadcast header

Signed-off-by: chyezh <[email protected]>

Author: chyezh

internal/flushcommon/pipeline/flow_graph_dd_node.go

@@ -274,6 +274,7 @@ func (ddn *ddNode) Operate(in []Msg) []Msg {
 				zap.Int32("msgType", int32(msg.Type())),
 				zap.Uint64("timetick", manualFlushMsg.ManualFlushMessage.TimeTick()),
 				zap.Uint64("flushTs", manualFlushMsg.ManualFlushMessage.Header().FlushTs),
+				zap.Int64s("segmentIDs", manualFlushMsg.ManualFlushMessage.Header().SegmentIds),
 			)
 			logger.Info("receive manual flush message")
 			if err := ddn.msgHandler.HandleManualFlush(manualFlushMsg.ManualFlushMessage); err != nil {
@@ -302,7 +303,12 @@ func (ddn *ddNode) Operate(in []Msg) []Msg {
 			if header.GetCollectionId() != ddn.collectionID {
 				continue
 			}
-			logger := log.With(zap.String("vchannel", ddn.Name()))
+			logger := log.With(
+				zap.String("vchannel", ddn.Name()),
+				zap.Int32("msgType", int32(msg.Type())),
+				zap.Uint64("timetick", schemaMsg.SchemaChangeMessage.TimeTick()),
+				zap.Int64s("segmentIDs", schemaMsg.SchemaChangeMessage.Header().FlushedSegmentIds),
+			)
 			logger.Info("receive schema change message")
 			body, err := schemaMsg.SchemaChangeMessage.Body()
 			if err != nil {

internal/mocks/streamingnode/server/wal/interceptors/shard/mock_shards/mock_ShardManager.go

@@ -49,6 +49,7 @@ func (dt *deleteTaskByStreamingService) Execute(ctx context.Context) (err error)
 			msg, err := message.NewDeleteMessageBuilderV1().
 				WithHeader(&message.DeleteMessageHeader{
 					CollectionId: dt.collectionID,
+					Rows:         uint64(deleteMsg.NumRows),
 				}).
 				WithBody(deleteMsg.DeleteRequest).
 				WithVChannel(vchannel).

internal/proxy/task_delete_streaming.go

@@ -22,11 +22,11 @@ import (
 func TestBroadcast(t *testing.T) {
 	s := newMockServer(t, 0)
 	bs := NewGRPCBroadcastService(walimplstest.WALName, s)
-	msg, _ := message.NewDropCollectionMessageBuilderV1().
+	msg := message.NewDropCollectionMessageBuilderV1().
 		WithHeader(&message.DropCollectionMessageHeader{}).
 		WithBody(&msgpb.DropCollectionRequest{}).
 		WithBroadcast([]string{"v1"}, message.NewCollectionNameResourceKey("r1")).
-		BuildBroadcast()
+		MustBuildBroadcast()
 	_, err := bs.Broadcast(context.Background(), msg)
 	assert.NoError(t, err)
 	err = bs.Ack(context.Background(), types.BroadcastAckRequest{

internal/streamingcoord/client/broadcast/broadcast_test.go

@@ -38,15 +38,16 @@ func TestFlushMsgHandler_HandleFlush(t *testing.T) {
 	wbMgr := writebuffer.NewMockBufferManager(t)
 	wbMgr.EXPECT().SealSegments(mock.Anything, mock.Anything, mock.Anything).Return(errors.New("mock err"))
 
-	msg, err := message.NewFlushMessageBuilderV2().
-		WithVChannel(vchannel).
+	msg := message.NewFlushMessageBuilderV2().
+		WithBroadcast([]string{vchannel}).
 		WithHeader(&message.FlushMessageHeader{
 			CollectionId: 0,
 			SegmentId:    1,
 		}).
 		WithBody(&message.FlushMessageBody{}).
-		BuildMutable()
-	assert.NoError(t, err)
+		MustBuildBroadcast().
+		WithBroadcastID(1).
+		SplitIntoMutableMessage()[0]
 
 	handler := newMsgHandler(wbMgr)
 	msgID := mock_message.NewMockMessageID(t)
@@ -72,15 +73,16 @@ func TestFlushMsgHandler_HandleManualFlush(t *testing.T) {
 	wbMgr.EXPECT().SealSegments(mock.Anything, mock.Anything, mock.Anything).Return(errors.New("mock err"))
 	wbMgr.EXPECT().FlushChannel(mock.Anything, mock.Anything, mock.Anything).Return(errors.New("mock err"))
 
-	msg, err := message.NewManualFlushMessageBuilderV2().
-		WithVChannel(vchannel).
+	msg := message.NewManualFlushMessageBuilderV2().
+		WithBroadcast([]string{vchannel}).
 		WithHeader(&message.ManualFlushMessageHeader{
 			CollectionId: 0,
 			FlushTs:      1000,
 		}).
 		WithBody(&message.ManualFlushMessageBody{}).
-		BuildMutable()
-	assert.NoError(t, err)
+		MustBuildBroadcast().
+		WithBroadcastID(1).
+		SplitIntoMutableMessage()[0]
 
 	handler := newMsgHandler(wbMgr)
 	msgID := mock_message.NewMockMessageID(t)
@@ -115,20 +117,21 @@ func TestFlushMsgHandler_HandlSchemaChange(t *testing.T) {
 	wbMgr := writebuffer.NewMockBufferManager(t)
 	wbMgr.EXPECT().SealSegments(mock.Anything, mock.Anything, mock.Anything).Return(errors.New("mock err"))
 
-	msg, err := message.NewSchemaChangeMessageBuilderV2().
-		WithVChannel(vchannel).
+	msg := message.NewSchemaChangeMessageBuilderV2().
+		WithBroadcast([]string{vchannel}).
 		WithHeader(&message.SchemaChangeMessageHeader{
 			CollectionId:      0,
 			FlushedSegmentIds: []int64{1},
 		}).
 		WithBody(&message.SchemaChangeMessageBody{}).
-		BuildMutable()
-	assert.NoError(t, err)
+		MustBuildBroadcast().
+		WithBroadcastID(1).
+		SplitIntoMutableMessage()[0]
 
 	handler := newMsgHandler(wbMgr)
 	msgID := mock_message.NewMockMessageID(t)
 	im := message.MustAsImmutableCollectionSchemaChangeV2(msg.IntoImmutableMessage(msgID))
-	err = handler.HandleSchemaChange(context.Background(), im)
+	err := handler.HandleSchemaChange(context.Background(), im)
 	assert.Error(t, err)
 
 	// test normal

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

@@ -42,7 +42,7 @@ func newOldVersionImmutableMessage(
 	case *msgstream.InsertMsg:
 		mutableMessage = newV1InsertMsgFromV0(underlyingMsg, uint64(len(msg.Payload())))
 	case *msgstream.DeleteMsg:
-		mutableMessage = newV1DeleteMsgFromV0(underlyingMsg)
+		mutableMessage = newV1DeleteMsgFromV0(underlyingMsg, uint64(underlyingMsg.NumRows))
 	case *msgstream.TimeTickMsg:
 		mutableMessage = newV1TimeTickMsgFromV0(underlyingMsg)
 	case *msgstream.CreatePartitionMsg:
@@ -131,11 +131,12 @@ func newV1InsertMsgFromV0(msg *msgstream.InsertMsg, binarySize uint64) message.M
 }
 
 // newV1DeleteMsgFromV0 creates a new delete message from the old version delete message.
-func newV1DeleteMsgFromV0(msg *msgstream.DeleteMsg) message.MutableMessage {
+func newV1DeleteMsgFromV0(msg *msgstream.DeleteMsg, rows uint64) message.MutableMessage {
 	mutableMessage, err := message.NewDeleteMessageBuilderV1().
 		WithVChannel(msg.ShardName).
 		WithHeader(&message.DeleteMessageHeader{
 			CollectionId: msg.CollectionID,
+			Rows:         rows,
 		}).
 		WithBody(msg.DeleteRequest).
 		BuildMutable()

internal/streamingnode/server/wal/adaptor/old_version_message.go

@@ -201,6 +201,7 @@ func (impl *shardInterceptor) handleDeleteMessage(ctx context.Context, msg messa
 		return nil, status.NewUnrecoverableError(err.Error())
 	}
 
+	impl.shardManager.ApplyDelete(deleteMessage)
 	return appendOp(ctx, msg)
 }
 

internal/streamingnode/server/wal/interceptors/shard/shard_interceptor.go

@@ -218,6 +218,7 @@ func TestShardInterceptor(t *testing.T) {
 
 	shardManager.EXPECT().CheckIfCollectionExists(mock.Anything).Unset()
 	shardManager.EXPECT().CheckIfCollectionExists(mock.Anything).Return(nil)
+	shardManager.EXPECT().ApplyDelete(mock.Anything).Return(nil)
 	msgID, err = i.DoAppend(ctx, msg, appender)
 	assert.NoError(t, err)
 	assert.NotNil(t, msgID)

internal/streamingnode/server/wal/interceptors/shard/shard_interceptor_test.go

@@ -38,6 +38,8 @@ type ShardManager interface {
 
 	AssignSegment(req *AssignSegmentRequest) (*AssignSegmentResult, error)
 
+	ApplyDelete(msg message.MutableDeleteMessageV1) error
+
 	WaitUntilGrowingSegmentReady(collectionID int64, partitonID int64) (<-chan struct{}, error)
 
 	FlushAndFenceSegmentAllocUntil(collectionID int64, timetick uint64) ([]int64, error)

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

@@ -122,11 +122,25 @@ func (m *shardManagerImpl) AssignSegment(req *AssignSegmentRequest) (*AssignSegm
 	m.mu.Lock()
 	defer m.mu.Unlock()
 
-	if pm, ok := m.partitionManagers[req.PartitionID]; ok {
-		return pm.AssignSegment(req)
-	} else {
+	pm, ok := m.partitionManagers[req.PartitionID]
+	if !ok {
 		return nil, ErrPartitionNotFound
 	}
+	result, err := pm.AssignSegment(req)
+	if err == nil {
+		m.metrics.ObserveInsert(req.InsertMetrics.Rows, req.InsertMetrics.BinarySize)
+		return result, nil
+	}
+	return nil, err
+}
+
+// ApplyDelete: TODO move the L0 flush operation here.
+func (m *shardManagerImpl) ApplyDelete(msg message.MutableDeleteMessageV1) error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	m.metrics.ObserveDelete(msg.Header().GetRows())
+	return nil
 }
 
 // WaitUntilGrowingSegmentReady waits until the growing segment is ready.

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

@@ -349,7 +349,7 @@ func (m *StatsManager) selectSegmentsWithTimePolicy() map[int64]policy.SealPolic
 // selectSegmentsUntilLessThanLWM selects segments until the total size is less than the threshold.
 func (m *StatsManager) selectSegmentsUntilLessThanLWM() []int64 {
 	m.mu.Lock()
-	restSpace := m.totalStats.BinarySize - uint64(m.cfg.growingBytesLWM)
+	restSpace := int64(m.totalStats.BinarySize) - m.cfg.growingBytesLWM
 	m.mu.Unlock()
 
 	if restSpace <= 0 {
@@ -363,7 +363,7 @@ func (m *StatsManager) selectSegmentsUntilLessThanLWM() []int64 {
 	})
 	for restSpace > 0 && statsHeap.Len() > 0 {
 		nextOne := statsHeap.Pop()
-		restSpace -= nextOne.binarySize
+		restSpace -= int64(nextOne.binarySize)
 		segmentIDs = append(segmentIDs, nextOne.segmentID)
 	}
 	return segmentIDs

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

@@ -22,6 +22,10 @@ func NewSegmentAssignMetrics(pchannel string) *SegmentAssignMetrics {
 		flushedTotal:     metrics.WALSegmentFlushedTotal.MustCurryWith(constLabel),
 		partitionTotal:   metrics.WALPartitionTotal.With(constLabel),
 		collectionTotal:  metrics.WALCollectionTotal.With(constLabel),
+
+		insertRowsTotal: metrics.WALInsertRowsTotal.With(constLabel),
+		insertBytes:     metrics.WALInsertBytes.With(constLabel),
+		deleteRowsTotal: metrics.WALDeleteRowsTotal.With(constLabel),
 	}
 }
 
@@ -37,6 +41,10 @@ type SegmentAssignMetrics struct {
 	flushedTotal     *prometheus.CounterVec
 	partitionTotal   prometheus.Gauge
 	collectionTotal  prometheus.Gauge
+
+	insertRowsTotal prometheus.Counter
+	insertBytes     prometheus.Counter
+	deleteRowsTotal prometheus.Counter
 }
 
 // ObserveOnAllocating observe a allocating operation and return a guard function.
@@ -55,6 +63,17 @@ func (m *SegmentAssignMetrics) ObseveOnFlushing() func() {
 	}
 }
 
+// ObserveInsert observe a insert operation
+func (m *SegmentAssignMetrics) ObserveInsert(rows uint64, bytes uint64) {
+	m.insertRowsTotal.Add(float64(rows))
+	m.insertBytes.Add(float64(bytes))
+}
+
+// ObserveDelete observe a insert operation
+func (m *SegmentAssignMetrics) ObserveDelete(rows uint64) {
+	m.deleteRowsTotal.Add(float64(rows))
+}
+
 // ObserveCreateSegment increments the total number of growing segment.
 func (m *SegmentAssignMetrics) ObserveCreateSegment() {
 	m.allocTotal.Inc()
@@ -87,4 +106,7 @@ func (m *SegmentAssignMetrics) Close() {
 	metrics.WALSegmentBytes.Delete(m.constLabel)
 	metrics.WALPartitionTotal.Delete(m.constLabel)
 	metrics.WALCollectionTotal.Delete(m.constLabel)
+	metrics.WALInsertRowsTotal.Delete(m.constLabel)
+	metrics.WALDeleteRowsTotal.Delete(m.constLabel)
+	metrics.WALInsertBytes.Delete(m.constLabel)
 }

internal/streamingnode/server/wal/metricsutil/segment.go

@@ -231,6 +231,22 @@ var (
 		Buckets: secondsBuckets,
 	}, WALChannelLabelName, WALTxnStateLabelName)
 
+	// Rows level counter.
+	WALInsertRowsTotal = newWALGaugeVec(prometheus.GaugeOpts{
+		Name: "insert_rows_total",
+		Help: "Rows of growing insert on wal",
+	}, WALChannelLabelName)
+
+	WALInsertBytes = newWALGaugeVec(prometheus.GaugeOpts{
+		Name: "insert_bytes",
+		Help: "Bytes of growing insert on wal",
+	}, WALChannelLabelName)
+
+	WALDeleteRowsTotal = newWALGaugeVec(prometheus.GaugeOpts{
+		Name: "delete_rows_total",
+		Help: "Rows of growing delete on wal",
+	}, WALChannelLabelName)
+
 	// Segment related metrics
 	WALGrowingSegmentRowsTotal = newWALGaugeVec(prometheus.GaugeOpts{
 		Name: "growing_segment_rows_total",
@@ -495,6 +511,9 @@ func registerWAL(registry *prometheus.Registry) {
 	registry.MustRegister(WALTimeTickSyncTimeTick)
 	registry.MustRegister(WALInflightTxn)
 	registry.MustRegister(WALTxnDurationSeconds)
+	registry.MustRegister(WALInsertRowsTotal)
+	registry.MustRegister(WALInsertBytes)
+	registry.MustRegister(WALDeleteRowsTotal)
 	registry.MustRegister(WALGrowingSegmentBytes)
 	registry.MustRegister(WALGrowingSegmentRowsTotal)
 	registry.MustRegister(WALGrowingSegmentHWMBytes)

pkg/metrics/streaming_service_metrics.go

@@ -137,6 +137,7 @@ message SegmentAssignment {
 // DeleteMessageHeader
 message DeleteMessageHeader {
     int64 collection_id = 1;
+    uint64 rows = 2;
 }
 
 // FlushMessageHeader just nothing.